feat(spider-storage): Implement scheduler registration gRPC service and re-enqueue ready tasks when a scheduler reconnects.

components/spider-storage/src/grpc.rs

@@ -6,6 +6,7 @@ use spider_proto_rust::{
     common,
     storage::{
         self,
+        SchedulerRegistration,
         execution_manager_liveness_service_server::ExecutionManagerLivenessService,
         inbound_queue_service_server::InboundQueueService,
         job_orchestration_service_server::JobOrchestrationService,
@@ -36,18 +37,18 @@ use crate::{
 /// * `TaskInstancePoolConnectorType` - The task instance pool connector type.
 #[derive(Clone)]
 pub struct GrpcServiceState<
-    ReadyQueueSenderType: ReadyQueueSender,
-    DbConnectorType: DbStorage,
-    TaskInstancePoolConnectorType: TaskInstancePoolConnector,
+    ReadyQueueSenderType: ReadyQueueSender + 'static,
+    DbConnectorType: DbStorage + 'static,
+    TaskInstancePoolConnectorType: TaskInstancePoolConnector + 'static,
 > {
     inner: ServiceState<ReadyQueueSenderType, DbConnectorType, TaskInstancePoolConnectorType>,
     cancellation_token: CancellationToken,
 }
 
 impl<
-    ReadyQueueSenderType: ReadyQueueSender,
-    DbConnectorType: DbStorage,
-    TaskInstancePoolConnectorType: TaskInstancePoolConnector,
+    ReadyQueueSenderType: ReadyQueueSender + 'static,
+    DbConnectorType: DbStorage + 'static,
+    TaskInstancePoolConnectorType: TaskInstancePoolConnector + 'static,
 > GrpcServiceState<ReadyQueueSenderType, DbConnectorType, TaskInstancePoolConnectorType>
 {
     /// Factory function.
@@ -330,6 +331,27 @@ impl<
         }
     }
 
+    /// 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 any failure happened on the server side.
+    #[must_use]
+    #[allow(clippy::needless_pass_by_value)]
+    fn scheduler_registration_service_error_handler(
+        &self,
+        error: StorageServerError,
+        tag: &'static str,
+    ) -> Status {
+        const SERVICE_NAME: &str = "SchedulerRegistration";
+        self.default_error_handler(SERVICE_NAME, tag, &error, false)
+    }
+
     /// Handles generic storage server errors.
     ///
     /// This handler maps every [`StorageServerError`] to an `INTERNAL` [`Status`] with a generic
@@ -835,16 +857,30 @@ impl<
 {
     async fn register_scheduler(
         &self,
-        _request: Request<storage::RegisterSchedulerRequest>,
+        request: Request<storage::RegisterSchedulerRequest>,
     ) -> Result<Response<storage::RegisterSchedulerResponse>, Status> {
-        todo!("Not implemented")
+        let (ip_addr, port) = request.into_inner().unpack()?;
+        tracing::info!(% ip_addr, port, "Scheduler registration request received.");
+        let scheduler_id = self
+            .inner
+            .register_scheduler(ip_addr, port)
+            .await
+            .map_err(|error| {
+                self.scheduler_registration_service_error_handler(error, "register_scheduler")
+            })?;
+        Ok(Response::new(storage::RegisterSchedulerResponse {
+            registration: Some(SchedulerRegistration {
+                scheduler_id: scheduler_id.get(),
+                session_id: self.inner.session_id(),
+            }),
+        }))
     }
 
     async fn get_schedulers(
         &self,
         _request: Request<common::Void>,
     ) -> Result<Response<storage::GetSchedulersResponse>, Status> {
-        todo!("Not implemented")
+        Err(Status::unimplemented("not implemented"))
     }
 }
 

components/spider-storage/src/state.rs

@@ -8,7 +8,7 @@ pub use error::StorageServerError;
 pub use job_cache::JobCache;
 pub use job_cache_gc::{JobCacheGcConfig, JobCacheGcHandle, create_job_cache_gc};
 pub use runtime::{Runtime, create_runtime};
-pub use service::ServiceState;
+pub use service::{ServiceState, ServiceStateParams};
 
 #[cfg(test)]
 mod test_utils;

components/spider-storage/src/state/runtime.rs

@@ -12,7 +12,14 @@ use crate::{
     config::DatabaseConfig,
     db::{DbStorage, MariaDbStorageConnector, SessionManagement},
     ready_queue::{ReadyQueueConfig, ReadyQueueSender, ReadyQueueSenderHandle, create_ready_queue},
-    state::{JobCache, JobCacheGcConfig, ServiceState, StorageServerError, create_job_cache_gc},
+    state::{
+        JobCache,
+        JobCacheGcConfig,
+        ServiceState,
+        ServiceStateParams,
+        StorageServerError,
+        create_job_cache_gc,
+    },
     task_instance_pool::{
         TaskInstancePoolConfig,
         TaskInstancePoolConnector,
@@ -41,9 +48,9 @@ pub struct RuntimeConfig {
 /// * `DbConnectorType` - The database connector type.
 /// * `TaskInstancePoolConnectorType` - The task instance pool connector type.
 pub struct Runtime<
-    ReadyQueueSenderType: ReadyQueueSender,
-    DbConnectorType: DbStorage,
-    TaskInstancePoolConnectorType: TaskInstancePoolConnector,
+    ReadyQueueSenderType: ReadyQueueSender + 'static,
+    DbConnectorType: DbStorage + 'static,
+    TaskInstancePoolConnectorType: TaskInstancePoolConnector + 'static,
 > {
     service_state:
         ServiceState<ReadyQueueSenderType, DbConnectorType, TaskInstancePoolConnectorType>,
@@ -54,9 +61,9 @@ pub struct Runtime<
 }
 
 impl<
-    ReadyQueueSenderType: ReadyQueueSender,
-    DbConnectorType: DbStorage,
-    TaskInstancePoolConnectorType: TaskInstancePoolConnector,
+    ReadyQueueSenderType: ReadyQueueSender + 'static,
+    DbConnectorType: DbStorage + 'static,
+    TaskInstancePoolConnectorType: TaskInstancePoolConnector + 'static,
 > Runtime<ReadyQueueSenderType, DbConnectorType, TaskInstancePoolConnectorType>
 {
     /// Stops the runtime.
@@ -172,15 +179,16 @@ pub async fn create_runtime(
         &config.job_cache_gc_config,
     )
     .map_err(CacheError::from)?;
-    let service_state = ServiceState::new(
+    let service_state = ServiceState::new(ServiceStateParams {
         db,
         session_id,
         job_cache,
         ready_queue_sender,
         ready_queue_receiver,
         task_instance_pool_connector,
         job_cache_gc_handle,
-    );
+        cancellation_token: cancellation_token.clone(),
+    });
 
     Ok((
         Runtime {
@@ -257,6 +265,7 @@ mod tests {
         state::{
             JobCache,
             ServiceState,
+            ServiceStateParams,
             StorageServerError,
             test_utils::{MockDbConnector, MockTaskInstancePoolConnector},
         },
@@ -281,15 +290,16 @@ mod tests {
             &JobCacheGcConfig::default(),
         )
         .expect("job cache GC creation");
-        let service_state = ServiceState::new(
+        let service_state = ServiceState::new(ServiceStateParams {
             db,
             session_id,
             job_cache,
-            sender,
-            receiver,
-            MockTaskInstancePoolConnector,
+            ready_queue_sender: sender,
+            ready_queue_receiver: receiver,
+            task_instance_pool_connector: MockTaskInstancePoolConnector,
             job_cache_gc_handle,
-        );
+            cancellation_token: cancellation_token.clone(),
+        });
 
         // Wired with a real job cache GC task, which should always be terminated without errors.
         Runtime {