From 4290c0a84580d81dcafafb5bfc057443bdd92529 Mon Sep 17 00:00:00 2001 From: LinZhihao-723 Date: Sun, 5 Jul 2026 17:15:29 -0400 Subject: [PATCH] Done. --- components/spider-storage/src/grpc.rs | 54 +++- components/spider-storage/src/state.rs | 2 +- .../spider-storage/src/state/runtime.rs | 38 ++- .../spider-storage/src/state/service.rs | 241 +++++++++++++----- .../spider-storage/src/state/test_utils.rs | 5 +- 5 files changed, 255 insertions(+), 85 deletions(-) diff --git a/components/spider-storage/src/grpc.rs b/components/spider-storage/src/grpc.rs index 36132eab4..ccbccc7a3 100644 --- a/components/spider-storage/src/grpc.rs +++ b/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, cancellation_token: CancellationToken, } impl< - ReadyQueueSenderType: ReadyQueueSender, - DbConnectorType: DbStorage, - TaskInstancePoolConnectorType: TaskInstancePoolConnector, + ReadyQueueSenderType: ReadyQueueSender + 'static, + DbConnectorType: DbStorage + 'static, + TaskInstancePoolConnectorType: TaskInstancePoolConnector + 'static, > GrpcServiceState { /// 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, + request: Request, ) -> Result, 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, ) -> Result, Status> { - todo!("Not implemented") + Err(Status::unimplemented("not implemented")) } } diff --git a/components/spider-storage/src/state.rs b/components/spider-storage/src/state.rs index 4b76ec055..7e83c191c 100644 --- a/components/spider-storage/src/state.rs +++ b/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; diff --git a/components/spider-storage/src/state/runtime.rs b/components/spider-storage/src/state/runtime.rs index 8d718e3ef..23f425ba7 100644 --- a/components/spider-storage/src/state/runtime.rs +++ b/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, @@ -54,9 +61,9 @@ pub struct Runtime< } impl< - ReadyQueueSenderType: ReadyQueueSender, - DbConnectorType: DbStorage, - TaskInstancePoolConnectorType: TaskInstancePoolConnector, + ReadyQueueSenderType: ReadyQueueSender + 'static, + DbConnectorType: DbStorage + 'static, + TaskInstancePoolConnectorType: TaskInstancePoolConnector + 'static, > Runtime { /// Stops the runtime. @@ -172,7 +179,7 @@ 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, @@ -180,7 +187,8 @@ pub async fn create_runtime( 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 { diff --git a/components/spider-storage/src/state/service.rs b/components/spider-storage/src/state/service.rs index 6d4d5d306..4ea7e9a64 100644 --- a/components/spider-storage/src/state/service.rs +++ b/components/spider-storage/src/state/service.rs @@ -18,6 +18,7 @@ use spider_core::{ }, }; use spider_tdl::error::TdlError; +use tokio_util::sync::CancellationToken; use crate::{ cache::{ @@ -37,6 +38,30 @@ use crate::{ task_instance_pool::TaskInstancePoolConnector, }; +/// Bundle of constructor parameters for [`ServiceState::new`]. +/// +/// This is a work-around for silencing the ` clippy::too_many_arguments ` warning. +/// +/// # Type Parameters +/// +/// * `ReadyQueueSenderType` - The type of the ready queue sender. +/// * `DbConnectorType` - The type of the DB-layer connector. +/// * `TaskInstancePoolConnectorType` - The type of the task instance pool connector. +pub struct ServiceStateParams< + ReadyQueueSenderType: ReadyQueueSender + 'static, + DbConnectorType: DbStorage + 'static, + TaskInstancePoolConnectorType: TaskInstancePoolConnector + 'static, +> { + pub db: DbConnectorType, + pub session_id: SessionId, + pub job_cache: JobCache, + pub ready_queue_sender: ReadyQueueSenderType, + pub ready_queue_receiver: ReadyQueueReceiverHandle, + pub task_instance_pool_connector: TaskInstancePoolConnectorType, + pub job_cache_gc_handle: JobCacheGcHandle, + pub cancellation_token: CancellationToken, +} + /// Per-request service state providing access to the storage layer. /// /// Internally wraps a single [`Arc`] around [`ServiceStateInner`] so that cloning is cheap (one @@ -49,9 +74,9 @@ use crate::{ /// * `TaskInstancePoolConnectorType` - The type of the task instance pool connector. #[derive(Clone)] pub struct ServiceState< - ReadyQueueSenderType: ReadyQueueSender, - DbConnectorType: DbStorage, - TaskInstancePoolConnectorType: TaskInstancePoolConnector, + ReadyQueueSenderType: ReadyQueueSender + 'static, + DbConnectorType: DbStorage + 'static, + TaskInstancePoolConnectorType: TaskInstancePoolConnector + 'static, > { inner: Arc< ServiceStateInner, @@ -59,9 +84,9 @@ pub struct ServiceState< } impl< - ReadyQueueSenderType: ReadyQueueSender, - DbConnectorType: DbStorage, - TaskInstancePoolConnectorType: TaskInstancePoolConnector, + ReadyQueueSenderType: ReadyQueueSender + 'static, + DbConnectorType: DbStorage + 'static, + TaskInstancePoolConnectorType: TaskInstancePoolConnector + 'static, > ServiceState { /// Factory function. @@ -71,14 +96,22 @@ impl< /// A newly created [`ServiceState`] that notifies the GC actor when cached jobs terminate. #[must_use] pub fn new( - db: DbConnectorType, - session_id: SessionId, - job_cache: JobCache, - ready_queue_sender: ReadyQueueSenderType, - ready_queue_receiver: ReadyQueueReceiverHandle, - task_instance_pool_connector: TaskInstancePoolConnectorType, - job_cache_gc_handle: JobCacheGcHandle, + params: ServiceStateParams< + ReadyQueueSenderType, + DbConnectorType, + TaskInstancePoolConnectorType, + >, ) -> Self { + let ServiceStateParams { + db, + session_id, + job_cache, + ready_queue_sender, + ready_queue_receiver, + task_instance_pool_connector, + job_cache_gc_handle, + cancellation_token, + } = params; Self { inner: Arc::new(ServiceStateInner { db, @@ -88,6 +121,8 @@ impl< ready_queue_receiver, task_instance_pool_connector, job_cache_gc_handle, + has_previous_scheduler_connection: tokio::sync::Mutex::new(false), + cancellation_token, }), } } @@ -654,9 +689,13 @@ impl< Ok(()) } - /// Registers the scheduler. + /// Registers a scheduler. + /// + /// Scheduler registration is mutually exclusive: only one registration request can be processed + /// at a time. Registering a scheduler invalidates any scheduler that was previously registered. /// - /// Registering a scheduler invalidates any previously registered scheduler. + /// If this replaces an existing scheduler, all ready tasks are re-enqueued in a background task + /// so they become visible in the inbound queue for the newly registered scheduler. /// /// # Returns /// @@ -673,6 +712,8 @@ impl< ip_address: IpAddr, port: u16, ) -> Result { + let mut has_previous_scheduler_connection = + self.inner.has_previous_scheduler_connection.lock().await; let scheduler_id = self.inner.db.register_scheduler(ip_address, port).await?; tracing::info!( scheduler_id = ? scheduler_id, @@ -680,6 +721,26 @@ impl< port, "Scheduler registered.", ); + if *has_previous_scheduler_connection { + tracing::info!( + "Previous scheduler connection has been invalidated. Resending all ready-tasks in \ + a background task." + ); + let job_cache = self.inner.job_cache.clone(); + let cancellation_token = self.inner.cancellation_token.clone(); + tokio::spawn(async move { + if let Err(e) = job_cache.resend_ready_tasks().await { + tracing::error!( + error = % e, + "Failed to resend ready-tasks after scheduler registration. Cancelling the \ + service." + ); + cancellation_token.cancel(); + } + }); + } + *has_previous_scheduler_connection = true; + drop(has_previous_scheduler_connection); Ok(scheduler_id) } @@ -733,9 +794,9 @@ impl< /// * `DbConnectorType` - The type of the DB-layer connector. /// * `TaskInstancePoolConnectorType` - The type of the task instance pool connector. struct ServiceStateInner< - ReadyQueueSenderType: ReadyQueueSender, - DbConnectorType: DbStorage, - TaskInstancePoolConnectorType: TaskInstancePoolConnector, + ReadyQueueSenderType: ReadyQueueSender + 'static, + DbConnectorType: DbStorage + 'static, + TaskInstancePoolConnectorType: TaskInstancePoolConnector + 'static, > { db: DbConnectorType, session_id: SessionId, @@ -744,6 +805,8 @@ struct ServiceStateInner< ready_queue_receiver: ReadyQueueReceiverHandle, task_instance_pool_connector: TaskInstancePoolConnectorType, job_cache_gc_handle: JobCacheGcHandle, + has_previous_scheduler_connection: tokio::sync::Mutex, + cancellation_token: CancellationToken, } #[cfg(test)] @@ -798,15 +861,16 @@ mod tests { db: MockDbConnector, session_id: SessionId, ) -> TestServiceState { - TestServiceState::new( + TestServiceState::new(ServiceStateParams { db, session_id, - JobCache::new(), - MockReadyQueueSender, - create_ready_queue_receiver(), - MockTaskInstancePoolConnector, - JobCacheGcHandle::new(tokio::sync::mpsc::unbounded_channel().0), - ) + job_cache: JobCache::new(), + ready_queue_sender: MockReadyQueueSender, + ready_queue_receiver: create_ready_queue_receiver(), + task_instance_pool_connector: MockTaskInstancePoolConnector, + job_cache_gc_handle: JobCacheGcHandle::new(tokio::sync::mpsc::unbounded_channel().0), + cancellation_token: CancellationToken::new(), + }) } fn create_ready_queue_receiver() -> ReadyQueueReceiverHandle { @@ -827,15 +891,16 @@ mod tests { use crate::ready_queue::{ReadyQueueConfig, create_ready_queue}; let (sender, receiver) = create_ready_queue(&ReadyQueueConfig::default()).expect("ready queue creation"); - let service = TestServiceStateWithReadyQueue::new( + let service = TestServiceStateWithReadyQueue::new(ServiceStateParams { db, - 0, - JobCache::new(), - sender.clone(), - receiver, - MockTaskInstancePoolConnector, - JobCacheGcHandle::new(tokio::sync::mpsc::unbounded_channel().0), - ); + session_id: 0, + job_cache: JobCache::new(), + ready_queue_sender: sender.clone(), + ready_queue_receiver: receiver, + task_instance_pool_connector: MockTaskInstancePoolConnector, + job_cache_gc_handle: JobCacheGcHandle::new(tokio::sync::mpsc::unbounded_channel().0), + cancellation_token: CancellationToken::new(), + }); (service, sender) } @@ -1367,15 +1432,16 @@ mod tests { #[tokio::test] async fn cancel_job_enqueues_terminal_job_for_cache_gc() -> anyhow::Result<()> { let (sender, mut receiver) = tokio::sync::mpsc::unbounded_channel(); - let service = TestServiceState::new( - MockDbConnector::default(), - TEST_SESSION_ID, - JobCache::new(), - MockReadyQueueSender, - create_ready_queue_receiver(), - MockTaskInstancePoolConnector, - JobCacheGcHandle::new(sender), - ); + let service = TestServiceState::new(ServiceStateParams { + db: MockDbConnector::default(), + session_id: TEST_SESSION_ID, + job_cache: JobCache::new(), + ready_queue_sender: MockReadyQueueSender, + ready_queue_receiver: create_ready_queue_receiver(), + task_instance_pool_connector: MockTaskInstancePoolConnector, + job_cache_gc_handle: JobCacheGcHandle::new(sender), + cancellation_token: CancellationToken::new(), + }); let job_id = JobId::random(); let jcb = create_test_jcb(job_id).await; service.inner.job_cache.insert(jcb).await?; @@ -1393,15 +1459,16 @@ mod tests { #[tokio::test] async fn succeed_task_instance_enqueues_terminal_job_for_cache_gc() -> anyhow::Result<()> { let (sender, mut receiver) = tokio::sync::mpsc::unbounded_channel(); - let service = TestServiceState::new( - MockDbConnector::default(), - TEST_SESSION_ID, - JobCache::new(), - MockReadyQueueSender, - create_ready_queue_receiver(), - MockTaskInstancePoolConnector, - JobCacheGcHandle::new(sender), - ); + let service = TestServiceState::new(ServiceStateParams { + db: MockDbConnector::default(), + session_id: TEST_SESSION_ID, + job_cache: JobCache::new(), + ready_queue_sender: MockReadyQueueSender, + ready_queue_receiver: create_ready_queue_receiver(), + task_instance_pool_connector: MockTaskInstancePoolConnector, + job_cache_gc_handle: JobCacheGcHandle::new(sender), + cancellation_token: CancellationToken::new(), + }); let (compressed_serialized_task_graph, compressed_serialized_inputs) = create_test_job_submission(); let job_id = service @@ -1442,15 +1509,16 @@ mod tests { #[tokio::test] async fn fail_task_instance_enqueues_terminal_job_for_cache_gc() -> anyhow::Result<()> { let (sender, mut receiver) = tokio::sync::mpsc::unbounded_channel(); - let service = TestServiceState::new( - MockDbConnector::default(), - TEST_SESSION_ID, - JobCache::new(), - MockReadyQueueSender, - create_ready_queue_receiver(), - MockTaskInstancePoolConnector, - JobCacheGcHandle::new(sender), - ); + let service = TestServiceState::new(ServiceStateParams { + db: MockDbConnector::default(), + session_id: TEST_SESSION_ID, + job_cache: JobCache::new(), + ready_queue_sender: MockReadyQueueSender, + ready_queue_receiver: create_ready_queue_receiver(), + task_instance_pool_connector: MockTaskInstancePoolConnector, + job_cache_gc_handle: JobCacheGcHandle::new(sender), + cancellation_token: CancellationToken::new(), + }); let (compressed_serialized_task_graph, compressed_serialized_inputs) = create_test_job_submission(); let job_id = service @@ -1680,6 +1748,59 @@ mod tests { Ok(()) } + #[tokio::test] + async fn register_scheduler_resends_ready_tasks_only_when_replacing_previous_scheduler() + -> anyhow::Result<()> { + let (service, _sender) = create_test_service_with_ready_queue(MockDbConnector::default()); + + let (task_graph, inputs) = create_test_job_submission(); + let job_id = service + .register_job(ResourceGroupId::random(), task_graph, inputs) + .await?; + service.start_job(job_id).await?; + + // Starting the job enqueues its initial ready task; drain it so the queue is empty before + // probing whether a registration triggers a resend. + let initial = service + .poll_ready_tasks(10, Duration::from_millis(100)) + .await?; + assert_eq!( + initial.len(), + 1, + "starting the job should enqueue its initial ready task" + ); + + // The first registration has no previous scheduler to replace, so it must not resend. + service + .register_scheduler("127.0.0.1".parse()?, 8080) + .await?; + tokio::task::yield_now().await; + let after_first = service + .poll_ready_tasks(10, Duration::from_millis(100)) + .await?; + assert!( + after_first.is_empty(), + "the first scheduler registration must not resend ready tasks" + ); + + // The second registration replaces the first scheduler, so it must resend ready tasks. The + // resend runs in a spawned background task, so yield to let it run before polling. + service + .register_scheduler("127.0.0.1".parse()?, 8081) + .await?; + tokio::task::yield_now().await; + let after_second = service + .poll_ready_tasks(10, Duration::from_millis(100)) + .await?; + assert_eq!( + after_second.len(), + 1, + "the second scheduler registration must resend the job's ready tasks" + ); + assert_eq!(after_second[0].job_id, job_id); + Ok(()) + } + #[tokio::test] async fn update_execution_manager_heartbeat_succeeds_for_registered_em() -> anyhow::Result<()> { let service = create_test_service(); diff --git a/components/spider-storage/src/state/test_utils.rs b/components/spider-storage/src/state/test_utils.rs index 491ec6be0..7cd4d9072 100644 --- a/components/spider-storage/src/state/test_utils.rs +++ b/components/spider-storage/src/state/test_utils.rs @@ -79,6 +79,7 @@ pub struct MockDbConnector { pub next_resource_group_id: Arc, pub execution_managers: Arc>, pub next_execution_manager_id: Arc, + pub next_scheduler_id: Arc, pub session_id: SessionId, } @@ -92,6 +93,7 @@ impl Default for MockDbConnector { next_resource_group_id: Arc::new(AtomicUsize::new(1)), execution_managers: Arc::new(DashMap::new()), next_execution_manager_id: Arc::new(AtomicUsize::new(1)), + next_scheduler_id: Arc::new(AtomicUsize::new(1)), session_id: 0, } } @@ -260,7 +262,8 @@ impl SchedulerRegistrationManagement for MockDbConnector { _ip_address: IpAddr, _port: u16, ) -> Result { - unreachable!("not implemented for mock connector") + let counter = self.next_scheduler_id.fetch_add(1, Ordering::Relaxed); + Ok(SchedulerId::from(counter as u64)) } async fn get_schedulers(&self) -> Result, DbError> {