diff --git a/components/spider-storage/src/cache/error.rs b/components/spider-storage/src/cache/error.rs index 865b92079..e0101b9c9 100644 --- a/components/spider-storage/src/cache/error.rs +++ b/components/spider-storage/src/cache/error.rs @@ -75,12 +75,15 @@ 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 job cache GC 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..d9a02b316 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; @@ -21,21 +24,13 @@ use crate::{ /// * `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. +#[derive(Clone)] pub struct JobCache< ReadyQueueSenderType: ReadyQueueSender, DbConnectorType: InternalJobOrchestration, TaskInstancePoolConnectorType: TaskInstancePoolConnector, > { - jobs: RwLock< - HashMap< - JobId, - SharedJobControlBlock< - ReadyQueueSenderType, - DbConnectorType, - TaskInstancePoolConnectorType, - >, - >, - >, + jobs: SharedJobMap, } impl< @@ -48,7 +43,7 @@ impl< #[must_use] pub fn new() -> Self { Self { - jobs: RwLock::new(HashMap::new()), + jobs: Arc::new(RwLock::new(HashMap::new())), } } @@ -106,6 +101,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 @@ -132,6 +140,14 @@ impl< } } +type JobMap = HashMap< + JobId, + SharedJobControlBlock, +>; + +type SharedJobMap = + Arc>>; + #[cfg(test)] mod tests { use std::sync::Arc; @@ -225,6 +241,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, 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 new file mode 100644 index 000000000..6d99ce3d8 --- /dev/null +++ b/components/spider-storage/src/state/job_cache_gc.rs @@ -0,0 +1,348 @@ +//! 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. + 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 + /// + /// A new [`JobCacheGcHandle`] backed by the given channel sender. + pub(crate) const fn new(sender: UnboundedSender) -> Self { + Self { sender } + } +} + +/// 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, +} + +struct JobCacheGc< + ReadyQueueSenderType: ReadyQueueSender, + DbConnectorType: InternalJobOrchestration, + TaskInstancePoolConnectorType: TaskInstancePoolConnector, +> { + job_cache: JobCache, + terminated_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, + terminated_jobs: VecDeque::new(), + terminated_job_retention, + receiver, + } + } + + /// Runs the actor loop until cancellation or sender shutdown. + /// + /// # Errors + /// + /// 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! { + 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); + } + } + } + } + + /// Adds a terminated job to the actor-owned GC queue. + fn enqueue_terminated_job(&mut self, job_id: JobId) { + self.terminated_jobs.push_back(TerminatedJob { + job_id, + enqueued_at: Instant::now(), + }); + } + + /// 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 expired_job_ids = Vec::new(); + while let Some(terminated_job) = self.terminated_jobs.front() { + if now.duration_since(terminated_job.enqueued_at) < self.terminated_job_retention { + break; + } + tracing::info!( + 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 terminated jobs currently queued for retention. + #[cfg(test)] + fn get_num_queued_terminated_jobs(&self) -> usize { + self.terminated_jobs.len() + } +} + +#[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); + + 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.get_num_queued_terminated_jobs(), + 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 job_id = JobId::random(); + 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.get_num_queued_terminated_jobs(), + 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, + ); + gc.enqueue_terminated_job(JobId::random()); + gc.enqueue_terminated_job(JobId::random()); + let now = Instant::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.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 9350bce13..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; @@ -11,7 +12,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, @@ -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 @@ -36,6 +45,7 @@ pub struct Runtime< ServiceState, cancellation_token: CancellationToken, task_instance_pool_join_handle: JoinHandle>, + job_cache_gc_join_handle: JoinHandle>, stop_timeout: Duration, } @@ -47,30 +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 = &mut self.task_instance_pool_join_handle => { - 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))) + + 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(); - Err(StorageServerError::Stopping( - "task instance pool 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 @@ -100,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, @@ -112,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)?; @@ -130,6 +162,12 @@ pub async fn create_runtime( task_instance_pool_connector.clone(), ) .await?; + let (job_cache_gc_handle, job_cache_gc_join_handle) = create_job_cache_gc( + job_cache.clone(), + cancellation_token.clone(), + &config.job_cache_gc_config, + ) + .map_err(CacheError::from)?; let service_state = ServiceState::new( db, session_id, @@ -137,6 +175,7 @@ pub async fn create_runtime( ready_queue_sender, ready_queue_receiver, task_instance_pool_connector, + job_cache_gc_handle, ); Ok(( @@ -144,6 +183,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, @@ -230,19 +270,29 @@ 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, ); + // Wired with a real job cache GC task, which should always be terminated without errors. 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), } } @@ -306,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 9bdab838c..eb2e26a8b 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, }; @@ -62,7 +62,8 @@ impl< /// /// # Returns /// - /// A newly created [`ServiceState`] from its constituent parts. + /// A newly created [`ServiceState`] that notifies the GC actor when cached jobs terminate. + #[must_use] pub fn new( db: DbConnectorType, session_id: SessionId, @@ -70,6 +71,7 @@ impl< ready_queue_sender: ReadyQueueSenderType, ready_queue_receiver: ReadyQueueReceiverHandle, task_instance_pool_connector: TaskInstancePoolConnectorType, + job_cache_gc_handle: JobCacheGcHandle, ) -> Self { Self { inner: Arc::new(ServiceStateInner { @@ -79,6 +81,7 @@ impl< ready_queue_sender, ready_queue_receiver, task_instance_pool_connector, + job_cache_gc_handle, }), } } @@ -195,6 +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_for_gc_if_terminal(job_id, state); tracing::info!( job_id = ? job_id, "Job cancelled.", @@ -357,6 +361,7 @@ impl< let state = jcb .succeed_task_instance(task_instance_id, task_index, task_outputs) .await?; + self.enqueue_for_gc_if_terminal(job_id, state); tracing::info!( job_id = ? job_id, task_id = ? TaskId::Index(task_index), @@ -395,6 +400,7 @@ impl< .await .ok_or(StorageServerError::JobNotFound(job_id))?; let state = jcb.succeed_commit_task_instance(task_instance_id).await?; + self.enqueue_for_gc_if_terminal(job_id, state); tracing::info!( job_id = ? job_id, task_id = ? TaskId::Commit, @@ -433,6 +439,7 @@ impl< .await .ok_or(StorageServerError::JobNotFound(job_id))?; let state = jcb.succeed_cleanup_task_instance(task_instance_id).await?; + self.enqueue_for_gc_if_terminal(job_id, state); tracing::info!( job_id = ? job_id, task_id = ? TaskId::Cleanup, @@ -474,6 +481,7 @@ impl< let state = jcb .fail_task_instance(task_instance_id, task_id, error) .await?; + self.enqueue_for_gc_if_terminal(job_id, state); tracing::info!( job_id = ? job_id, task_id = ? task_id, @@ -657,6 +665,15 @@ impl< } Ok(()) } + + /// Enqueues a job for delayed cache GC if it has reached a terminal state. + 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); + } + } } /// Inner data for [`ServiceState`], holding all storage services. @@ -677,6 +694,7 @@ struct ServiceStateInner< ready_queue_sender: ReadyQueueSenderType, ready_queue_receiver: ReadyQueueReceiverHandle, task_instance_pool_connector: TaskInstancePoolConnectorType, + job_cache_gc_handle: JobCacheGcHandle, } #[cfg(test)] @@ -703,6 +721,7 @@ mod tests { db::DbError, ready_queue::ReadyQueueSenderHandle, state::{ + JobCacheGcHandle, StorageServerError, test_utils::{MockDbConnector, MockReadyQueueSender, MockTaskInstancePoolConnector}, }, @@ -728,19 +747,24 @@ 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, + JobCacheGcHandle::new(tokio::sync::mpsc::unbounded_channel().0), ) } + 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 @@ -759,6 +783,7 @@ mod tests { sender.clone(), receiver, MockTaskInstancePoolConnector, + JobCacheGcHandle::new(tokio::sync::mpsc::unbounded_channel().0), ); (service, sender) } @@ -1280,6 +1305,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( + 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.try_recv(), + Ok(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( + 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.try_recv(), + Ok(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( + 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. 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)) }