diff --git a/components/spider-proto-rust/src/generated/storage.rs b/components/spider-proto-rust/src/generated/storage.rs index b4d4fcc5e..1c38d945a 100644 --- a/components/spider-proto-rust/src/generated/storage.rs +++ b/components/spider-proto-rust/src/generated/storage.rs @@ -1635,6 +1635,32 @@ pub mod inbound_queue_service_client { ); self.inner.unary(req, path, codec).await } + pub async fn resend_ready_tasks( + &mut self, + request: impl tonic::IntoRequest, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + > { + self.inner + .ready() + .await + .map_err(|e| { + tonic::Status::unknown( + format!("Service was not ready: {}", e.into()), + ) + })?; + let codec = tonic_prost::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/storage.InboundQueueService/ResendReadyTasks", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert( + GrpcMethod::new("storage.InboundQueueService", "ResendReadyTasks"), + ); + self.inner.unary(req, path, codec).await + } } } /// Generated server implementations. @@ -1671,6 +1697,13 @@ pub mod inbound_queue_service_server { tonic::Response, tonic::Status, >; + async fn resend_ready_tasks( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; } #[derive(Debug)] pub struct InboundQueueServiceServer { @@ -1895,6 +1928,55 @@ pub mod inbound_queue_service_server { }; Box::pin(fut) } + "/storage.InboundQueueService/ResendReadyTasks" => { + #[allow(non_camel_case_types)] + struct ResendReadyTasksSvc(pub Arc); + impl< + T: InboundQueueService, + > tonic::server::UnaryService + for ResendReadyTasksSvc { + type Response = super::super::common::Void; + type Future = BoxFuture< + tonic::Response, + tonic::Status, + >; + fn call( + &mut self, + request: tonic::Request, + ) -> Self::Future { + let inner = Arc::clone(&self.0); + let fut = async move { + ::resend_ready_tasks( + &inner, + request, + ) + .await + }; + Box::pin(fut) + } + } + let accept_compression_encodings = self.accept_compression_encodings; + let send_compression_encodings = self.send_compression_encodings; + let max_decoding_message_size = self.max_decoding_message_size; + let max_encoding_message_size = self.max_encoding_message_size; + let inner = self.inner.clone(); + let fut = async move { + let method = ResendReadyTasksSvc(inner); + let codec = tonic_prost::ProstCodec::default(); + let mut grpc = tonic::server::Grpc::new(codec) + .apply_compression_config( + accept_compression_encodings, + send_compression_encodings, + ) + .apply_max_message_size_config( + max_decoding_message_size, + max_encoding_message_size, + ); + let res = grpc.unary(method, req).await; + Ok(res) + }; + Box::pin(fut) + } _ => { Box::pin(async move { let mut response = http::Response::new( diff --git a/components/spider-proto-rust/src/lib.rs b/components/spider-proto-rust/src/lib.rs index 6016d9045..e19242bb7 100644 --- a/components/spider-proto-rust/src/lib.rs +++ b/components/spider-proto-rust/src/lib.rs @@ -5,6 +5,7 @@ pub mod error; pub mod id; pub mod io; pub mod job; +pub mod scheduler_registration; pub mod unpack; #[allow(clippy::all, clippy::nursery, clippy::pedantic)] diff --git a/components/spider-proto-rust/src/scheduler_registration.rs b/components/spider-proto-rust/src/scheduler_registration.rs new file mode 100644 index 000000000..e01403e0b --- /dev/null +++ b/components/spider-proto-rust/src/scheduler_registration.rs @@ -0,0 +1,40 @@ +//! Conversions between protobuf scheduler messages and their Spider core representations. + +use spider_core::types::scheduler::RegisteredScheduler; + +use crate::storage; + +impl From for storage::Scheduler { + fn from(scheduler: RegisteredScheduler) -> Self { + Self { + scheduler_id: scheduler.id.get(), + ip_address: scheduler.ip_address.to_string(), + port: u32::from(scheduler.port), + } + } +} + +#[cfg(test)] +mod tests { + use std::net::IpAddr; + + use spider_core::types::{id::SchedulerId, scheduler::RegisteredScheduler}; + + use crate::storage; + + #[test] + fn registered_scheduler_to_protocol_carries_id_ip_and_port() { + const SCHEDULER_ID: SchedulerId = SchedulerId::from(42); + const PORT: u16 = 5678; + + let scheduler = storage::Scheduler::from(RegisteredScheduler { + id: SCHEDULER_ID, + ip_address: IpAddr::V4("127.0.0.1".parse().expect("valid IP")), + port: PORT, + }); + + assert_eq!(scheduler.scheduler_id, SCHEDULER_ID.get()); + assert_eq!(scheduler.ip_address, "127.0.0.1"); + assert_eq!(scheduler.port, u32::from(PORT)); + } +} diff --git a/components/spider-proto-rust/src/unpack/storage.rs b/components/spider-proto-rust/src/unpack/storage.rs index ded53f6d2..ebeba4910 100644 --- a/components/spider-proto-rust/src/unpack/storage.rs +++ b/components/spider-proto-rust/src/unpack/storage.rs @@ -1,5 +1,7 @@ //! [`RequestUnpack`] implementations for `storage.proto` requests. +use std::{net::IpAddr, time::Duration}; + use spider_core::types::id::{ ExecutionManagerId, JobId, @@ -8,14 +10,21 @@ use spider_core::types::id::{ TaskId, TaskInstanceId, }; +use tonic::Code; use crate::{ storage::{ + AddResourceGroupRequest, + ExecutionManagerIdRequest, JobIdRequest, + PollReadyTasksRequest, + RegisterExecutionManagerRequest, RegisterJobRequest, + RegisterSchedulerRequest, RegisterTaskInstanceRequest, ReportTaskFailureRequest, ReportTaskSuccessRequest, + VerifyResourceGroupRequest, }, unpack::{RequestUnpack, UnpackError, common::unpack_task_id}, }; @@ -134,3 +143,95 @@ impl RequestUnpack for ReportTaskFailureRequest { )) } } + +/// Unpacks [`AddResourceGroupRequest`] into a tuple containing: +/// +/// * The external resource group ID. +/// * The password. +impl RequestUnpack for AddResourceGroupRequest { + type Unpacked = (String, Vec); + + fn unpack(self) -> Result { + Ok((self.external_resource_group_id, self.password)) + } +} + +/// Unpacks [`VerifyResourceGroupRequest`] into a tuple containing: +/// +/// * The resource group ID. +/// * The password. +impl RequestUnpack for VerifyResourceGroupRequest { + type Unpacked = (ResourceGroupId, Vec); + + fn unpack(self) -> Result { + Ok((ResourceGroupId::from(self.resource_group_id), self.password)) + } +} + +/// Unpacks [`RegisterExecutionManagerRequest`] into the execution manager's IP address. +impl RequestUnpack for RegisterExecutionManagerRequest { + type Unpacked = IpAddr; + + fn unpack(self) -> Result { + self.ip_address + .parse::() + .map_err(|error| invalid_argument(format!("invalid IP address: {error}"))) + } +} + +/// Unpacks [`ExecutionManagerIdRequest`] into an [`ExecutionManagerId`]. +impl RequestUnpack for ExecutionManagerIdRequest { + type Unpacked = ExecutionManagerId; + + fn unpack(self) -> Result { + Ok(ExecutionManagerId::from(self.execution_manager_id)) + } +} + +/// Unpacks [`RegisterSchedulerRequest`] into a tuple containing: +/// +/// * The scheduler IP address. +/// * The scheduler port. +impl RequestUnpack for RegisterSchedulerRequest { + type Unpacked = (IpAddr, u16); + + fn unpack(self) -> Result { + let ip_address = self + .ip_address + .parse::() + .map_err(|error| invalid_argument(format!("invalid IP address: {error}")))?; + let port = u16::try_from(self.port) + .map_err(|_| invalid_argument(format!("port does not fit in `u16`: {}", self.port)))?; + Ok((ip_address, port)) + } +} + +/// Unpacks [`PollReadyTasksRequest`] into a tuple containing: +/// +/// * The maximum number of entries to return. +/// * The maximum duration to block waiting for entries. +impl RequestUnpack for PollReadyTasksRequest { + type Unpacked = (usize, Duration); + + fn unpack(self) -> Result { + let max_items = usize::try_from(self.max_items).map_err(|_| { + invalid_argument(format!( + "max_items does not fit in `usize`: {}", + self.max_items + )) + })?; + Ok((max_items, Duration::from_millis(self.wait_ms))) + } +} + +/// Builds an [`UnpackError`] carrying [`Code::InvalidArgument`] and the given message. +/// +/// # Returns +/// +/// An [`UnpackError`] whose [`Code`] is [`Code::InvalidArgument`] and whose message is `message`. +const fn invalid_argument(message: String) -> UnpackError { + UnpackError { + code: Code::InvalidArgument, + message, + } +} diff --git a/components/spider-proto/storage/storage.proto b/components/spider-proto/storage/storage.proto index 11caf6e3f..1913b1237 100644 --- a/components/spider-proto/storage/storage.proto +++ b/components/spider-proto/storage/storage.proto @@ -23,6 +23,7 @@ service InboundQueueService { rpc PollReadyTasks(PollReadyTasksRequest) returns (PollReadyTasksResponse); rpc PollReadyCommitTasks(PollReadyTasksRequest) returns (PollReadyTasksResponse); rpc PollReadyCleanupTasks(PollReadyTasksRequest) returns (PollReadyTasksResponse); + rpc ResendReadyTasks(common.Void) returns (common.Void); } service ResourceGroupManagementService { diff --git a/components/spider-scheduler/src/core_impl/round_robin/tests.rs b/components/spider-scheduler/src/core_impl/round_robin/tests.rs index a98228013..b8a259442 100644 --- a/components/spider-scheduler/src/core_impl/round_robin/tests.rs +++ b/components/spider-scheduler/src/core_impl/round_robin/tests.rs @@ -178,6 +178,10 @@ impl SchedulerStorageClient for MockStorageClient { async fn job_state(&self, _job_id: JobId) -> Result { Ok(JobState::Running) } + + async fn resend_ready_tasks(&self) -> Result<(), StorageClientError> { + Ok(()) + } } /// # Returns diff --git a/components/spider-scheduler/src/error.rs b/components/spider-scheduler/src/error.rs index a10689f41..c0d61762a 100644 --- a/components/spider-scheduler/src/error.rs +++ b/components/spider-scheduler/src/error.rs @@ -11,13 +11,6 @@ pub enum StorageClientError { #[error("job not found: {0:?}")] JobNotFound(JobId), - /// The scheduler's storage session is stale. - #[error("stale storage session: {storage_session:?}")] - StaleSession { - /// Storage's current session ID. - storage_session: SessionId, - }, - /// The storage server returned an invalid input error. #[error("invalid storage request: {0}")] InvalidInput(String), diff --git a/components/spider-scheduler/src/runtime.rs b/components/spider-scheduler/src/runtime.rs index d514498a7..aa8c42bb0 100644 --- a/components/spider-scheduler/src/runtime.rs +++ b/components/spider-scheduler/src/runtime.rs @@ -125,6 +125,8 @@ pub async fn create_runtime Result { Ok(JobState::Running) } + + async fn resend_ready_tasks(&self) -> Result<(), StorageClientError> { + Ok(()) + } } /// # Returns diff --git a/components/spider-scheduler/src/storage_client/grpc.rs b/components/spider-scheduler/src/storage_client/grpc.rs index 7c5239698..c9ca8c2a0 100644 --- a/components/spider-scheduler/src/storage_client/grpc.rs +++ b/components/spider-scheduler/src/storage_client/grpc.rs @@ -7,11 +7,14 @@ use spider_core::{ job::JobState, types::id::{JobId, ResourceGroupId, SchedulerId, SessionId, TaskId}, }; -use spider_proto_rust::storage::{ - self, - inbound_queue_service_client::InboundQueueServiceClient, - job_orchestration_service_client::JobOrchestrationServiceClient, - scheduler_registration_service_client::SchedulerRegistrationServiceClient, +use spider_proto_rust::{ + common, + storage::{ + self, + inbound_queue_service_client::InboundQueueServiceClient, + job_orchestration_service_client::JobOrchestrationServiceClient, + scheduler_registration_service_client::SchedulerRegistrationServiceClient, + }, }; use spider_utils::grpc::client::ConnectionPool; use tonic::{ @@ -167,6 +170,15 @@ impl SchedulerStorageClient for GrpcSchedulerStorageClient { .into_inner(); job_state_response_to_result(response) } + + async fn resend_ready_tasks(&self) -> Result<(), StorageClientError> { + self.inbound_queue + .get_client() + .resend_ready_tasks(common::Void {}) + .await + .map_err(|status| inbound_status_to_error(&status))?; + Ok(()) + } } /// Maps an inbound-queue gRPC [`Status`] to a [`StorageClientError`]. diff --git a/components/spider-scheduler/src/storage_client/mod.rs b/components/spider-scheduler/src/storage_client/mod.rs index 3f3f75246..c4c66ad23 100644 --- a/components/spider-scheduler/src/storage_client/mod.rs +++ b/components/spider-scheduler/src/storage_client/mod.rs @@ -145,4 +145,15 @@ pub trait SchedulerStorageClient: Send + Sync + Clone { /// * [`StorageClientError::Server`] if the storage server returns an error. /// * [`StorageClientError::Transport`] if the storage server returns malformed data. async fn job_state(&self, job_id: JobId) -> Result; + + /// Asks storage to resend ready tasks for all jobs in the cache to the ready queue. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`StorageClientError::Server`] if the storage server returns an error. + /// * [`StorageClientError::Transport`] if the storage transport fails or returns malformed + /// data. + async fn resend_ready_tasks(&self) -> Result<(), StorageClientError>; } diff --git a/components/spider-storage/src/grpc.rs b/components/spider-storage/src/grpc.rs index ea0a145ed..a9f0b47ef 100644 --- a/components/spider-storage/src/grpc.rs +++ b/components/spider-storage/src/grpc.rs @@ -20,9 +20,9 @@ use tokio_util::sync::CancellationToken; use tonic::{Request, Response, Status}; use crate::{ - cache::error::CacheError, + cache::error::{CacheError, InternalError}, db::{DbError, DbStorage}, - ready_queue::ReadyQueueSender, + ready_queue::{ReadyQueueEntry, ReadyQueueSender}, state::{ServiceState, StorageServerError}, task_instance_pool::TaskInstancePoolConnector, }; @@ -78,7 +78,7 @@ impl< /// * `INTERNAL` for: /// * A fatal cache-internal error (the service will restart). /// * Any other (database or otherwise unexpected) error. - /// * `UNAUTHENTICATED` for an unknown or unauthorized resource group. + /// * `UNAUTHENTICATED` for an unknown or unauthorized resource group, or a wrong password. /// * `NOT_FOUND` for a missing job. /// * `FAILED_PRECONDITION` for operations on an invalid job state. /// * `INVALID_ARGUMENT` for a malformed task graph, inputs, or request. @@ -90,14 +90,7 @@ impl< const SERVICE_NAME: &str = "JobOrchestration"; match error { StorageServerError::Cache(CacheError::Internal(e)) => { - tracing::error!( - error = % e, - service = SERVICE_NAME, - tag, - "Internal error in the cache layer. Cancelling service." - ); - self.cancellation_token.cancel(); - Status::internal("storage service unavailable") + self.fatal_internal_status(SERVICE_NAME, tag, &e) } StorageServerError::Db(db_error) => match &db_error { @@ -159,15 +152,7 @@ impl< Status::invalid_argument(error.to_string()) } - _ => { - tracing::error!( - error = % error, - service = SERVICE_NAME, - tag, - "Unexpected internal error." - ); - Status::internal("internal error") - } + _ => self.unexpected_internal_status(SERVICE_NAME, tag, &error, false), } } @@ -194,14 +179,7 @@ impl< const SERVICE_NAME: &str = "TaskInstanceManagement"; match error { StorageServerError::Cache(CacheError::Internal(e)) => { - tracing::error!( - error = % e, - service = SERVICE_NAME, - tag, - "Internal error in the cache layer. Cancelling service." - ); - self.cancellation_token.cancel(); - Status::internal("storage service unavailable") + self.fatal_internal_status(SERVICE_NAME, tag, &e) } StorageServerError::StaleSession(storage_session) => { @@ -247,16 +225,241 @@ impl< Status::invalid_argument(error.to_string()) } - _ => { - tracing::error!( + _ => self.unexpected_internal_status(SERVICE_NAME, tag, &error, true), + } + } + + /// Error handler for inbound queue service errors. + /// + /// This function maps the given [`StorageServerError`] to a [`Status`] that can be sent to the + /// client. The errors are logged for observability. + /// + /// # Returns + /// + /// The [`Status`] to send to the client: + /// + /// * `INTERNAL` for any failure happened on the server side. This method should never fail + /// under the service's assumption. + #[must_use] + #[allow(clippy::needless_pass_by_value)] + pub fn inbound_queue_service_error_handler( + &self, + error: StorageServerError, + tag: &'static str, + ) -> Status { + const SERVICE_NAME: &str = "InboundQueue"; + self.unexpected_internal_status(SERVICE_NAME, tag, &error, false) + } + + /// Error handler for resource group management service errors. + /// + /// This function maps the given [`StorageServerError`] to a [`Status`] that can be sent to the + /// client. The errors are logged for observability. + /// + /// # Returns + /// + /// The [`Status`] to send to the client: + /// + /// * `UNAUTHENTICATED` for an unknown resource group or a wrong password. + /// * `ALREADY_EXISTS` for a duplicate external resource group ID. + /// * `INTERNAL` for: + /// * A fatal cache-internal error (the service will restart). + /// * Any other unexpected failure. + pub fn resource_group_management_service_error_handler( + &self, + error: StorageServerError, + tag: &'static str, + ) -> Status { + const SERVICE_NAME: &str = "ResourceGroupManagement"; + match error { + error @ StorageServerError::Db( + DbError::ResourceGroupNotFound(_) | DbError::InvalidPassword(_), + ) => { + tracing::warn!( error = % error, service = SERVICE_NAME, tag, - "Unexpected internal error. Cancelling service to avoid cache corruption." + "Invalid resource group." ); - self.cancellation_token.cancel(); - Status::internal("internal error") + Status::unauthenticated("invalid resource group") + } + + error @ StorageServerError::Db(DbError::ResourceGroupAlreadyExists(_)) => { + tracing::warn!( + error = % error, + service = SERVICE_NAME, + tag, + "Resource group already exists." + ); + Status::already_exists(error.to_string()) + } + + StorageServerError::Cache(CacheError::Internal(e)) => { + self.fatal_internal_status(SERVICE_NAME, tag, &e) + } + + error => self.unexpected_internal_status(SERVICE_NAME, tag, &error, false), + } + } + + /// Error handler for execution manager liveness service errors. + /// + /// This function maps the given [`StorageServerError`] to a [`Status`] that can be sent to the + /// client. The errors are logged for observability. + /// + /// # Returns + /// + /// The [`Status`] to send to the client: + /// + /// * `FAILED_PRECONDITION` when the execution manager has already been reaped. + /// * `INVALID_ARGUMENT` for an illegal execution manager ID. + /// * `INTERNAL` for: + /// * A fatal cache-internal error (the service will restart). + /// * Any other unexpected failure. + pub fn execution_manager_liveness_service_error_handler( + &self, + error: StorageServerError, + tag: &'static str, + ) -> Status { + const SERVICE_NAME: &str = "ExecutionManagerLiveness"; + match error { + error @ StorageServerError::Db(DbError::ExecutionManagerAlreadyDead(_)) => { + tracing::warn!( + error = % error, + service = SERVICE_NAME, + tag, + "Execution manager already marked dead." + ); + Status::failed_precondition(error.to_string()) + } + + error @ StorageServerError::Db(DbError::IllegalExecutionManagerId(_)) => { + tracing::warn!( + error = % error, + service = SERVICE_NAME, + tag, + "Illegal execution manager id." + ); + Status::invalid_argument(error.to_string()) + } + + StorageServerError::Cache(CacheError::Internal(e)) => { + self.fatal_internal_status(SERVICE_NAME, tag, &e) + } + + error => self.unexpected_internal_status(SERVICE_NAME, tag, &error, false), + } + } + + /// Error handler for scheduler registration service errors. + /// + /// This function maps the given [`StorageServerError`] to a [`Status`] that can be sent to the + /// client. The errors are logged for observability. + /// + /// # Returns + /// + /// The [`Status`] to send to the client: + /// + /// * `INTERNAL` for: + /// * A fatal cache-internal error (the service will restart). + /// * Any other failure; scheduler registration currently has no caller-visible error + /// classification beyond a generic server error. + #[must_use] + pub fn scheduler_registration_service_error_handler( + &self, + error: StorageServerError, + tag: &'static str, + ) -> Status { + const SERVICE_NAME: &str = "SchedulerRegistration"; + match error { + StorageServerError::Cache(CacheError::Internal(e)) => { + self.fatal_internal_status(SERVICE_NAME, tag, &e) } + + error => self.unexpected_internal_status(SERVICE_NAME, tag, &error, false), + } + } + + /// Logs a fatal cache-internal error, cancels the service, and returns an `INTERNAL` status. + /// + /// # Returns + /// + /// An `INTERNAL` [`Status`] carrying the message `"storage service unavailable"`. + fn fatal_internal_status( + &self, + service_name: &'static str, + tag: &'static str, + error: &InternalError, + ) -> Status { + tracing::error!( + error = % error, + service = service_name, + tag, + "Internal error in the cache layer. Cancelling service." + ); + self.cancellation_token.cancel(); + Status::internal("storage service unavailable") + } + + /// Logs an unexpected error, cancels the service, and returns an `INTERNAL` status. + /// + /// If `is_fatal` flag is raised, the service cancellation token will be fired. + /// + /// # Returns + /// + /// An `INTERNAL` [`Status`] carrying the message `"internal error"`. + fn unexpected_internal_status( + &self, + service_name: &'static str, + tag: &'static str, + error: &StorageServerError, + is_fatal: bool, + ) -> Status { + tracing::error!( + error = % error, + service = service_name, + tag, + "Unexpected internal error. Cancelling service to avoid cache corruption." + ); + if is_fatal { + self.cancellation_token.cancel(); + } + Status::internal("internal error") + } + + /// Builds a [`storage::ReadyTasks`] message from a batch of ready-queue entries. + /// + /// # Type Parameters + /// + /// * `TaskKindType` - The kind of ready-queue task carried by each entry: + /// * [`spider_core::task::TaskIndex`] for the regular lane. + /// * [`crate::ready_queue::CommitTaskMarker`] for the commit lane. + /// * [`crate::ready_queue::CleanupTaskMarker`] for the cleanup lane. + /// + /// # Returns + /// + /// A [`storage::ReadyTasks`] carrying the storage session and the flattened ready tasks. + fn build_ready_tasks( + &self, + entries: Vec>, + to_task_id: impl Fn(TaskKindType) -> common::TaskId, + ) -> storage::ReadyTasks { + let tasks = entries + .into_iter() + .map(|entry| { + let resource_group_id = entry.resource_group_id.get(); + let job_id = entry.job_id.get(); + let task_id = to_task_id(entry.task_kind); + storage::ReadyTask { + resource_group_id, + job_id, + task_id: Some(task_id), + } + }) + .collect(); + storage::ReadyTasks { + session_id: self.inner.session_id(), + tasks, } } } @@ -495,23 +698,67 @@ impl< { async fn poll_ready_tasks( &self, - _request: Request, + request: Request, ) -> Result, Status> { - todo!("Not implemented") + let (max_items, wait) = request.into_inner().unpack()?; + tracing::info!(max_items, "Poll ready tasks request received."); + let entries = self + .inner + .poll_ready_tasks(max_items, wait) + .await + .map_err(|error| self.inbound_queue_service_error_handler(error, "poll_ready_tasks"))?; + Ok(Response::new(storage::PollReadyTasksResponse { + tasks: Some(self.build_ready_tasks(entries, |task_index| { + common::TaskId::from(TaskId::Index(task_index)) + })), + })) } async fn poll_ready_commit_tasks( &self, - _request: Request, + request: Request, ) -> Result, Status> { - todo!("Not implemented") + let (max_items, wait) = request.into_inner().unpack()?; + tracing::info!(max_items, "Poll ready commit tasks request received."); + let entries = self + .inner + .poll_commit_ready_tasks(max_items, wait) + .await + .map_err(|error| { + self.inbound_queue_service_error_handler(error, "poll_ready_commit_tasks") + })?; + Ok(Response::new(storage::PollReadyTasksResponse { + tasks: Some(self.build_ready_tasks(entries, |_| common::TaskId::from(TaskId::Commit))), + })) } async fn poll_ready_cleanup_tasks( &self, - _request: Request, + request: Request, ) -> Result, Status> { - todo!("Not implemented") + let (max_items, wait) = request.into_inner().unpack()?; + tracing::info!(max_items, "Poll ready cleanup tasks request received."); + let entries = self + .inner + .poll_cleanup_ready_tasks(max_items, wait) + .await + .map_err(|error| { + self.inbound_queue_service_error_handler(error, "poll_ready_cleanup_tasks") + })?; + Ok(Response::new(storage::PollReadyTasksResponse { + tasks: Some(self.build_ready_tasks(entries, |_| common::TaskId::from(TaskId::Cleanup))), + })) + } + + async fn resend_ready_tasks( + &self, + _request: Request, + ) -> Result, Status> { + tracing::info!("Resend ready tasks request received."); + self.inner.resend_ready_tasks().await.map_err(|error| { + self.inbound_queue_service_error_handler(error, "resend_ready_tasks") + })?; + Ok(Response::new(common::Void {})) } } @@ -525,16 +772,38 @@ impl< { async fn add_resource_group( &self, - _request: Request, + request: Request, ) -> Result, Status> { - todo!("Not implemented") + let (external_id, password) = request.into_inner().unpack()?; + tracing::info!(external_id = % external_id, "Add resource group request received."); + let rg_id = self + .inner + .add_resource_group(external_id, password) + .await + .map_err(|error| { + self.resource_group_management_service_error_handler(error, "add_resource_group") + })?; + Ok(Response::new(storage::ResourceGroupIdResponse { + resource_group_id: rg_id.get(), + })) } async fn verify_resource_group( &self, - _request: Request, + request: Request, ) -> Result, Status> { - todo!("Not implemented") + let (rg_id, password) = request.into_inner().unpack()?; + tracing::info!( + rg_id = rg_id.get(), + "Verify resource group request received." + ); + self.inner + .verify_resource_group(rg_id, &password) + .await + .map_err(|error| { + self.resource_group_management_service_error_handler(error, "verify_resource_group") + })?; + Ok(Response::new(common::Void {})) } } @@ -548,16 +817,51 @@ impl< { async fn register_execution_manager( &self, - _request: Request, + request: Request, ) -> Result, Status> { - todo!("Not implemented") + let ip_address = request.into_inner().unpack()?; + tracing::info!(% ip_address, "Execution manager registration request received."); + let em_id = self + .inner + .register_execution_manager(ip_address) + .await + .map_err(|error| { + self.execution_manager_liveness_service_error_handler( + error, + "register_execution_manager", + ) + })?; + Ok(Response::new(storage::RegisterExecutionManagerResponse { + registration: Some(storage::ExecutionManagerRegistration { + execution_manager_id: em_id.get(), + session_id: self.inner.session_id(), + }), + })) } async fn update_execution_manager_heartbeat( &self, - _request: Request, + request: Request, ) -> Result, Status> { - todo!("Not implemented") + let em_id = request.into_inner().unpack()?; + tracing::info!( + em_id = em_id.get(), + "Execution manager heartbeat request received." + ); + self.inner + .update_execution_manager_heartbeat(em_id) + .await + .map_err(|error| { + self.execution_manager_liveness_service_error_handler( + error, + "update_execution_manager_heartbeat", + ) + })?; + Ok(Response::new( + storage::UpdateExecutionManagerHeartbeatResponse { + session_id: self.inner.session_id(), + }, + )) } } @@ -571,16 +875,41 @@ impl< { async fn register_scheduler( &self, - _request: Request, + request: Request, ) -> Result, Status> { - todo!("Not implemented") + let (ip_address, port) = request.into_inner().unpack()?; + tracing::info!(% ip_address, port, "Scheduler registration request received."); + let scheduler_id = self + .inner + .register_scheduler(ip_address, port) + .await + .map_err(|error| { + self.scheduler_registration_service_error_handler(error, "register_scheduler") + })?; + Ok(Response::new(storage::RegisterSchedulerResponse { + registration: Some(storage::SchedulerRegistration { + scheduler_id: scheduler_id.get(), + session_id: self.inner.session_id(), + }), + })) } async fn get_schedulers( &self, _request: Request, ) -> Result, Status> { - todo!("Not implemented") + tracing::info!("Get schedulers request received."); + let schedulers = self.inner.get_schedulers().await.map_err(|error| { + self.scheduler_registration_service_error_handler(error, "get_schedulers") + })?; + Ok(Response::new(storage::GetSchedulersResponse { + schedulers: Some(storage::SchedulerRegistrations { + schedulers: schedulers + .into_iter() + .map(storage::Scheduler::from) + .collect(), + }), + })) } } @@ -596,7 +925,9 @@ impl< &self, _request: Request, ) -> Result, Status> { - todo!("Not implemented") + let session_id = self.inner.session_id(); + tracing::info!(session_id, "Get session request received."); + Ok(Response::new(storage::GetSessionResponse { session_id })) } } diff --git a/components/spider-storage/src/state.rs b/components/spider-storage/src/state.rs index 4b76ec055..3ea42a3b5 100644 --- a/components/spider-storage/src/state.rs +++ b/components/spider-storage/src/state.rs @@ -11,4 +11,4 @@ pub use runtime::{Runtime, create_runtime}; pub use service::ServiceState; #[cfg(test)] -mod test_utils; +pub(crate) mod test_utils;