Repository navigation
feat(spider-storage): Add storage service runtime; Replace ExecutionManagerLivenessStore with ExecutionManagerLivenessManagement DB protocol in TaskInstancePool.
#324
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
sitaowang1998
merged 12 commits into
y-scope:storage-service-dev
from
sitaowang1998:graceful-shutdown
Jun 8, 2026
Merged
Changes from 7 commits
Commits
Show all changes
12 commits
Select commit
Hold shift + click to select a range
6d421d5
Add graceful shutdown & server runtime
sitaowang1998 fc6e11e
Merge two em liveness trait into one
sitaowang1998 6a3be1e
Use u64 instead of duration
sitaowang1998 6bba976
Fix test
sitaowang1998 5ccfa8d
Rename variable
sitaowang1998 e772967
Refactor server runtime
sitaowang1998 5b8eccf
Fix import
sitaowang1998 76dafa0
Address comment
sitaowang1998 cda4e6b
Review .
LinZhihao-723 cf7235a
Review task instance pool changes.
LinZhihao-723 4f721c7
Rename factory API.
LinZhihao-723 b7c8c70
Toml fix.
LinZhihao-723 File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,241 @@ | ||
| use std::time::Duration; | ||
|
|
||
| use tokio::task::JoinHandle; | ||
| use tokio_util::sync::CancellationToken; | ||
|
|
||
| use crate::{ | ||
| cache::error::{CacheError, InternalError}, | ||
| config::DatabaseConfig, | ||
| db::{DbStorage, MariaDbStorageConnector, SessionManagement}, | ||
| ready_queue::{ReadyQueueConfig, ReadyQueueSender, ReadyQueueSenderHandle, create_ready_queue}, | ||
| state::{JobCache, ServiceState, StorageServerError}, | ||
| task_instance_pool::{ | ||
| TaskInstancePoolConfig, | ||
| TaskInstancePoolConnector, | ||
| TaskInstancePoolHandle, | ||
| create_task_instance_pool, | ||
| }, | ||
| }; | ||
|
|
||
| /// Per-process storage server runtime. | ||
| /// | ||
| /// # Type Parameters | ||
| /// | ||
| /// * `ReadyQueueSenderType` - The ready queue sender type. | ||
| /// * `DbConnectorType` - The database connector type. | ||
| /// * `TaskInstancePoolConnectorType` - The task instance pool connector type. | ||
| pub struct ServerRuntime< | ||
| ReadyQueueSenderType: ReadyQueueSender, | ||
| DbConnectorType: DbStorage, | ||
| TaskInstancePoolConnectorType: TaskInstancePoolConnector, | ||
| > { | ||
| service_state: | ||
| ServiceState<ReadyQueueSenderType, DbConnectorType, TaskInstancePoolConnectorType>, | ||
| cancellation_token: CancellationToken, | ||
| task_instance_pool_join_handle: JoinHandle<Result<(), InternalError>>, | ||
| stop_timeout_sec: u64, | ||
| } | ||
|
|
||
| /// Creates a storage server runtime from the database configuration. | ||
| /// | ||
| /// # Returns | ||
| /// | ||
| /// A newly created [`ServerRuntime`] on success. | ||
| /// | ||
| /// # Errors | ||
| /// | ||
| /// Returns an error if: | ||
| /// | ||
| /// * Forwards [`MariaDbStorageConnector::connect`]'s return values on failure. | ||
| /// * Forwards [`create_ready_queue`]'s return values on failure. | ||
| pub async fn create_server_runtime( | ||
| db_config: &DatabaseConfig, | ||
| ) -> Result< | ||
| ServerRuntime<ReadyQueueSenderHandle, MariaDbStorageConnector, TaskInstancePoolHandle>, | ||
| StorageServerError, | ||
| > { | ||
| let cancellation_token = CancellationToken::new(); | ||
| let db = MariaDbStorageConnector::connect(db_config).await?; | ||
| let session_id = db.session_id(); | ||
| let (ready_queue_sender, ready_queue_receiver) = | ||
| create_ready_queue(ReadyQueueConfig::default()).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(), | ||
| TaskInstancePoolConfig::default(), | ||
| ); | ||
| let service_state = ServiceState::new( | ||
| db, | ||
| session_id, | ||
| JobCache::new(), | ||
| ready_queue_sender, | ||
| ready_queue_receiver, | ||
| task_instance_pool_connector, | ||
| ); | ||
|
|
||
| Ok(ServerRuntime { | ||
| service_state, | ||
| cancellation_token, | ||
| task_instance_pool_join_handle, | ||
| stop_timeout_sec: STOP_BACKGROUND_TASKS_TIMEOUT_SEC, | ||
| }) | ||
| } | ||
|
|
||
| impl<ReadyQueueSenderType, DbConnectorType, TaskInstancePoolConnectorType> | ||
| ServerRuntime<ReadyQueueSenderType, DbConnectorType, TaskInstancePoolConnectorType> | ||
| where | ||
| ReadyQueueSenderType: ReadyQueueSender, | ||
| DbConnectorType: DbStorage, | ||
| TaskInstancePoolConnectorType: TaskInstancePoolConnector, | ||
| { | ||
| /// Stops background tasks owned by the runtime. | ||
| /// | ||
| /// # 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 fails or cannot be joined. | ||
| pub async fn stop_background_tasks(self) -> Result<(), StorageServerError> { | ||
| self.cancellation_token.cancel(); | ||
| let join_result = tokio::time::timeout( | ||
| Duration::from_secs(self.stop_timeout_sec), | ||
| self.task_instance_pool_join_handle, | ||
| ) | ||
| .await | ||
| .map_err(|_| { | ||
| StorageServerError::Stopping("task instance pool stop timed out".to_owned()) | ||
| })?; | ||
| let pool_result = join_result.map_err(|e| { | ||
| StorageServerError::Cache(CacheError::Internal( | ||
| InternalError::TaskInstancePoolCorrupted(format!("task join error: {e}")), | ||
| )) | ||
| })?; | ||
| pool_result.map_err(|e| StorageServerError::Cache(CacheError::Internal(e))) | ||
| } | ||
|
|
||
| /// # Returns | ||
| /// | ||
| /// A clone of the runtime's [`ServiceState`]. | ||
| #[must_use] | ||
| pub fn service_state( | ||
| &self, | ||
| ) -> ServiceState<ReadyQueueSenderType, DbConnectorType, TaskInstancePoolConnectorType> { | ||
| self.service_state.clone() | ||
| } | ||
| } | ||
|
|
||
| const STOP_BACKGROUND_TASKS_TIMEOUT_SEC: u64 = 30; | ||
|
|
||
| #[cfg(test)] | ||
| mod tests { | ||
| use std::time::Duration; | ||
|
|
||
| use tokio::task::JoinHandle; | ||
| use tokio_util::sync::CancellationToken; | ||
|
|
||
| use super::ServerRuntime; | ||
| use crate::{ | ||
| cache::error::InternalError, | ||
| db::SessionManagement, | ||
| ready_queue::{ReadyQueueConfig, ReadyQueueSenderHandle, create_ready_queue}, | ||
| state::{ | ||
| JobCache, | ||
| ServiceState, | ||
| StorageServerError, | ||
| test_utils::{MockDbConnector, MockTaskInstancePoolConnector}, | ||
| }, | ||
| }; | ||
|
|
||
| type TestServerRuntime = | ||
| ServerRuntime<ReadyQueueSenderHandle, MockDbConnector, MockTaskInstancePoolConnector>; | ||
|
|
||
| fn create_test_server_runtime( | ||
| cancellation_token: CancellationToken, | ||
| task: JoinHandle<Result<(), InternalError>>, | ||
| stop_timeout_sec: u64, | ||
| ) -> TestServerRuntime { | ||
| let db = MockDbConnector::default(); | ||
| let session_id = db.session_id(); | ||
| let (sender, receiver) = | ||
| create_ready_queue(ReadyQueueConfig::default()).expect("ready queue creation"); | ||
| let service_state = ServiceState::new( | ||
| db, | ||
| session_id, | ||
| JobCache::new(), | ||
| sender, | ||
| receiver, | ||
| MockTaskInstancePoolConnector, | ||
| ); | ||
|
|
||
| ServerRuntime { | ||
| service_state, | ||
| cancellation_token, | ||
| task_instance_pool_join_handle: task, | ||
| stop_timeout_sec, | ||
| } | ||
| } | ||
|
|
||
| #[tokio::test] | ||
| async fn stop_background_tasks_cancels_and_joins_task() -> anyhow::Result<()> { | ||
| let cancellation_token = CancellationToken::new(); | ||
| let task_cancellation_token = cancellation_token.clone(); | ||
| let task: JoinHandle<Result<(), InternalError>> = tokio::spawn(async move { | ||
| task_cancellation_token.cancelled().await; | ||
| Ok(()) | ||
| }); | ||
|
|
||
| let runtime = create_test_server_runtime( | ||
| cancellation_token, | ||
| task, | ||
| super::STOP_BACKGROUND_TASKS_TIMEOUT_SEC, | ||
| ); | ||
| runtime | ||
| .stop_background_tasks() | ||
| .await | ||
| .expect("stop_background_tasks should succeed"); | ||
| Ok(()) | ||
| } | ||
|
|
||
| #[tokio::test] | ||
| async fn stop_background_tasks_returns_stopping_on_timeout() -> anyhow::Result<()> { | ||
| let cancellation_token = CancellationToken::new(); | ||
| let task: JoinHandle<Result<(), InternalError>> = tokio::spawn(async move { | ||
| tokio::time::sleep(Duration::from_secs(10)).await; | ||
| Ok(()) | ||
| }); | ||
|
|
||
| let runtime = create_test_server_runtime(cancellation_token, task, 0); | ||
| let result = runtime.stop_background_tasks().await; | ||
|
|
||
| assert!( | ||
| matches!(result, Err(StorageServerError::Stopping(_))), | ||
| "timeout should return Stopping" | ||
| ); | ||
| Ok(()) | ||
| } | ||
|
|
||
| #[tokio::test] | ||
| async fn stop_background_tasks_returns_cache_error_on_pool_error() -> anyhow::Result<()> { | ||
| let cancellation_token = CancellationToken::new(); | ||
| let task: JoinHandle<Result<(), InternalError>> = tokio::spawn(async move { | ||
| Err(InternalError::TaskInstancePoolCorrupted( | ||
| "test failure".to_owned(), | ||
| )) | ||
| }); | ||
|
|
||
| let runtime = create_test_server_runtime( | ||
| cancellation_token, | ||
| task, | ||
| super::STOP_BACKGROUND_TASKS_TIMEOUT_SEC, | ||
| ); | ||
| let result = runtime.stop_background_tasks().await; | ||
|
|
||
| assert!( | ||
| matches!(result, Err(StorageServerError::Cache(_))), | ||
| "pool task failure should return Cache error" | ||
| ); | ||
| Ok(()) | ||
| } | ||
| } | ||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.