From 05af73a63a62a682ba377168507b5e5a0c6fe628 Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Fri, 12 Jun 2026 16:29:42 -0400 Subject: [PATCH 1/3] Add cache gc --- components/spider-storage/src/cache/error.rs | 3 + components/spider-storage/src/state.rs | 2 + .../spider-storage/src/state/job_cache.rs | 39 ++- .../spider-storage/src/state/job_cache_gc.rs | 328 ++++++++++++++++++ .../spider-storage/src/state/runtime.rs | 59 +++- .../spider-storage/src/state/service.rs | 188 +++++++++- 6 files changed, 593 insertions(+), 26 deletions(-) create mode 100644 components/spider-storage/src/state/job_cache_gc.rs diff --git a/components/spider-storage/src/cache/error.rs b/components/spider-storage/src/cache/error.rs index 865b92079..fa0631b75 100644 --- a/components/spider-storage/src/cache/error.rs +++ b/components/spider-storage/src/cache/error.rs @@ -81,6 +81,9 @@ pub enum InternalError { #[error("invalid config: {0}")] TaskInstancePoolInvalidConfig(&'static str), + #[error("invalid config: {0}")] + JobCacheGcInvalidConfig(&'static str), + #[error("ready queue channel is closed")] ReadyQueueChannelClosed, diff --git a/components/spider-storage/src/state.rs b/components/spider-storage/src/state.rs index 587f16ad2..4b76ec055 100644 --- a/components/spider-storage/src/state.rs +++ b/components/spider-storage/src/state.rs @@ -1,10 +1,12 @@ pub mod error; pub mod job_cache; +pub mod job_cache_gc; pub mod runtime; pub mod service; 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; diff --git a/components/spider-storage/src/state/job_cache.rs b/components/spider-storage/src/state/job_cache.rs index be38381eb..af67fbe52 100644 --- a/components/spider-storage/src/state/job_cache.rs +++ b/components/spider-storage/src/state/job_cache.rs @@ -1,4 +1,7 @@ -use std::collections::{HashMap, hash_map::Entry}; +use std::{ + collections::{HashMap, hash_map::Entry}, + sync::Arc, +}; use spider_core::types::id::JobId; use tokio::sync::RwLock; @@ -11,6 +14,14 @@ use crate::{ task_instance_pool::TaskInstancePoolConnector, }; +type JobMap = HashMap< + JobId, + SharedJobControlBlock, +>; + +type SharedJobMap = + Arc>>; + /// An in-memory cache for job control blocks. /// /// This type provides concurrent access to job control blocks via a [`tokio::sync::RwLock`] over a @@ -26,16 +37,7 @@ pub struct JobCache< DbConnectorType: InternalJobOrchestration, TaskInstancePoolConnectorType: TaskInstancePoolConnector, > { - jobs: RwLock< - HashMap< - JobId, - SharedJobControlBlock< - ReadyQueueSenderType, - DbConnectorType, - TaskInstancePoolConnectorType, - >, - >, - >, + jobs: SharedJobMap, } impl< @@ -48,7 +50,7 @@ impl< #[must_use] pub fn new() -> Self { Self { - jobs: RwLock::new(HashMap::new()), + jobs: Arc::new(RwLock::new(HashMap::new())), } } @@ -132,6 +134,19 @@ impl< } } +impl< + ReadyQueueSenderType: ReadyQueueSender, + DbConnectorType: InternalJobOrchestration, + TaskInstancePoolConnectorType: TaskInstancePoolConnector, +> Clone for JobCache +{ + fn clone(&self) -> Self { + Self { + jobs: Arc::clone(&self.jobs), + } + } +} + #[cfg(test)] mod tests { use std::sync::Arc; diff --git a/components/spider-storage/src/state/job_cache_gc.rs b/components/spider-storage/src/state/job_cache_gc.rs new file mode 100644 index 000000000..677e958bd --- /dev/null +++ b/components/spider-storage/src/state/job_cache_gc.rs @@ -0,0 +1,328 @@ +//! Background actor for removing terminated jobs from the in-memory job cache. + +use std::{ + collections::VecDeque, + time::{Duration, Instant}, +}; + +use spider_core::types::id::JobId; +use tokio::{ + sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel}, + task::JoinHandle, +}; +use tokio_util::sync::CancellationToken; + +use crate::{ + cache::error::InternalError, + db::InternalJobOrchestration, + ready_queue::ReadyQueueSender, + state::JobCache, + task_instance_pool::TaskInstancePoolConnector, +}; + +/// Configuration for the job-cache GC actor. +#[derive(Debug, Clone, Copy)] +pub struct JobCacheGcConfig { + /// Seconds to keep a terminated job in the in-memory cache before GC can remove it. + pub terminated_job_retention_sec: u64, + + /// Interval in seconds between GC cycles. + pub gc_interval_sec: u64, +} + +impl JobCacheGcConfig { + /// Validates the configuration parameters. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`InternalError::JobCacheGcInvalidConfig`] if `terminated_job_retention_sec` or + /// `gc_interval_sec` is zero. + pub const fn validate(&self) -> Result<(), InternalError> { + if self.terminated_job_retention_sec == 0 { + return Err(InternalError::JobCacheGcInvalidConfig( + "terminated_job_retention_sec must be greater than zero", + )); + } + if self.gc_interval_sec == 0 { + return Err(InternalError::JobCacheGcInvalidConfig( + "gc_interval_sec must be greater than zero", + )); + } + Ok(()) + } +} + +impl Default for JobCacheGcConfig { + fn default() -> Self { + Self { + terminated_job_retention_sec: 300, + gc_interval_sec: 30, + } + } +} + +/// Handle for enqueueing terminated jobs into the job-cache GC actor. +#[derive(Clone)] +pub struct JobCacheGcHandle { + sender: UnboundedSender, +} + +impl JobCacheGcHandle { + /// Enqueues a terminated job for delayed cache removal. + /// + /// # Errors + /// + /// Returns an error if the GC actor has stopped. + pub fn enqueue_terminated_job(&self, job_id: JobId) -> Result<(), JobId> { + self.sender.send(job_id).map_err(|e| e.0) + } + + /// # Returns + /// + /// A new [`JobCacheGcHandle`] backed by the given channel sender. + pub(crate) const fn new(sender: UnboundedSender) -> Self { + Self { sender } + } +} + +struct TerminatedJob { + job_id: JobId, + enqueued_at: Instant, +} + +struct JobCacheGc< + ReadyQueueSenderType: ReadyQueueSender, + DbConnectorType: InternalJobOrchestration, + TaskInstancePoolConnectorType: TaskInstancePoolConnector, +> { + job_cache: JobCache, + pending_jobs: VecDeque, + terminated_job_retention: Duration, + receiver: UnboundedReceiver, +} + +impl< + ReadyQueueSenderType: ReadyQueueSender, + DbConnectorType: InternalJobOrchestration, + TaskInstancePoolConnectorType: TaskInstancePoolConnector, +> JobCacheGc +{ + /// # Returns + /// + /// A new [`JobCacheGc`] actor over the given cache and message receiver. + const fn new( + job_cache: JobCache, + terminated_job_retention: Duration, + receiver: UnboundedReceiver, + ) -> Self { + Self { + job_cache, + pending_jobs: VecDeque::new(), + terminated_job_retention, + receiver, + } + } + + /// Runs the actor loop until cancellation or sender shutdown. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * No errors are returned by the current implementation. + async fn run( + mut self, + cancellation_token: CancellationToken, + gc_interval_sec: u64, + ) -> Result<(), InternalError> { + let mut gc_interval = tokio::time::interval(Duration::from_secs(gc_interval_sec)); + gc_interval.tick().await; + + loop { + tokio::select! { + () = cancellation_token.cancelled() => { + return Ok(()); + } + job_id = self.receiver.recv() => { + let Some(job_id) = job_id else { + return Ok(()); + }; + self.enqueue_terminated_job(job_id, Instant::now()); + } + _ = gc_interval.tick() => { + let _removed_jobs = self.run_gc_cycle_at(Instant::now()).await; + } + } + } + } + + /// Adds a terminated job to the actor-owned GC queue. + fn enqueue_terminated_job(&mut self, job_id: JobId, enqueued_at: Instant) { + self.pending_jobs.push_back(TerminatedJob { + job_id, + enqueued_at, + }); + } + + /// Removes expired terminated jobs from the cache. + /// + /// # Returns + /// + /// The number of expired entries processed by this GC cycle. + async fn run_gc_cycle_at(&mut self, now: Instant) -> usize { + let mut num_removed_jobs = 0; + while let Some(terminated_job) = self.pending_jobs.front() { + if now.duration_since(terminated_job.enqueued_at) < self.terminated_job_retention { + break; + } + let terminated_job = self + .pending_jobs + .pop_front() + .expect("pending terminated job should exist"); + self.job_cache.remove(terminated_job.job_id).await; + num_removed_jobs += 1; + tracing::info!( + job_id = ? terminated_job.job_id, + "Terminated job removed from cache.", + ); + } + num_removed_jobs + } + + /// # Returns + /// + /// The number of jobs currently queued for delayed GC. + #[cfg(test)] + fn pending_len(&self) -> usize { + self.pending_jobs.len() + } +} + +/// Creates a job-cache GC actor. +/// +/// # Returns +/// +/// A tuple containing the enqueue handle and spawned actor join handle on success. +/// +/// # Errors +/// +/// Returns an error if: +/// +/// * Forwards [`JobCacheGcConfig::validate`]'s return values on failure. +pub fn create_job_cache_gc< + ReadyQueueSenderType: ReadyQueueSender + 'static, + DbConnectorType: InternalJobOrchestration + 'static, + TaskInstancePoolConnectorType: TaskInstancePoolConnector + 'static, +>( + job_cache: JobCache, + cancellation_token: CancellationToken, + config: &JobCacheGcConfig, +) -> Result<(JobCacheGcHandle, JoinHandle>), InternalError> { + config.validate()?; + let (sender, receiver) = unbounded_channel(); + let gc_interval_sec = config.gc_interval_sec; + let gc = JobCacheGc::new( + job_cache, + Duration::from_secs(config.terminated_job_retention_sec), + receiver, + ); + let join_handle = + tokio::spawn(async move { gc.run(cancellation_token, gc_interval_sec).await }); + Ok((JobCacheGcHandle::new(sender), join_handle)) +} + +#[cfg(test)] +mod tests { + use std::time::{Duration, Instant}; + + use spider_core::types::id::JobId; + + use super::{JobCacheGc, JobCacheGcConfig}; + use crate::state::{ + JobCache, + test_utils::{MockDbConnector, MockReadyQueueSender, MockTaskInstancePoolConnector}, + }; + + type TestJobCache = + JobCache; + + #[test] + fn config_rejects_zero_values() { + let config = JobCacheGcConfig { + terminated_job_retention_sec: 0, + gc_interval_sec: 1, + }; + assert!( + config.validate().is_err(), + "zero retention should be invalid" + ); + + let config = JobCacheGcConfig { + terminated_job_retention_sec: 1, + gc_interval_sec: 0, + }; + assert!( + config.validate().is_err(), + "zero GC interval should be invalid" + ); + } + + #[tokio::test] + async fn gc_cycle_keeps_jobs_until_retention_expires() -> anyhow::Result<()> { + let cache = TestJobCache::new(); + let mut gc = JobCacheGc::new( + cache, + Duration::from_secs(10), + tokio::sync::mpsc::unbounded_channel().1, + ); + let now = Instant::now(); + let job_id = JobId::random(); + gc.enqueue_terminated_job(job_id, now); + + let removed = gc.run_gc_cycle_at(now + Duration::from_secs(9)).await; + + assert_eq!(removed, 0, "job should not be removed before retention"); + assert_eq!(gc.pending_len(), 1, "job should remain queued for GC"); + Ok(()) + } + + #[tokio::test] + async fn gc_cycle_removes_jobs_after_retention_expires() -> anyhow::Result<()> { + let cache = TestJobCache::new(); + let mut gc = JobCacheGc::new( + cache, + Duration::from_secs(10), + tokio::sync::mpsc::unbounded_channel().1, + ); + let now = Instant::now(); + let job_id = JobId::random(); + gc.enqueue_terminated_job(job_id, now); + + let removed = gc.run_gc_cycle_at(now + Duration::from_secs(10)).await; + + assert_eq!(removed, 1, "job should be removed after retention"); + assert_eq!(gc.pending_len(), 0, "job should no longer be queued"); + Ok(()) + } + + #[tokio::test] + async fn gc_cycle_removes_all_expired_jobs() -> anyhow::Result<()> { + let cache = TestJobCache::new(); + let mut gc = JobCacheGc::new( + cache, + Duration::from_secs(10), + tokio::sync::mpsc::unbounded_channel().1, + ); + let now = Instant::now(); + gc.enqueue_terminated_job(JobId::random(), now); + gc.enqueue_terminated_job(JobId::random(), now); + + let removed = gc.run_gc_cycle_at(now + Duration::from_secs(10)).await; + + assert_eq!(removed, 2, "all expired jobs should be removed"); + assert_eq!(gc.pending_len(), 0, "no expired job should remain queued"); + Ok(()) + } +} diff --git a/components/spider-storage/src/state/runtime.rs b/components/spider-storage/src/state/runtime.rs index 9350bce13..59862ae54 100644 --- a/components/spider-storage/src/state/runtime.rs +++ b/components/spider-storage/src/state/runtime.rs @@ -11,7 +11,7 @@ use crate::{ config::DatabaseConfig, db::{DbStorage, MariaDbStorageConnector, SessionManagement}, ready_queue::{ReadyQueueConfig, ReadyQueueSender, ReadyQueueSenderHandle, create_ready_queue}, - state::{JobCache, ServiceState, StorageServerError}, + state::{JobCache, JobCacheGcConfig, ServiceState, StorageServerError, create_job_cache_gc}, task_instance_pool::{ TaskInstancePoolConfig, TaskInstancePoolConnector, @@ -36,6 +36,7 @@ pub struct Runtime< ServiceState, cancellation_token: CancellationToken, task_instance_pool_join_handle: JoinHandle>, + job_cache_gc_join_handle: JoinHandle>, stop_timeout: Duration, } @@ -56,18 +57,27 @@ impl< pub async fn stop(mut self) -> Result<(), StorageServerError> { self.cancellation_token.cancel(); tokio::select! { - result = &mut self.task_instance_pool_join_handle => { + result = async { + let task_instance_pool_result = Self::wait_for_background_task( + "task instance pool", + &mut self.task_instance_pool_join_handle, + ) + .await; + let job_cache_gc_result = Self::wait_for_background_task( + "job cache GC", + &mut self.job_cache_gc_join_handle, + ) + .await; + task_instance_pool_result?; + job_cache_gc_result + } => { result - .map_err(|e| { - let msg = format!("task instance pool panic: {e}"); - CacheError::Internal(InternalError::TaskInstancePoolCorrupted(msg)) - })? - .map_err(|e| StorageServerError::Cache(CacheError::Internal(e))) } () = tokio::time::sleep(self.stop_timeout) => { self.task_instance_pool_join_handle.abort(); + self.job_cache_gc_join_handle.abort(); Err(StorageServerError::Stopping( - "task instance pool stop timed out".to_owned(), + "background task stop timed out".to_owned(), )) } } @@ -82,6 +92,26 @@ impl< ) -> ServiceState { self.service_state.clone() } + + /// Waits for a background task to stop. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`StorageServerError::Cache`] if the task terminated on error or panic. + async fn wait_for_background_task( + name: &'static str, + join_handle: &mut JoinHandle>, + ) -> Result<(), StorageServerError> { + join_handle + .await + .map_err(|e| { + let msg = format!("{name} panic: {e}"); + CacheError::Internal(InternalError::TaskInstancePoolCorrupted(msg)) + })? + .map_err(|e| StorageServerError::Cache(CacheError::Internal(e))) + } } /// Creates a storage server runtime from the given configurations. @@ -130,13 +160,20 @@ pub async fn create_runtime( task_instance_pool_connector.clone(), ) .await?; - let service_state = ServiceState::new( + let (job_cache_gc_handle, job_cache_gc_join_handle) = create_job_cache_gc( + job_cache.clone(), + cancellation_token.clone(), + &JobCacheGcConfig::default(), + ) + .map_err(CacheError::from)?; + let service_state = ServiceState::new_with_job_cache_gc( db, session_id, job_cache, ready_queue_sender, ready_queue_receiver, task_instance_pool_connector, + job_cache_gc_handle, ); Ok(( @@ -144,6 +181,7 @@ pub async fn create_runtime( service_state, cancellation_token: cancellation_token.clone(), task_instance_pool_join_handle, + job_cache_gc_join_handle, stop_timeout: Duration::from_secs(STOP_BACKGROUND_TASKS_TIMEOUT_SEC), }, cancellation_token, @@ -238,11 +276,14 @@ mod tests { receiver, MockTaskInstancePoolConnector, ); + let job_cache_gc_join_handle: JoinHandle> = + tokio::spawn(async { Ok(()) }); Runtime { service_state, cancellation_token, task_instance_pool_join_handle: mock_task_instance_pool_handle, + job_cache_gc_join_handle, stop_timeout: Duration::from_secs(stop_timeout_sec), } } diff --git a/components/spider-storage/src/state/service.rs b/components/spider-storage/src/state/service.rs index 9bdab838c..2c8a6d2d3 100644 --- a/components/spider-storage/src/state/service.rs +++ b/components/spider-storage/src/state/service.rs @@ -27,7 +27,7 @@ use crate::{ ReadyQueueReceiverHandle, ReadyQueueSender, }, - state::{JobCache, StorageServerError}, + state::{JobCache, JobCacheGcHandle, StorageServerError}, task_instance_pool::TaskInstancePoolConnector, }; @@ -79,6 +79,35 @@ impl< ready_queue_sender, ready_queue_receiver, task_instance_pool_connector, + job_cache_gc_handle: None, + }), + } + } + + /// Factory function with a job-cache GC enqueue handle. + /// + /// # Returns + /// + /// A newly created [`ServiceState`] that notifies the GC actor when cached jobs terminate. + #[must_use] + pub fn new_with_job_cache_gc( + 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, + ) -> Self { + Self { + inner: Arc::new(ServiceStateInner { + db, + session_id, + job_cache, + ready_queue_sender, + ready_queue_receiver, + task_instance_pool_connector, + job_cache_gc_handle: Some(job_cache_gc_handle), }), } } @@ -195,6 +224,7 @@ impl< pub async fn cancel_job(&self, job_id: JobId) -> Result { if let Some(jcb) = self.inner.job_cache.get(job_id).await { let state = jcb.cancel().await?; + self.enqueue_terminal_job_if_needed(job_id, state); tracing::info!( job_id = ? job_id, "Job cancelled.", @@ -357,6 +387,7 @@ impl< let state = jcb .succeed_task_instance(task_instance_id, task_index, task_outputs) .await?; + self.enqueue_terminal_job_if_needed(job_id, state); tracing::info!( job_id = ? job_id, task_id = ? TaskId::Index(task_index), @@ -395,6 +426,7 @@ impl< .await .ok_or(StorageServerError::JobNotFound(job_id))?; let state = jcb.succeed_commit_task_instance(task_instance_id).await?; + self.enqueue_terminal_job_if_needed(job_id, state); tracing::info!( job_id = ? job_id, task_id = ? TaskId::Commit, @@ -433,6 +465,7 @@ impl< .await .ok_or(StorageServerError::JobNotFound(job_id))?; let state = jcb.succeed_cleanup_task_instance(task_instance_id).await?; + self.enqueue_terminal_job_if_needed(job_id, state); tracing::info!( job_id = ? job_id, task_id = ? TaskId::Cleanup, @@ -474,6 +507,7 @@ impl< let state = jcb .fail_task_instance(task_instance_id, task_id, error) .await?; + self.enqueue_terminal_job_if_needed(job_id, state); tracing::info!( job_id = ? job_id, task_id = ? task_id, @@ -657,6 +691,22 @@ impl< } Ok(()) } + + /// Enqueues a job for delayed cache GC if it has reached a terminal state. + fn enqueue_terminal_job_if_needed(&self, job_id: JobId, state: JobState) { + if !state.is_terminal() { + return; + } + let Some(job_cache_gc_handle) = &self.inner.job_cache_gc_handle else { + return; + }; + if let Err(job_id) = job_cache_gc_handle.enqueue_terminated_job(job_id) { + tracing::warn!( + job_id = ? job_id, + "Failed to enqueue terminated job for cache GC.", + ); + } + } } /// Inner data for [`ServiceState`], holding all storage services. @@ -677,6 +727,7 @@ struct ServiceStateInner< ready_queue_sender: ReadyQueueSenderType, ready_queue_receiver: ReadyQueueReceiverHandle, task_instance_pool_connector: TaskInstancePoolConnectorType, + job_cache_gc_handle: Option, } #[cfg(test)] @@ -703,6 +754,7 @@ mod tests { db::DbError, ready_queue::ReadyQueueSenderHandle, state::{ + JobCacheGcHandle, StorageServerError, test_utils::{MockDbConnector, MockReadyQueueSender, MockTaskInstancePoolConnector}, }, @@ -728,19 +780,23 @@ mod tests { db: MockDbConnector, session_id: SessionId, ) -> TestServiceState { - use crate::ready_queue::{ReadyQueueConfig, create_ready_queue}; - let (_sender, receiver) = - create_ready_queue(&ReadyQueueConfig::default()).expect("ready queue creation"); TestServiceState::new( db, session_id, JobCache::new(), MockReadyQueueSender, - receiver, + create_ready_queue_receiver(), MockTaskInstancePoolConnector, ) } + fn create_ready_queue_receiver() -> ReadyQueueReceiverHandle { + use crate::ready_queue::{ReadyQueueConfig, create_ready_queue}; + let (_sender, receiver) = + create_ready_queue(&ReadyQueueConfig::default()).expect("ready queue creation"); + receiver + } + /// Creates a [`ServiceState`] backed by [`ReadyQueueReceiverHandle`]. /// /// # Returns @@ -1280,6 +1336,128 @@ mod tests { Ok(()) } + #[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_with_job_cache_gc( + MockDbConnector::default(), + TEST_SESSION_ID, + JobCache::new(), + MockReadyQueueSender, + create_ready_queue_receiver(), + MockTaskInstancePoolConnector, + JobCacheGcHandle::new(sender), + ); + let job_id = JobId::random(); + let jcb = create_test_jcb(job_id).await; + service.inner.job_cache.insert(jcb).await?; + + service.cancel_job(job_id).await?; + + assert_eq!( + receiver.recv().await, + Some(job_id), + "cancelled job should be enqueued for cache GC" + ); + Ok(()) + } + + #[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_with_job_cache_gc( + MockDbConnector::default(), + TEST_SESSION_ID, + JobCache::new(), + MockReadyQueueSender, + create_ready_queue_receiver(), + MockTaskInstancePoolConnector, + JobCacheGcHandle::new(sender), + ); + let (serialized_task_graph, serialized_inputs) = create_test_job_submission(); + let job_id = service + .register_job( + ResourceGroupId::random(), + serialized_task_graph, + serialized_inputs, + ) + .await?; + service.start_job(job_id).await?; + let context = service + .create_task_instance( + TEST_SESSION_ID, + job_id, + TaskId::Index(0), + ExecutionManagerId::random(), + ) + .await?; + + service + .succeed_task_instance( + TEST_SESSION_ID, + job_id, + context.task_instance_id, + 0, + create_test_serialized_outputs(), + ) + .await?; + + assert_eq!( + receiver.recv().await, + Some(job_id), + "succeeded job should be enqueued for cache GC" + ); + Ok(()) + } + + #[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_with_job_cache_gc( + MockDbConnector::default(), + TEST_SESSION_ID, + JobCache::new(), + MockReadyQueueSender, + create_ready_queue_receiver(), + MockTaskInstancePoolConnector, + JobCacheGcHandle::new(sender), + ); + let (serialized_task_graph, serialized_inputs) = create_test_job_submission(); + let job_id = service + .register_job( + ResourceGroupId::random(), + serialized_task_graph, + serialized_inputs, + ) + .await?; + service.start_job(job_id).await?; + let context = service + .create_task_instance( + TEST_SESSION_ID, + job_id, + TaskId::Index(0), + ExecutionManagerId::random(), + ) + .await?; + + service + .fail_task_instance( + TEST_SESSION_ID, + job_id, + context.task_instance_id, + TaskId::Index(0), + "test failure".to_owned(), + ) + .await?; + + assert_eq!( + receiver.recv().await, + Some(job_id), + "failed job should be enqueued for cache GC" + ); + Ok(()) + } + #[tokio::test] async fn task_instance_orchestration_return_stale_session_on_mismatch() -> anyhow::Result<()> { // Create a service with a higher session ID to simulate a server restart. From 8ddca52522677ea688691a72b342a750055a25f0 Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Fri, 12 Jun 2026 16:41:17 -0400 Subject: [PATCH 2/3] Add batch job removal --- .../spider-storage/src/state/job_cache.rs | 43 +++++++++++++++++++ .../spider-storage/src/state/job_cache_gc.rs | 16 ++++--- 2 files changed, 53 insertions(+), 6 deletions(-) diff --git a/components/spider-storage/src/state/job_cache.rs b/components/spider-storage/src/state/job_cache.rs index af67fbe52..5440c3f87 100644 --- a/components/spider-storage/src/state/job_cache.rs +++ b/components/spider-storage/src/state/job_cache.rs @@ -108,6 +108,19 @@ impl< self.jobs.write().await.remove(&job_id) } + /// Removes multiple job control blocks from the cache. + /// + /// # Returns + /// + /// The number of job control blocks that existed and were removed. + pub async fn remove_batch(&self, job_ids: &[JobId]) -> usize { + let mut jobs = self.jobs.write().await; + job_ids + .iter() + .filter(|job_id| jobs.remove(job_id).is_some()) + .count() + } + /// Resends all ready tasks for every job in the cache to the ready queue. /// /// # Errors @@ -240,6 +253,36 @@ mod tests { Ok(()) } + #[tokio::test] + async fn job_cache_remove_batch_removes_existing_jobs_once() -> anyhow::Result<()> { + let cache: JobCache = + JobCache::new(); + let first_job_id = JobId::random(); + let second_job_id = JobId::random(); + let missing_job_id = JobId::random(); + + cache.insert(create_test_jcb(first_job_id).await).await?; + cache.insert(create_test_jcb(second_job_id).await).await?; + + let num_removed_jobs = cache + .remove_batch(&[first_job_id, missing_job_id, second_job_id]) + .await; + + assert_eq!( + num_removed_jobs, 2, + "remove_batch should count only existing jobs" + ); + assert!( + cache.get(first_job_id).await.is_none(), + "first job should be removed" + ); + assert!( + cache.get(second_job_id).await.is_none(), + "second job should be removed" + ); + Ok(()) + } + #[tokio::test] async fn job_cache_get_returns_none_for_nonexistent_job() -> anyhow::Result<()> { let cache: JobCache = diff --git a/components/spider-storage/src/state/job_cache_gc.rs b/components/spider-storage/src/state/job_cache_gc.rs index 677e958bd..b73deb246 100644 --- a/components/spider-storage/src/state/job_cache_gc.rs +++ b/components/spider-storage/src/state/job_cache_gc.rs @@ -172,7 +172,7 @@ impl< /// /// The number of expired entries processed by this GC cycle. async fn run_gc_cycle_at(&mut self, now: Instant) -> usize { - let mut num_removed_jobs = 0; + let mut expired_job_ids = Vec::new(); while let Some(terminated_job) = self.pending_jobs.front() { if now.duration_since(terminated_job.enqueued_at) < self.terminated_job_retention { break; @@ -181,14 +181,18 @@ impl< .pending_jobs .pop_front() .expect("pending terminated job should exist"); - self.job_cache.remove(terminated_job.job_id).await; - num_removed_jobs += 1; + expired_job_ids.push(terminated_job.job_id); + } + let num_expired_jobs = expired_job_ids.len(); + let num_removed_jobs = self.job_cache.remove_batch(&expired_job_ids).await; + if num_expired_jobs > 0 { tracing::info!( - job_id = ? terminated_job.job_id, - "Terminated job removed from cache.", + num_expired_jobs, + num_removed_jobs, + "Terminated jobs removed from cache.", ); } - num_removed_jobs + num_expired_jobs } /// # Returns From ce5f53b2900ade777ec7dc4501849db169526aeb Mon Sep 17 00:00:00 2001 From: LinZhihao-723 Date: Sun, 14 Jun 2026 19:02:39 -0400 Subject: [PATCH 3/3] Reviewed. --- components/spider-storage/src/cache/error.rs | 6 +- .../spider-storage/src/state/job_cache.rs | 30 +--- .../spider-storage/src/state/job_cache_gc.rs | 166 ++++++++++-------- .../spider-storage/src/state/runtime.rs | 127 +++++++------- .../spider-storage/src/state/service.rs | 75 +++----- .../tests/runtime_recovery_test.rs | 114 +++++------- 6 files changed, 235 insertions(+), 283 deletions(-) diff --git a/components/spider-storage/src/cache/error.rs b/components/spider-storage/src/cache/error.rs index fa0631b75..e0101b9c9 100644 --- a/components/spider-storage/src/cache/error.rs +++ b/components/spider-storage/src/cache/error.rs @@ -75,13 +75,13 @@ pub enum InternalError { #[error("task instance pool corrupted: {0}")] TaskInstancePoolCorrupted(String), - #[error("invalid config: {0}")] + #[error("invalid ready-queue config: {0}")] ReadyQueueInvalidConfig(&'static str), - #[error("invalid config: {0}")] + #[error("invalid task instance pool config: {0}")] TaskInstancePoolInvalidConfig(&'static str), - #[error("invalid config: {0}")] + #[error("invalid job cache GC config: {0}")] JobCacheGcInvalidConfig(&'static str), #[error("ready queue channel is closed")] diff --git a/components/spider-storage/src/state/job_cache.rs b/components/spider-storage/src/state/job_cache.rs index 5440c3f87..d9a02b316 100644 --- a/components/spider-storage/src/state/job_cache.rs +++ b/components/spider-storage/src/state/job_cache.rs @@ -14,14 +14,6 @@ use crate::{ task_instance_pool::TaskInstancePoolConnector, }; -type JobMap = HashMap< - JobId, - SharedJobControlBlock, ->; - -type SharedJobMap = - Arc>>; - /// An in-memory cache for job control blocks. /// /// This type provides concurrent access to job control blocks via a [`tokio::sync::RwLock`] over a @@ -32,6 +24,7 @@ type SharedJobMap Clone for JobCache -{ - fn clone(&self) -> Self { - Self { - jobs: Arc::clone(&self.jobs), - } - } -} +type JobMap = HashMap< + JobId, + SharedJobControlBlock, +>; + +type SharedJobMap = + Arc>>; #[cfg(test)] mod tests { @@ -265,7 +253,7 @@ mod tests { cache.insert(create_test_jcb(second_job_id).await).await?; let num_removed_jobs = cache - .remove_batch(&[first_job_id, missing_job_id, second_job_id]) + .remove_batch(&[first_job_id, missing_job_id, second_job_id, second_job_id]) .await; assert_eq!( diff --git a/components/spider-storage/src/state/job_cache_gc.rs b/components/spider-storage/src/state/job_cache_gc.rs index b73deb246..6d99ce3d8 100644 --- a/components/spider-storage/src/state/job_cache_gc.rs +++ b/components/spider-storage/src/state/job_cache_gc.rs @@ -71,12 +71,10 @@ pub struct JobCacheGcHandle { impl JobCacheGcHandle { /// Enqueues a terminated job for delayed cache removal. - /// - /// # Errors - /// - /// Returns an error if the GC actor has stopped. - pub fn enqueue_terminated_job(&self, job_id: JobId) -> Result<(), JobId> { - self.sender.send(job_id).map_err(|e| e.0) + pub fn enqueue_terminated_job(&self, job_id: JobId) { + // Fire-and-forget: if the channel has been closed, just leave it. The GC coroutine is + // closed by a cancellation token. + let _ = self.sender.send(job_id); } /// # Returns @@ -87,6 +85,49 @@ impl JobCacheGcHandle { } } +/// Creates a job-cache GC actor. +/// +/// # Type Parameters +/// +/// * `ReadyQueueSenderType` - The type of the ready queue sender required by the job cache. +/// * `DbConnectorType` - The type of the DB-layer connector required by the job cache. +/// * `TaskInstancePoolConnectorType` - The type of the task instance pool connector required by the +/// job cache. +/// +/// # Returns +/// +/// A tuple on success, containing: +/// +/// * The handle for enqueueing terminated jobs into the GC actor. +/// * The join handle for the GC actor. +/// +/// # Errors +/// +/// Returns an error if: +/// +/// * Forwards [`JobCacheGcConfig::validate`]'s return values on failure. +pub fn create_job_cache_gc< + ReadyQueueSenderType: ReadyQueueSender + 'static, + DbConnectorType: InternalJobOrchestration + 'static, + TaskInstancePoolConnectorType: TaskInstancePoolConnector + 'static, +>( + job_cache: JobCache, + cancellation_token: CancellationToken, + config: &JobCacheGcConfig, +) -> Result<(JobCacheGcHandle, JoinHandle>), InternalError> { + config.validate()?; + let (sender, receiver) = unbounded_channel(); + let gc_interval_sec = config.gc_interval_sec; + let gc = JobCacheGc::new( + job_cache, + Duration::from_secs(config.terminated_job_retention_sec), + receiver, + ); + let join_handle = + tokio::spawn(async move { gc.run(cancellation_token, gc_interval_sec).await }); + Ok((JobCacheGcHandle::new(sender), join_handle)) +} + struct TerminatedJob { job_id: JobId, enqueued_at: Instant, @@ -98,7 +139,7 @@ struct JobCacheGc< TaskInstancePoolConnectorType: TaskInstancePoolConnector, > { job_cache: JobCache, - pending_jobs: VecDeque, + terminated_jobs: VecDeque, terminated_job_retention: Duration, receiver: UnboundedReceiver, } @@ -119,7 +160,7 @@ impl< ) -> Self { Self { job_cache, - pending_jobs: VecDeque::new(), + terminated_jobs: VecDeque::new(), terminated_job_retention, receiver, } @@ -129,9 +170,7 @@ impl< /// /// # Errors /// - /// Returns an error if: - /// - /// * No errors are returned by the current implementation. + /// No errors are returned by the current implementation. async fn run( mut self, cancellation_token: CancellationToken, @@ -142,27 +181,28 @@ impl< loop { tokio::select! { + biased; () = cancellation_token.cancelled() => { return Ok(()); } + _ = gc_interval.tick() => { + let _removed_jobs = self.run_gc_cycle_at(Instant::now()).await; + } job_id = self.receiver.recv() => { let Some(job_id) = job_id else { return Ok(()); }; - self.enqueue_terminated_job(job_id, Instant::now()); - } - _ = gc_interval.tick() => { - let _removed_jobs = self.run_gc_cycle_at(Instant::now()).await; + self.enqueue_terminated_job(job_id); } } } } /// Adds a terminated job to the actor-owned GC queue. - fn enqueue_terminated_job(&mut self, job_id: JobId, enqueued_at: Instant) { - self.pending_jobs.push_back(TerminatedJob { + fn enqueue_terminated_job(&mut self, job_id: JobId) { + self.terminated_jobs.push_back(TerminatedJob { job_id, - enqueued_at, + enqueued_at: Instant::now(), }); } @@ -173,70 +213,34 @@ impl< /// The number of expired entries processed by this GC cycle. async fn run_gc_cycle_at(&mut self, now: Instant) -> usize { let mut expired_job_ids = Vec::new(); - while let Some(terminated_job) = self.pending_jobs.front() { + while let Some(terminated_job) = self.terminated_jobs.front() { if now.duration_since(terminated_job.enqueued_at) < self.terminated_job_retention { break; } - let terminated_job = self - .pending_jobs - .pop_front() - .expect("pending terminated job should exist"); - expired_job_ids.push(terminated_job.job_id); - } - let num_expired_jobs = expired_job_ids.len(); - let num_removed_jobs = self.job_cache.remove_batch(&expired_job_ids).await; - if num_expired_jobs > 0 { tracing::info!( - num_expired_jobs, - num_removed_jobs, - "Terminated jobs removed from cache.", + job_id = % terminated_job.job_id, + "Terminated job expired, removing from cache." ); + expired_job_ids.push(terminated_job.job_id); + self.terminated_jobs.pop_front(); + } + if expired_job_ids.is_empty() { + return 0; } + let num_expired_jobs = expired_job_ids.len(); + self.job_cache.remove_batch(&expired_job_ids).await; num_expired_jobs } /// # Returns /// - /// The number of jobs currently queued for delayed GC. + /// The number of terminated jobs currently queued for retention. #[cfg(test)] - fn pending_len(&self) -> usize { - self.pending_jobs.len() + fn get_num_queued_terminated_jobs(&self) -> usize { + self.terminated_jobs.len() } } -/// Creates a job-cache GC actor. -/// -/// # Returns -/// -/// A tuple containing the enqueue handle and spawned actor join handle on success. -/// -/// # Errors -/// -/// Returns an error if: -/// -/// * Forwards [`JobCacheGcConfig::validate`]'s return values on failure. -pub fn create_job_cache_gc< - ReadyQueueSenderType: ReadyQueueSender + 'static, - DbConnectorType: InternalJobOrchestration + 'static, - TaskInstancePoolConnectorType: TaskInstancePoolConnector + 'static, ->( - job_cache: JobCache, - cancellation_token: CancellationToken, - config: &JobCacheGcConfig, -) -> Result<(JobCacheGcHandle, JoinHandle>), InternalError> { - config.validate()?; - let (sender, receiver) = unbounded_channel(); - let gc_interval_sec = config.gc_interval_sec; - let gc = JobCacheGc::new( - job_cache, - Duration::from_secs(config.terminated_job_retention_sec), - receiver, - ); - let join_handle = - tokio::spawn(async move { gc.run(cancellation_token, gc_interval_sec).await }); - Ok((JobCacheGcHandle::new(sender), join_handle)) -} - #[cfg(test)] mod tests { use std::time::{Duration, Instant}; @@ -283,12 +287,16 @@ mod tests { ); let now = Instant::now(); let job_id = JobId::random(); - gc.enqueue_terminated_job(job_id, now); + gc.enqueue_terminated_job(job_id); let removed = gc.run_gc_cycle_at(now + Duration::from_secs(9)).await; assert_eq!(removed, 0, "job should not be removed before retention"); - assert_eq!(gc.pending_len(), 1, "job should remain queued for GC"); + assert_eq!( + gc.get_num_queued_terminated_jobs(), + 1, + "job should remain queued for GC" + ); Ok(()) } @@ -300,14 +308,18 @@ mod tests { Duration::from_secs(10), tokio::sync::mpsc::unbounded_channel().1, ); - let now = Instant::now(); let job_id = JobId::random(); - gc.enqueue_terminated_job(job_id, now); + gc.enqueue_terminated_job(job_id); + let now = Instant::now(); let removed = gc.run_gc_cycle_at(now + Duration::from_secs(10)).await; assert_eq!(removed, 1, "job should be removed after retention"); - assert_eq!(gc.pending_len(), 0, "job should no longer be queued"); + assert_eq!( + gc.get_num_queued_terminated_jobs(), + 0, + "job should no longer be queued" + ); Ok(()) } @@ -319,14 +331,18 @@ mod tests { Duration::from_secs(10), tokio::sync::mpsc::unbounded_channel().1, ); + gc.enqueue_terminated_job(JobId::random()); + gc.enqueue_terminated_job(JobId::random()); let now = Instant::now(); - gc.enqueue_terminated_job(JobId::random(), now); - gc.enqueue_terminated_job(JobId::random(), now); let removed = gc.run_gc_cycle_at(now + Duration::from_secs(10)).await; assert_eq!(removed, 2, "all expired jobs should be removed"); - assert_eq!(gc.pending_len(), 0, "no expired job should remain queued"); + assert_eq!( + gc.get_num_queued_terminated_jobs(), + 0, + "no expired job should remain queued" + ); Ok(()) } } diff --git a/components/spider-storage/src/state/runtime.rs b/components/spider-storage/src/state/runtime.rs index 59862ae54..59d017c1c 100644 --- a/components/spider-storage/src/state/runtime.rs +++ b/components/spider-storage/src/state/runtime.rs @@ -1,5 +1,6 @@ use std::time::Duration; +use sqlx::__rt::timeout; use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; @@ -20,6 +21,14 @@ use crate::{ }, }; +/// Runtime configuration for the storage service. +pub struct RuntimeConfig { + pub db_config: DatabaseConfig, + pub ready_queue_config: ReadyQueueConfig, + pub task_instance_pool_config: TaskInstancePoolConfig, + pub job_cache_gc_config: JobCacheGcConfig, +} + /// Runtime state for the storage service. /// /// # Type Parameters @@ -48,39 +57,53 @@ impl< { /// Stops the runtime. /// + /// The background tasks will be cancelled and joined. The errors of the background tasks are + /// logged and will not be returned through this method. + /// /// # Errors /// /// Returns an error if: /// - /// * [`StorageServerError::Stopping`] if the task instance pool does not stop before timeout. - /// * [`StorageServerError::Cache`] if the task instance pool task terminated on error or panic. - pub async fn stop(mut self) -> Result<(), StorageServerError> { + /// * [`StorageServerError::Stopping`] if any of the background tasks does not stop before + /// timeout. + pub async fn stop(self) -> Result<(), StorageServerError> { self.cancellation_token.cancel(); - tokio::select! { - result = async { - let task_instance_pool_result = Self::wait_for_background_task( - "task instance pool", - &mut self.task_instance_pool_join_handle, - ) - .await; - let job_cache_gc_result = Self::wait_for_background_task( - "job cache GC", - &mut self.job_cache_gc_join_handle, - ) - .await; - task_instance_pool_result?; - job_cache_gc_result - } => { - result + + let join_task_instance_pool = async { + match self.task_instance_pool_join_handle.await { + Ok(Ok(())) => { + tracing::info!("Task instance pool stopped."); + } + Ok(Err(e)) => { + tracing::error!(error = ? e, "Task instance pool exited on error."); + } + Err(e) => { + tracing::error!(error = ? e, "Task instance pool exited on panic."); + } } - () = tokio::time::sleep(self.stop_timeout) => { - self.task_instance_pool_join_handle.abort(); - self.job_cache_gc_join_handle.abort(); - Err(StorageServerError::Stopping( - "background task stop timed out".to_owned(), - )) + }; + + let join_job_cache_gc = async { + match self.job_cache_gc_join_handle.await { + Ok(Ok(())) => { + tracing::info!("Job cache GC stopped."); + } + Ok(Err(e)) => { + tracing::error!(error = ? e, "Job cache GC exited on error."); + } + Err(e) => { + tracing::error!(error = ? e, "Job cache GC exited on panic."); + } } - } + }; + + let _ = timeout(self.stop_timeout, async { + tokio::join!(join_task_instance_pool, join_job_cache_gc,) + }) + .await + .map_err(|_| StorageServerError::Stopping("background task stop timed out".to_owned()))?; + + Ok(()) } /// # Returns @@ -92,26 +115,6 @@ impl< ) -> ServiceState { self.service_state.clone() } - - /// Waits for a background task to stop. - /// - /// # Errors - /// - /// Returns an error if: - /// - /// * [`StorageServerError::Cache`] if the task terminated on error or panic. - async fn wait_for_background_task( - name: &'static str, - join_handle: &mut JoinHandle>, - ) -> Result<(), StorageServerError> { - join_handle - .await - .map_err(|e| { - let msg = format!("{name} panic: {e}"); - CacheError::Internal(InternalError::TaskInstancePoolCorrupted(msg)) - })? - .map_err(|e| StorageServerError::Cache(CacheError::Internal(e))) - } } /// Creates a storage server runtime from the given configurations. @@ -130,10 +133,9 @@ impl< /// * Forwards [`MariaDbStorageConnector::connect`]'s return values on failure. /// * Forwards [`create_task_instance_pool`]'s return values on failure. /// * Forwards [`create_ready_queue`]'s return values on failure. +/// * Forwards [`create_job_cache_gc`]'s return values on failure. pub async fn create_runtime( - db_config: &DatabaseConfig, - ready_queue_config: &ReadyQueueConfig, - task_instance_pool_config: &TaskInstancePoolConfig, + config: &RuntimeConfig, ) -> Result< ( Runtime, @@ -142,15 +144,15 @@ pub async fn create_runtime( StorageServerError, > { let cancellation_token = CancellationToken::new(); - let db = MariaDbStorageConnector::connect(db_config).await?; + let db = MariaDbStorageConnector::connect(&config.db_config).await?; let session_id = db.session_id(); let (ready_queue_sender, ready_queue_receiver) = - create_ready_queue(ready_queue_config).map_err(CacheError::from)?; + create_ready_queue(&config.ready_queue_config).map_err(CacheError::from)?; let (task_instance_pool_connector, task_instance_pool_join_handle) = create_task_instance_pool( ready_queue_sender.clone(), db.clone(), cancellation_token.clone(), - task_instance_pool_config, + &config.task_instance_pool_config, ) .map_err(CacheError::from)?; @@ -163,10 +165,10 @@ pub async fn create_runtime( let (job_cache_gc_handle, job_cache_gc_join_handle) = create_job_cache_gc( job_cache.clone(), cancellation_token.clone(), - &JobCacheGcConfig::default(), + &config.job_cache_gc_config, ) .map_err(CacheError::from)?; - let service_state = ServiceState::new_with_job_cache_gc( + let service_state = ServiceState::new( db, session_id, job_cache, @@ -268,17 +270,24 @@ mod tests { let session_id = db.session_id(); let (sender, receiver) = create_ready_queue(&ReadyQueueConfig::default()).expect("ready queue creation"); + let job_cache = JobCache::new(); + let (job_cache_gc_handle, job_cache_gc_join_handle) = create_job_cache_gc( + job_cache.clone(), + cancellation_token.clone(), + &JobCacheGcConfig::default(), + ) + .expect("job cache GC creation"); let service_state = ServiceState::new( db, session_id, - JobCache::new(), + job_cache, sender, receiver, MockTaskInstancePoolConnector, + job_cache_gc_handle, ); - let job_cache_gc_join_handle: JoinHandle> = - tokio::spawn(async { Ok(()) }); + // Wired with a real job cache GC task, which should always be terminated without errors. Runtime { service_state, cancellation_token, @@ -347,8 +356,8 @@ mod tests { let result = runtime.stop().await; assert!( - matches!(result, Err(StorageServerError::Cache(_))), - "pool task failure should return Cache error" + matches!(result, Ok(())), + "pool task failure should not be forwarded as an error" ); Ok(()) } diff --git a/components/spider-storage/src/state/service.rs b/components/spider-storage/src/state/service.rs index 2c8a6d2d3..eb2e26a8b 100644 --- a/components/spider-storage/src/state/service.rs +++ b/components/spider-storage/src/state/service.rs @@ -62,35 +62,9 @@ impl< /// /// # Returns /// - /// A newly created [`ServiceState`] from its constituent parts. - pub fn new( - db: DbConnectorType, - session_id: SessionId, - job_cache: JobCache, - ready_queue_sender: ReadyQueueSenderType, - ready_queue_receiver: ReadyQueueReceiverHandle, - task_instance_pool_connector: TaskInstancePoolConnectorType, - ) -> Self { - Self { - inner: Arc::new(ServiceStateInner { - db, - session_id, - job_cache, - ready_queue_sender, - ready_queue_receiver, - task_instance_pool_connector, - job_cache_gc_handle: None, - }), - } - } - - /// Factory function with a job-cache GC enqueue handle. - /// - /// # Returns - /// /// A newly created [`ServiceState`] that notifies the GC actor when cached jobs terminate. #[must_use] - pub fn new_with_job_cache_gc( + pub fn new( db: DbConnectorType, session_id: SessionId, job_cache: JobCache, @@ -107,7 +81,7 @@ impl< ready_queue_sender, ready_queue_receiver, task_instance_pool_connector, - job_cache_gc_handle: Some(job_cache_gc_handle), + job_cache_gc_handle, }), } } @@ -224,7 +198,7 @@ impl< pub async fn cancel_job(&self, job_id: JobId) -> Result { if let Some(jcb) = self.inner.job_cache.get(job_id).await { let state = jcb.cancel().await?; - self.enqueue_terminal_job_if_needed(job_id, state); + self.enqueue_for_gc_if_terminal(job_id, state); tracing::info!( job_id = ? job_id, "Job cancelled.", @@ -387,7 +361,7 @@ impl< let state = jcb .succeed_task_instance(task_instance_id, task_index, task_outputs) .await?; - self.enqueue_terminal_job_if_needed(job_id, state); + self.enqueue_for_gc_if_terminal(job_id, state); tracing::info!( job_id = ? job_id, task_id = ? TaskId::Index(task_index), @@ -426,7 +400,7 @@ impl< .await .ok_or(StorageServerError::JobNotFound(job_id))?; let state = jcb.succeed_commit_task_instance(task_instance_id).await?; - self.enqueue_terminal_job_if_needed(job_id, state); + self.enqueue_for_gc_if_terminal(job_id, state); tracing::info!( job_id = ? job_id, task_id = ? TaskId::Commit, @@ -465,7 +439,7 @@ impl< .await .ok_or(StorageServerError::JobNotFound(job_id))?; let state = jcb.succeed_cleanup_task_instance(task_instance_id).await?; - self.enqueue_terminal_job_if_needed(job_id, state); + self.enqueue_for_gc_if_terminal(job_id, state); tracing::info!( job_id = ? job_id, task_id = ? TaskId::Cleanup, @@ -507,7 +481,7 @@ impl< let state = jcb .fail_task_instance(task_instance_id, task_id, error) .await?; - self.enqueue_terminal_job_if_needed(job_id, state); + self.enqueue_for_gc_if_terminal(job_id, state); tracing::info!( job_id = ? job_id, task_id = ? task_id, @@ -693,18 +667,11 @@ impl< } /// Enqueues a job for delayed cache GC if it has reached a terminal state. - fn enqueue_terminal_job_if_needed(&self, job_id: JobId, state: JobState) { - if !state.is_terminal() { - return; - } - let Some(job_cache_gc_handle) = &self.inner.job_cache_gc_handle else { - return; - }; - if let Err(job_id) = job_cache_gc_handle.enqueue_terminated_job(job_id) { - tracing::warn!( - job_id = ? job_id, - "Failed to enqueue terminated job for cache GC.", - ); + fn enqueue_for_gc_if_terminal(&self, job_id: JobId, state: JobState) { + if state.is_terminal() { + self.inner + .job_cache_gc_handle + .enqueue_terminated_job(job_id); } } } @@ -727,7 +694,7 @@ struct ServiceStateInner< ready_queue_sender: ReadyQueueSenderType, ready_queue_receiver: ReadyQueueReceiverHandle, task_instance_pool_connector: TaskInstancePoolConnectorType, - job_cache_gc_handle: Option, + job_cache_gc_handle: JobCacheGcHandle, } #[cfg(test)] @@ -787,6 +754,7 @@ mod tests { MockReadyQueueSender, create_ready_queue_receiver(), MockTaskInstancePoolConnector, + JobCacheGcHandle::new(tokio::sync::mpsc::unbounded_channel().0), ) } @@ -815,6 +783,7 @@ mod tests { sender.clone(), receiver, MockTaskInstancePoolConnector, + JobCacheGcHandle::new(tokio::sync::mpsc::unbounded_channel().0), ); (service, sender) } @@ -1339,7 +1308,7 @@ 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_with_job_cache_gc( + let service = TestServiceState::new( MockDbConnector::default(), TEST_SESSION_ID, JobCache::new(), @@ -1355,8 +1324,8 @@ mod tests { service.cancel_job(job_id).await?; assert_eq!( - receiver.recv().await, - Some(job_id), + receiver.try_recv(), + Ok(job_id), "cancelled job should be enqueued for cache GC" ); Ok(()) @@ -1365,7 +1334,7 @@ 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_with_job_cache_gc( + let service = TestServiceState::new( MockDbConnector::default(), TEST_SESSION_ID, JobCache::new(), @@ -1403,8 +1372,8 @@ mod tests { .await?; assert_eq!( - receiver.recv().await, - Some(job_id), + receiver.try_recv(), + Ok(job_id), "succeeded job should be enqueued for cache GC" ); Ok(()) @@ -1413,7 +1382,7 @@ 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_with_job_cache_gc( + let service = TestServiceState::new( MockDbConnector::default(), TEST_SESSION_ID, JobCache::new(), diff --git a/components/spider-storage/tests/runtime_recovery_test.rs b/components/spider-storage/tests/runtime_recovery_test.rs index 2a1860d7c..8f757a7f5 100644 --- a/components/spider-storage/tests/runtime_recovery_test.rs +++ b/components/spider-storage/tests/runtime_recovery_test.rs @@ -12,7 +12,14 @@ use spider_storage::{ cache::error::{CacheError, StaleStateError}, db::ExternalJobOrchestration, ready_queue::{CleanupTaskMarker, CommitTaskMarker, ReadyQueueConfig, ReadyQueueEntry}, - state::{Runtime, ServiceState, StorageServerError, create_runtime}, + state::{ + JobCacheGcConfig, + Runtime, + ServiceState, + StorageServerError, + create_runtime, + runtime::RuntimeConfig, + }, task_instance_pool::TaskInstancePoolConfig, }; use spider_tdl::wire::{TaskInputsSerializer, TaskOutputsSerializer}; @@ -26,24 +33,14 @@ use crate::{ #[ignore = "requires MariaDB"] #[serial_test::file_serial] async fn restarted_storage_cache_recovers_ready_job() -> anyhow::Result<()> { - let db_config = create_mariadb_config(); - let (runtime, _) = create_runtime( - &db_config, - &ReadyQueueConfig::default(), - &TaskInstancePoolConfig::default(), - ) - .await?; + let config = create_runtime_config(); + let (runtime, _) = create_runtime(&config).await?; let service = runtime.get_service_state(); let job_id = register_job(&service, false, false).await?; assert_eq!(service.get_job_state(job_id).await?, JobState::Ready); runtime.stop().await?; - let (recovered_runtime, _) = create_runtime( - &db_config, - &ReadyQueueConfig::default(), - &TaskInstancePoolConfig::default(), - ) - .await?; + let (recovered_runtime, _) = create_runtime(&config).await?; let recovered_service = recovered_runtime.get_service_state(); recovered_service.start_job(job_id).await?; assert_eq!( @@ -55,12 +52,7 @@ async fn restarted_storage_cache_recovers_ready_job() -> anyhow::Result<()> { recovered_runtime.stop().await?; // Create another runtime to test the job state and outputs are persisted. - let (recovered_runtime, _) = create_runtime( - &db_config, - &ReadyQueueConfig::default(), - &TaskInstancePoolConfig::default(), - ) - .await?; + let (recovered_runtime, _) = create_runtime(&config).await?; assert_job_outputs_on_success( &recovered_runtime.get_service_state(), job_id, @@ -76,8 +68,8 @@ async fn restarted_storage_cache_recovers_ready_job() -> anyhow::Result<()> { #[ignore = "requires MariaDB"] #[serial_test::file_serial] async fn restarted_storage_cache_recovers_running_job_from_start() -> anyhow::Result<()> { - let db_config = create_mariadb_config(); - let (job_id, recovered_runtime) = restart_after_starting_job(&db_config, false, false).await?; + let config = create_runtime_config(); + let (job_id, recovered_runtime) = restart_after_starting_job(&config, false, false).await?; let recovered_service = recovered_runtime.get_service_state(); recovered_service.resend_ready_tasks().await?; @@ -86,12 +78,7 @@ async fn restarted_storage_cache_recovers_running_job_from_start() -> anyhow::Re recovered_runtime.stop().await?; // Create another runtime to test the job state and outputs are persisted. - let (recovered_runtime, _) = create_runtime( - &db_config, - &ReadyQueueConfig::default(), - &TaskInstancePoolConfig::default(), - ) - .await?; + let (recovered_runtime, _) = create_runtime(&config).await?; assert_job_outputs_on_success( &recovered_runtime.get_service_state(), job_id, @@ -107,8 +94,8 @@ async fn restarted_storage_cache_recovers_running_job_from_start() -> anyhow::Re #[ignore = "requires MariaDB"] #[serial_test::file_serial] async fn restarted_storage_cache_recovers_commit_ready_job() -> anyhow::Result<()> { - let db_config = create_mariadb_config(); - let (job_id, recovered_runtime) = restart_after_commit_ready(&db_config).await?; + let config = create_runtime_config(); + let (job_id, recovered_runtime) = restart_after_commit_ready(&config).await?; let recovered_service = recovered_runtime.get_service_state(); recovered_service.resend_ready_tasks().await?; @@ -148,12 +135,7 @@ async fn restarted_storage_cache_recovers_commit_ready_job() -> anyhow::Result<( recovered_runtime.stop().await?; // Create another runtime to test the job state and outputs are persisted. - let (recovered_runtime, _) = create_runtime( - &db_config, - &ReadyQueueConfig::default(), - &TaskInstancePoolConfig::default(), - ) - .await?; + let (recovered_runtime, _) = create_runtime(&config).await?; assert_job_outputs_on_success( &recovered_runtime.get_service_state(), job_id, @@ -169,8 +151,8 @@ async fn restarted_storage_cache_recovers_commit_ready_job() -> anyhow::Result<( #[ignore = "requires MariaDB"] #[serial_test::file_serial] async fn restarted_storage_cache_recovers_cleanup_ready_job() -> anyhow::Result<()> { - let db_config = create_mariadb_config(); - let (job_id, recovered_runtime) = restart_after_cleanup_ready(&db_config).await?; + let config = create_runtime_config(); + let (job_id, recovered_runtime) = restart_after_cleanup_ready(&config).await?; let recovered_service = recovered_runtime.get_service_state(); recovered_service.resend_ready_tasks().await?; @@ -213,6 +195,19 @@ async fn restarted_storage_cache_recovers_cleanup_ready_job() -> anyhow::Result< Ok(()) } +/// # Returns +/// +/// A runtime configuration created for testing, with DB config given by [`create_mariadb_config`] +/// while other configurations set to default. +fn create_runtime_config() -> RuntimeConfig { + RuntimeConfig { + db_config: create_mariadb_config(), + ready_queue_config: ReadyQueueConfig::default(), + task_instance_pool_config: TaskInstancePoolConfig::default(), + job_cache_gc_config: JobCacheGcConfig::default(), + } +} + /// Starts a job, stops the runtime, and creates a replacement runtime over the same database. /// /// # Returns @@ -230,7 +225,7 @@ async fn restarted_storage_cache_recovers_cleanup_ready_job() -> anyhow::Result< /// * Forwards [`register_and_start_job`]'s return values on failure. /// * Forwards [`Runtime::stop`]'s return values on failure. async fn restart_after_starting_job( - db_config: &spider_storage::DatabaseConfig, + config: &RuntimeConfig, with_commit: bool, with_cleanup: bool, ) -> anyhow::Result<( @@ -241,22 +236,12 @@ async fn restart_after_starting_job( spider_storage::task_instance_pool::TaskInstancePoolHandle, >, )> { - let (runtime, _) = create_runtime( - db_config, - &ReadyQueueConfig::default(), - &TaskInstancePoolConfig::default(), - ) - .await?; + let (runtime, _) = create_runtime(config).await?; let service = runtime.get_service_state(); let job_id = register_and_start_job(&service, with_commit, with_cleanup).await?; runtime.stop().await?; - let (recovered_runtime, _) = create_runtime( - db_config, - &ReadyQueueConfig::default(), - &TaskInstancePoolConfig::default(), - ) - .await?; + let (recovered_runtime, _) = create_runtime(config).await?; Ok((job_id, recovered_runtime)) } @@ -281,7 +266,7 @@ async fn restart_after_starting_job( /// * Forwards [`Runtime::stop`]'s return values on failure. /// * Forwards [`create_runtime`]'s return values on failure. async fn restart_after_commit_ready( - db_config: &spider_storage::DatabaseConfig, + config: &RuntimeConfig, ) -> anyhow::Result<( JobId, Runtime< @@ -290,7 +275,7 @@ async fn restart_after_commit_ready( spider_storage::task_instance_pool::TaskInstancePoolHandle, >, )> { - let (job_id, runtime) = restart_after_starting_job(db_config, true, false).await?; + let (job_id, runtime) = restart_after_starting_job(config, true, false).await?; let service = runtime.get_service_state(); service.resend_ready_tasks().await?; @@ -313,12 +298,7 @@ async fn restart_after_commit_ready( assert_eq!(state, JobState::CommitReady); runtime.stop().await?; - let (recovered_runtime, _) = create_runtime( - db_config, - &ReadyQueueConfig::default(), - &TaskInstancePoolConfig::default(), - ) - .await?; + let (recovered_runtime, _) = create_runtime(config).await?; Ok((job_id, recovered_runtime)) } @@ -341,7 +321,7 @@ async fn restart_after_commit_ready( /// * Forwards [`ServiceState::cancel_job`]'s return values on failure. /// * Forwards [`Runtime::stop`]'s return values on failure. async fn restart_after_cleanup_ready( - db_config: &spider_storage::DatabaseConfig, + config: &RuntimeConfig, ) -> anyhow::Result<( JobId, Runtime< @@ -350,24 +330,14 @@ async fn restart_after_cleanup_ready( spider_storage::task_instance_pool::TaskInstancePoolHandle, >, )> { - let (runtime, _) = create_runtime( - db_config, - &ReadyQueueConfig::default(), - &TaskInstancePoolConfig::default(), - ) - .await?; + let (runtime, _) = create_runtime(config).await?; let service = runtime.get_service_state(); let job_id = register_and_start_job(&service, false, true).await?; let state = service.cancel_job(job_id).await?; assert_eq!(state, JobState::CleanupReady); runtime.stop().await?; - let (recovered_runtime, _) = create_runtime( - db_config, - &ReadyQueueConfig::default(), - &TaskInstancePoolConfig::default(), - ) - .await?; + let (recovered_runtime, _) = create_runtime(config).await?; Ok((job_id, recovered_runtime)) }