From 27c07010888432d0c0f0caf4d820f530a8159cf4 Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Fri, 26 Jun 2026 15:44:11 -0400 Subject: [PATCH 01/21] Migrate storage gRPC error codes to tonic::Status --- .../src/client/grpc/liveness.rs | 155 ++++++--- .../src/generated/storage.rs | 322 +----------------- components/spider-proto/storage/storage.proto | 86 +---- .../src/storage_client/grpc.rs | 129 +++---- 4 files changed, 197 insertions(+), 495 deletions(-) diff --git a/components/spider-execution-manager/src/client/grpc/liveness.rs b/components/spider-execution-manager/src/client/grpc/liveness.rs index af5c0fb16..9eeecb172 100644 --- a/components/spider-execution-manager/src/client/grpc/liveness.rs +++ b/components/spider-execution-manager/src/client/grpc/liveness.rs @@ -6,12 +6,12 @@ use async_trait::async_trait; use spider_core::types::id::{ExecutionManagerId, SessionId}; use spider_proto_rust::storage::{ self, - execution_manager_liveness_error, execution_manager_liveness_service_client::ExecutionManagerLivenessServiceClient, - register_execution_manager_response, - update_execution_manager_heartbeat_response, }; -use tonic::transport::{Channel, Endpoint}; +use tonic::{ + Code, + transport::{Channel, Endpoint}, +}; use crate::client::liveness::{LivenessClient, LivenessResponseError, RegistrationResponse}; @@ -52,7 +52,7 @@ impl LivenessClient for GrpcLivenessClient { .clone() .register_execution_manager(request) .await - .map_err(to_transport_error)? + .map_err(|status| map_liveness_status(&status))? .into_inner(); register_response_to_result(response) @@ -70,28 +70,28 @@ impl LivenessClient for GrpcLivenessClient { .clone() .update_execution_manager_heartbeat(request) .await - .map_err(to_transport_error)? + .map_err(|status| map_liveness_status(&status))? .into_inner(); heartbeat_response_to_result(response) } } -impl From for LivenessResponseError { - fn from(error: storage::ExecutionManagerLivenessError) -> Self { - match execution_manager_liveness_error::ErrCode::try_from(error.err_code) { - Ok(execution_manager_liveness_error::ErrCode::MarkedDead) => Self::MarkedDead, - Ok(execution_manager_liveness_error::ErrCode::InvalidInput) => { - Self::IllegalId(error.message) - } - Ok( - execution_manager_liveness_error::ErrCode::Server - | execution_manager_liveness_error::ErrCode::Unspecified, - ) => Self::Transport(error.message), - Err(error) => Self::Transport(format!( - "unknown execution manager liveness error kind: {error}" - )), - } +/// Maps a [`tonic::Status`] returned by an execution-manager-liveness RPC into a +/// [`LivenessResponseError`]. +/// +/// # Returns +/// +/// * [`LivenessResponseError::MarkedDead`] when storage has already reaped the execution manager, +/// signalled by `FAILED_PRECONDITION`. +/// * [`LivenessResponseError::IllegalId`] when storage rejects the execution manager id, signalled +/// by `INVALID_ARGUMENT`. +/// * [`LivenessResponseError::Transport`] for any other failure. +fn map_liveness_status(status: &tonic::Status) -> LivenessResponseError { + match status.code() { + Code::FailedPrecondition => LivenessResponseError::MarkedDead, + Code::InvalidArgument => LivenessResponseError::IllegalId(status.message().to_owned()), + _ => LivenessResponseError::Transport(status.message().to_owned()), } } @@ -102,16 +102,20 @@ impl From for LivenessResponseError { fn register_response_to_result( response: storage::RegisterExecutionManagerResponse, ) -> Result { - match response.result { - Some(register_execution_manager_response::Result::Registration(registration)) => { + match response.registration { + Some(registration) => { + if registration.session_id == 0 { + return Err(LivenessResponseError::Transport( + "register execution manager response carried a zero session id".to_owned(), + )); + } Ok(RegistrationResponse { em_id: ExecutionManagerId::from(registration.execution_manager_id), session_id: registration.session_id, }) } - Some(register_execution_manager_response::Result::Error(error)) => Err(error.into()), None => Err(LivenessResponseError::Transport( - "register execution manager response missing result".to_owned(), + "register execution manager response missing registration".to_owned(), )), } } @@ -123,17 +127,13 @@ fn register_response_to_result( fn heartbeat_response_to_result( response: storage::UpdateExecutionManagerHeartbeatResponse, ) -> Result { - match response.result { - Some(update_execution_manager_heartbeat_response::Result::SessionId(session_id)) => { - Ok(session_id) - } - Some(update_execution_manager_heartbeat_response::Result::Error(error)) => { - Err(error.into()) - } - None => Err(LivenessResponseError::Transport( - "update execution manager heartbeat response missing result".to_owned(), - )), + let session_id = response.session_id; + if session_id == 0 { + return Err(LivenessResponseError::Transport( + "update execution manager heartbeat response carried a zero session id".to_owned(), + )); } + Ok(session_id) } /// Converts a displayable transport-layer error into [`LivenessResponseError::Transport`]. @@ -158,12 +158,10 @@ mod tests { const EM_ID: ExecutionManagerId = ExecutionManagerId::from(5); let response = storage::RegisterExecutionManagerResponse { - result: Some(register_execution_manager_response::Result::Registration( - storage::ExecutionManagerRegistration { - execution_manager_id: EM_ID.get(), - session_id: SESSION_ID, - }, - )), + registration: Some(storage::ExecutionManagerRegistration { + execution_manager_id: EM_ID.get(), + session_id: SESSION_ID, + }), }; let registration = register_response_to_result(response) @@ -178,14 +176,37 @@ mod tests { ); } + #[test] + fn register_response_to_result_rejects_missing_registration() { + let response = storage::RegisterExecutionManagerResponse { registration: None }; + + assert!(matches!( + register_response_to_result(response), + Err(LivenessResponseError::Transport(_)) + )); + } + + #[test] + fn register_response_to_result_rejects_zero_session_id() { + let response = storage::RegisterExecutionManagerResponse { + registration: Some(storage::ExecutionManagerRegistration { + execution_manager_id: 5, + session_id: 0, + }), + }; + + assert!(matches!( + register_response_to_result(response), + Err(LivenessResponseError::Transport(_)) + )); + } + #[test] fn heartbeat_response_to_result_returns_session_id() { const SESSION_ID: SessionId = 9; let response = storage::UpdateExecutionManagerHeartbeatResponse { - result: Some( - update_execution_manager_heartbeat_response::Result::SessionId(SESSION_ID), - ), + session_id: SESSION_ID, }; let session_id = heartbeat_response_to_result(response) @@ -195,19 +216,43 @@ mod tests { } #[test] - fn liveness_storage_error_maps_invalid_input_to_illegal_id() { - const ERROR_MSG: &str = "bad em id"; + fn heartbeat_response_to_result_rejects_zero_session_id() { + let response = storage::UpdateExecutionManagerHeartbeatResponse { session_id: 0 }; - let error = storage::ExecutionManagerLivenessError { - err_code: execution_manager_liveness_error::ErrCode::InvalidInput.into(), - message: ERROR_MSG.to_owned(), - }; + assert!(matches!( + heartbeat_response_to_result(response), + Err(LivenessResponseError::Transport(_)) + )); + } - match LivenessResponseError::from(error) { - LivenessResponseError::IllegalId(message) => { - assert_eq!(message, ERROR_MSG); - } - error => panic!("unexpected liveness response error: {error:?}"), + #[test] + fn map_liveness_status_maps_failed_precondition_to_marked_dead() { + let status = tonic::Status::failed_precondition("already dead"); + + assert!(matches!( + map_liveness_status(&status), + LivenessResponseError::MarkedDead + )); + } + + #[test] + fn map_liveness_status_maps_invalid_argument_to_illegal_id() { + const ERROR_MSG: &str = "bad em id"; + let status = tonic::Status::invalid_argument(ERROR_MSG); + + match map_liveness_status(&status) { + LivenessResponseError::IllegalId(message) => assert_eq!(message, ERROR_MSG), + error => panic!("unexpected liveness status mapping: {error:?}"), } } + + #[test] + fn map_liveness_status_maps_other_codes_to_transport() { + let status = tonic::Status::internal("boom"); + + assert!(matches!( + map_liveness_status(&status), + LivenessResponseError::Transport(_) + )); + } } diff --git a/components/spider-proto-rust/src/generated/storage.rs b/components/spider-proto-rust/src/generated/storage.rs index a690a2bc1..4b9453458 100644 --- a/components/spider-proto-rust/src/generated/storage.rs +++ b/components/spider-proto-rust/src/generated/storage.rs @@ -42,18 +42,8 @@ pub struct PollReadyTasksRequest { } #[derive(Clone, PartialEq, ::prost::Message)] pub struct PollReadyTasksResponse { - #[prost(oneof = "poll_ready_tasks_response::Result", tags = "1, 2")] - pub result: ::core::option::Option, -} -/// Nested message and enum types in `PollReadyTasksResponse`. -pub mod poll_ready_tasks_response { - #[derive(Clone, PartialEq, ::prost::Oneof)] - pub enum Result { - #[prost(message, tag = "1")] - Tasks(super::ReadyTasks), - #[prost(message, tag = "2")] - Error(super::InboundQueueResponseError), - } + #[prost(message, optional, tag = "1")] + pub tasks: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct ReadyTasks { @@ -148,23 +138,11 @@ pub struct AddResourceGroupRequest { pub external_resource_group_id: ::prost::alloc::string::String, #[prost(bytes = "vec", tag = "2")] pub password: ::prost::alloc::vec::Vec, - #[prost(uint64, tag = "3")] - pub session_id: u64, } -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Clone, Copy, PartialEq, ::prost::Message)] pub struct ResourceGroupIdResponse { - #[prost(oneof = "resource_group_id_response::Result", tags = "1, 2")] - pub result: ::core::option::Option, -} -/// Nested message and enum types in `ResourceGroupIdResponse`. -pub mod resource_group_id_response { - #[derive(Clone, PartialEq, ::prost::Oneof)] - pub enum Result { - #[prost(uint64, tag = "1")] - ResourceGroupId(u64), - #[prost(message, tag = "2")] - Error(super::ResourceGroupManagementError), - } + #[prost(uint64, tag = "1")] + pub resource_group_id: u64, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct VerifyResourceGroupRequest { @@ -172,8 +150,6 @@ pub struct VerifyResourceGroupRequest { pub resource_group_id: u64, #[prost(bytes = "vec", tag = "2")] pub password: ::prost::alloc::vec::Vec, - #[prost(uint64, tag = "3")] - pub session_id: u64, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct RegisterExecutionManagerRequest { @@ -192,40 +168,15 @@ pub struct ExecutionManagerRegistration { #[prost(uint64, tag = "2")] pub session_id: u64, } -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Clone, Copy, PartialEq, ::prost::Message)] pub struct RegisterExecutionManagerResponse { - #[prost(oneof = "register_execution_manager_response::Result", tags = "1, 2")] - pub result: ::core::option::Option, -} -/// Nested message and enum types in `RegisterExecutionManagerResponse`. -pub mod register_execution_manager_response { - #[derive(Clone, PartialEq, ::prost::Oneof)] - pub enum Result { - #[prost(message, tag = "1")] - Registration(super::ExecutionManagerRegistration), - #[prost(message, tag = "2")] - Error(super::ExecutionManagerLivenessError), - } + #[prost(message, optional, tag = "1")] + pub registration: ::core::option::Option, } -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Clone, Copy, PartialEq, ::prost::Message)] pub struct UpdateExecutionManagerHeartbeatResponse { - #[prost( - oneof = "update_execution_manager_heartbeat_response::Result", - tags = "1, 2" - )] - pub result: ::core::option::Option< - update_execution_manager_heartbeat_response::Result, - >, -} -/// Nested message and enum types in `UpdateExecutionManagerHeartbeatResponse`. -pub mod update_execution_manager_heartbeat_response { - #[derive(Clone, PartialEq, ::prost::Oneof)] - pub enum Result { - #[prost(uint64, tag = "1")] - SessionId(u64), - #[prost(message, tag = "2")] - Error(super::ExecutionManagerLivenessError), - } + #[prost(uint64, tag = "1")] + pub session_id: u64, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct RegisterSchedulerRequest { @@ -241,20 +192,10 @@ pub struct SchedulerRegistration { #[prost(uint64, tag = "2")] pub session_id: u64, } -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Clone, Copy, PartialEq, ::prost::Message)] pub struct RegisterSchedulerResponse { - #[prost(oneof = "register_scheduler_response::Result", tags = "1, 2")] - pub result: ::core::option::Option, -} -/// Nested message and enum types in `RegisterSchedulerResponse`. -pub mod register_scheduler_response { - #[derive(Clone, PartialEq, ::prost::Oneof)] - pub enum Result { - #[prost(message, tag = "1")] - Registration(super::SchedulerRegistration), - #[prost(message, tag = "2")] - Error(super::SchedulerRegistrationError), - } + #[prost(message, optional, tag = "1")] + pub registration: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct Scheduler { @@ -272,18 +213,8 @@ pub struct SchedulerRegistrations { } #[derive(Clone, PartialEq, ::prost::Message)] pub struct GetSchedulersResponse { - #[prost(oneof = "get_schedulers_response::Result", tags = "1, 2")] - pub result: ::core::option::Option, -} -/// Nested message and enum types in `GetSchedulersResponse`. -pub mod get_schedulers_response { - #[derive(Clone, PartialEq, ::prost::Oneof)] - pub enum Result { - #[prost(message, tag = "1")] - Schedulers(super::SchedulerRegistrations), - #[prost(message, tag = "2")] - Error(super::SchedulerRegistrationError), - } + #[prost(message, optional, tag = "1")] + pub schedulers: ::core::option::Option, } #[derive(Clone, Copy, PartialEq, ::prost::Message)] pub struct GetSessionResponse { @@ -309,227 +240,10 @@ pub mod task_id { } #[derive(Clone, Copy, PartialEq, ::prost::Message)] pub struct TaskInstanceOperationResponse {} -#[derive(Clone, PartialEq, ::prost::Message)] -pub struct ResourceGroupOperationResponse { - #[prost(oneof = "resource_group_operation_response::Result", tags = "1, 2")] - pub result: ::core::option::Option, -} -/// Nested message and enum types in `ResourceGroupOperationResponse`. -pub mod resource_group_operation_response { - #[derive(Clone, PartialEq, ::prost::Oneof)] - pub enum Result { - #[prost(message, tag = "1")] - Ok(super::Void), - #[prost(message, tag = "2")] - Error(super::ResourceGroupManagementError), - } -} +#[derive(Clone, Copy, PartialEq, ::prost::Message)] +pub struct ResourceGroupOperationResponse {} #[derive(Clone, Copy, PartialEq, ::prost::Message)] pub struct Void {} -#[derive(Clone, PartialEq, ::prost::Message)] -pub struct InboundQueueResponseError { - #[prost(enumeration = "inbound_queue_response_error::ErrCode", tag = "1")] - pub err_code: i32, - #[prost(string, tag = "2")] - pub message: ::prost::alloc::string::String, -} -/// Nested message and enum types in `InboundQueueResponseError`. -pub mod inbound_queue_response_error { - #[derive( - Clone, - Copy, - Debug, - PartialEq, - Eq, - Hash, - PartialOrd, - Ord, - ::prost::Enumeration - )] - #[repr(i32)] - pub enum ErrCode { - Unspecified = 0, - InboundClosed = 1, - Server = 2, - InvalidInput = 3, - } - impl ErrCode { - /// String value of the enum field names used in the ProtoBuf definition. - /// - /// The values are not transformed in any way and thus are considered stable - /// (if the ProtoBuf definition does not change) and safe for programmatic use. - pub fn as_str_name(&self) -> &'static str { - match self { - Self::Unspecified => "ERR_CODE_UNSPECIFIED", - Self::InboundClosed => "INBOUND_CLOSED", - Self::Server => "SERVER", - Self::InvalidInput => "INVALID_INPUT", - } - } - /// Creates an enum from field names used in the ProtoBuf definition. - pub fn from_str_name(value: &str) -> ::core::option::Option { - match value { - "ERR_CODE_UNSPECIFIED" => Some(Self::Unspecified), - "INBOUND_CLOSED" => Some(Self::InboundClosed), - "SERVER" => Some(Self::Server), - "INVALID_INPUT" => Some(Self::InvalidInput), - _ => None, - } - } - } -} -#[derive(Clone, PartialEq, ::prost::Message)] -pub struct ResourceGroupManagementError { - #[prost(enumeration = "resource_group_management_error::ErrCode", tag = "1")] - pub err_code: i32, - #[prost(string, tag = "2")] - pub message: ::prost::alloc::string::String, - #[prost(uint64, tag = "3")] - pub storage_session: u64, -} -/// Nested message and enum types in `ResourceGroupManagementError`. -pub mod resource_group_management_error { - #[derive( - Clone, - Copy, - Debug, - PartialEq, - Eq, - Hash, - PartialOrd, - Ord, - ::prost::Enumeration - )] - #[repr(i32)] - pub enum ErrCode { - Unspecified = 0, - StaleSession = 1, - Server = 2, - InvalidInput = 3, - } - impl ErrCode { - /// String value of the enum field names used in the ProtoBuf definition. - /// - /// The values are not transformed in any way and thus are considered stable - /// (if the ProtoBuf definition does not change) and safe for programmatic use. - pub fn as_str_name(&self) -> &'static str { - match self { - Self::Unspecified => "ERR_CODE_UNSPECIFIED", - Self::StaleSession => "STALE_SESSION", - Self::Server => "SERVER", - Self::InvalidInput => "INVALID_INPUT", - } - } - /// Creates an enum from field names used in the ProtoBuf definition. - pub fn from_str_name(value: &str) -> ::core::option::Option { - match value { - "ERR_CODE_UNSPECIFIED" => Some(Self::Unspecified), - "STALE_SESSION" => Some(Self::StaleSession), - "SERVER" => Some(Self::Server), - "INVALID_INPUT" => Some(Self::InvalidInput), - _ => None, - } - } - } -} -#[derive(Clone, PartialEq, ::prost::Message)] -pub struct ExecutionManagerLivenessError { - #[prost(enumeration = "execution_manager_liveness_error::ErrCode", tag = "1")] - pub err_code: i32, - #[prost(string, tag = "2")] - pub message: ::prost::alloc::string::String, -} -/// Nested message and enum types in `ExecutionManagerLivenessError`. -pub mod execution_manager_liveness_error { - #[derive( - Clone, - Copy, - Debug, - PartialEq, - Eq, - Hash, - PartialOrd, - Ord, - ::prost::Enumeration - )] - #[repr(i32)] - pub enum ErrCode { - Unspecified = 0, - MarkedDead = 1, - InvalidInput = 2, - Server = 3, - } - impl ErrCode { - /// String value of the enum field names used in the ProtoBuf definition. - /// - /// The values are not transformed in any way and thus are considered stable - /// (if the ProtoBuf definition does not change) and safe for programmatic use. - pub fn as_str_name(&self) -> &'static str { - match self { - Self::Unspecified => "ERR_CODE_UNSPECIFIED", - Self::MarkedDead => "MARKED_DEAD", - Self::InvalidInput => "INVALID_INPUT", - Self::Server => "SERVER", - } - } - /// Creates an enum from field names used in the ProtoBuf definition. - pub fn from_str_name(value: &str) -> ::core::option::Option { - match value { - "ERR_CODE_UNSPECIFIED" => Some(Self::Unspecified), - "MARKED_DEAD" => Some(Self::MarkedDead), - "INVALID_INPUT" => Some(Self::InvalidInput), - "SERVER" => Some(Self::Server), - _ => None, - } - } - } -} -#[derive(Clone, PartialEq, ::prost::Message)] -pub struct SchedulerRegistrationError { - #[prost(enumeration = "scheduler_registration_error::ErrCode", tag = "1")] - pub err_code: i32, - #[prost(string, tag = "2")] - pub message: ::prost::alloc::string::String, -} -/// Nested message and enum types in `SchedulerRegistrationError`. -pub mod scheduler_registration_error { - #[derive( - Clone, - Copy, - Debug, - PartialEq, - Eq, - Hash, - PartialOrd, - Ord, - ::prost::Enumeration - )] - #[repr(i32)] - pub enum ErrCode { - Unspecified = 0, - Server = 1, - } - impl ErrCode { - /// String value of the enum field names used in the ProtoBuf definition. - /// - /// The values are not transformed in any way and thus are considered stable - /// (if the ProtoBuf definition does not change) and safe for programmatic use. - pub fn as_str_name(&self) -> &'static str { - match self { - Self::Unspecified => "ERR_CODE_UNSPECIFIED", - Self::Server => "SERVER", - } - } - /// Creates an enum from field names used in the ProtoBuf definition. - pub fn from_str_name(value: &str) -> ::core::option::Option { - match value { - "ERR_CODE_UNSPECIFIED" => Some(Self::Unspecified), - "SERVER" => Some(Self::Server), - _ => None, - } - } - } -} #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] #[repr(i32)] pub enum JobState { diff --git a/components/spider-proto/storage/storage.proto b/components/spider-proto/storage/storage.proto index 3437a4f67..cfc3e2ce6 100644 --- a/components/spider-proto/storage/storage.proto +++ b/components/spider-proto/storage/storage.proto @@ -76,10 +76,7 @@ message PollReadyTasksRequest { } message PollReadyTasksResponse { - oneof result { - ReadyTasks tasks = 1; - InboundQueueResponseError error = 2; - } + ReadyTasks tasks = 1; } message ReadyTasks { @@ -142,20 +139,15 @@ message ReportTaskFailureRequest { message AddResourceGroupRequest { string external_resource_group_id = 1; bytes password = 2; - uint64 session_id = 3; } message ResourceGroupIdResponse { - oneof result { - uint64 resource_group_id = 1; - ResourceGroupManagementError error = 2; - } + uint64 resource_group_id = 1; } message VerifyResourceGroupRequest { uint64 resource_group_id = 1; bytes password = 2; - uint64 session_id = 3; } message RegisterExecutionManagerRequest { @@ -172,17 +164,11 @@ message ExecutionManagerRegistration { } message RegisterExecutionManagerResponse { - oneof result { - ExecutionManagerRegistration registration = 1; - ExecutionManagerLivenessError error = 2; - } + ExecutionManagerRegistration registration = 1; } message UpdateExecutionManagerHeartbeatResponse { - oneof result { - uint64 session_id = 1; - ExecutionManagerLivenessError error = 2; - } + uint64 session_id = 1; } message RegisterSchedulerRequest { @@ -196,10 +182,7 @@ message SchedulerRegistration { } message RegisterSchedulerResponse { - oneof result { - SchedulerRegistration registration = 1; - SchedulerRegistrationError error = 2; - } + SchedulerRegistration registration = 1; } message Scheduler { @@ -213,10 +196,7 @@ message SchedulerRegistrations { } message GetSchedulersResponse { - oneof result { - SchedulerRegistrations schedulers = 1; - SchedulerRegistrationError error = 2; - } + SchedulerRegistrations schedulers = 1; } message GetSessionResponse { @@ -244,58 +224,6 @@ enum JobState { message TaskInstanceOperationResponse {} -message ResourceGroupOperationResponse { - oneof result { - Void ok = 1; - ResourceGroupManagementError error = 2; - } -} +message ResourceGroupOperationResponse {} message Void {} - -message InboundQueueResponseError { - enum ErrCode { - ERR_CODE_UNSPECIFIED = 0; - INBOUND_CLOSED = 1; - SERVER = 2; - INVALID_INPUT = 3; - } - - ErrCode err_code = 1; - string message = 2; -} - -message ResourceGroupManagementError { - enum ErrCode { - ERR_CODE_UNSPECIFIED = 0; - STALE_SESSION = 1; - SERVER = 2; - INVALID_INPUT = 3; - } - - ErrCode err_code = 1; - string message = 2; - uint64 storage_session = 3; -} - -message ExecutionManagerLivenessError { - enum ErrCode { - ERR_CODE_UNSPECIFIED = 0; - MARKED_DEAD = 1; - INVALID_INPUT = 2; - SERVER = 3; - } - - ErrCode err_code = 1; - string message = 2; -} - -message SchedulerRegistrationError { - enum ErrCode { - ERR_CODE_UNSPECIFIED = 0; - SERVER = 1; - } - - ErrCode err_code = 1; - string message = 2; -} diff --git a/components/spider-scheduler/src/storage_client/grpc.rs b/components/spider-scheduler/src/storage_client/grpc.rs index 78a6b1d5f..1d97c8109 100644 --- a/components/spider-scheduler/src/storage_client/grpc.rs +++ b/components/spider-scheduler/src/storage_client/grpc.rs @@ -9,10 +9,8 @@ use spider_core::{ }; use spider_proto_rust::storage::{ self, - inbound_queue_response_error, inbound_queue_service_client::InboundQueueServiceClient, job_orchestration_service_client::JobOrchestrationServiceClient, - poll_ready_tasks_response, }; use tonic::{ Code, @@ -67,7 +65,7 @@ impl SchedulerStorageClient for GrpcSchedulerStorageClient { .clone() .poll_ready_tasks(request) .await - .map_err(to_transport_error)? + .map_err(|status| map_inbound_status(&status))? .into_inner(); poll_ready_tasks_response_to_result(response) } @@ -83,7 +81,7 @@ impl SchedulerStorageClient for GrpcSchedulerStorageClient { .clone() .poll_ready_commit_tasks(request) .await - .map_err(to_transport_error)? + .map_err(|status| map_inbound_status(&status))? .into_inner(); poll_ready_tasks_response_to_result(response) } @@ -99,7 +97,7 @@ impl SchedulerStorageClient for GrpcSchedulerStorageClient { .clone() .poll_ready_cleanup_tasks(request) .await - .map_err(to_transport_error)? + .map_err(|status| map_inbound_status(&status))? .into_inner(); poll_ready_tasks_response_to_result(response) } @@ -122,19 +120,19 @@ impl SchedulerStorageClient for GrpcSchedulerStorageClient { } } -impl From for StorageClientError { - fn from(error: storage::InboundQueueResponseError) -> Self { - match inbound_queue_response_error::ErrCode::try_from(error.err_code) { - Ok(inbound_queue_response_error::ErrCode::InboundClosed) => Self::InboundClosed, - Ok(inbound_queue_response_error::ErrCode::InvalidInput) => { - Self::InvalidInput(error.message) - } - Ok( - inbound_queue_response_error::ErrCode::Server - | inbound_queue_response_error::ErrCode::Unspecified, - ) => Self::Server(error.message), - Err(error) => Self::Transport(format!("unknown scheduler storage error kind: {error}")), - } +/// Maps a [`tonic::Status`] returned by an inbound-queue RPC into a [`StorageClientError`]. +/// +/// # Returns +/// +/// * [`StorageClientError::InboundClosed`] when the inbound queue is closed, signalled by +/// `UNAVAILABLE`. +/// * [`StorageClientError::InvalidInput`] for a malformed request, signalled by `INVALID_ARGUMENT`. +/// * [`StorageClientError::Server`] for any other failure. +fn map_inbound_status(status: &tonic::Status) -> StorageClientError { + match status.code() { + Code::Unavailable => StorageClientError::InboundClosed, + Code::InvalidArgument => to_invalid_input_error(status.message()), + _ => StorageClientError::Server(status.message().to_owned()), } } @@ -164,13 +162,12 @@ fn poll_ready_tasks_request( fn poll_ready_tasks_response_to_result( response: storage::PollReadyTasksResponse, ) -> Result<(SessionId, Vec), StorageClientError> { - match response.result { - Some(poll_ready_tasks_response::Result::Tasks(tasks)) => ready_tasks_to_result(tasks), - Some(poll_ready_tasks_response::Result::Error(error)) => Err(error.into()), - None => Err(StorageClientError::Transport( - "poll ready tasks response missing `result` message".to_owned(), - )), - } + let tasks = response.tasks.ok_or_else(|| { + StorageClientError::Transport( + "poll ready tasks response missing `tasks` message".to_owned(), + ) + })?; + ready_tasks_to_result(tasks) } /// # Returns @@ -254,11 +251,7 @@ fn to_invalid_input_error(error: impl std::fmt::Display) -> StorageClientError { #[cfg(test)] mod tests { use spider_core::types::id::{JobId, ResourceGroupId, TaskId}; - use spider_proto_rust::storage::{ - self, - inbound_queue_response_error, - poll_ready_tasks_response, - }; + use spider_proto_rust::storage; use super::*; @@ -270,16 +263,14 @@ mod tests { #[test] fn poll_ready_tasks_response_converts_entries() { let response = storage::PollReadyTasksResponse { - result: Some(poll_ready_tasks_response::Result::Tasks( - storage::ReadyTasks { - session_id: SESSION_ID, - tasks: vec![storage::ReadyTask { - resource_group_id: RESOURCE_GROUP_ID, - job_id: JOB_ID, - task_id: Some(storage::TaskId::from(TaskId::Index(TASK_INDEX))), - }], - }, - )), + tasks: Some(storage::ReadyTasks { + session_id: SESSION_ID, + tasks: vec![storage::ReadyTask { + resource_group_id: RESOURCE_GROUP_ID, + job_id: JOB_ID, + task_id: Some(storage::TaskId::from(TaskId::Index(TASK_INDEX))), + }], + }), }; let (session_id, entries) = poll_ready_tasks_response_to_result(response) @@ -301,16 +292,14 @@ mod tests { const MISSING_TASK_ID_MESSAGE: &str = "missing task ID"; let response = storage::PollReadyTasksResponse { - result: Some(poll_ready_tasks_response::Result::Tasks( - storage::ReadyTasks { - session_id: SESSION_ID, - tasks: vec![storage::ReadyTask { - resource_group_id: RESOURCE_GROUP_ID, - job_id: JOB_ID, - task_id: None, - }], - }, - )), + tasks: Some(storage::ReadyTasks { + session_id: SESSION_ID, + tasks: vec![storage::ReadyTask { + resource_group_id: RESOURCE_GROUP_ID, + job_id: JOB_ID, + task_id: None, + }], + }), }; match poll_ready_tasks_response_to_result(response) { @@ -322,17 +311,43 @@ mod tests { } #[test] - fn inbound_queue_response_error_maps_inbound_closed() { - const ERROR_MESSAGE: &str = "closed"; + fn poll_ready_tasks_response_rejects_missing_tasks() { + let response = storage::PollReadyTasksResponse { tasks: None }; - let error = storage::InboundQueueResponseError { - err_code: inbound_queue_response_error::ErrCode::InboundClosed.into(), - message: ERROR_MESSAGE.to_owned(), - }; + assert!(matches!( + poll_ready_tasks_response_to_result(response), + Err(StorageClientError::Transport(_)) + )); + } + + #[test] + fn map_inbound_status_maps_unavailable_to_inbound_closed() { + let status = tonic::Status::unavailable("inbound queue is closed"); assert!(matches!( - StorageClientError::from(error), + map_inbound_status(&status), StorageClientError::InboundClosed )); } + + #[test] + fn map_inbound_status_maps_invalid_argument_to_invalid_input() { + const MESSAGE: &str = "bad max_items"; + let status = tonic::Status::invalid_argument(MESSAGE); + + match map_inbound_status(&status) { + StorageClientError::InvalidInput(message) => assert_eq!(message, MESSAGE), + error => panic!("unexpected inbound status mapping: {error:?}"), + } + } + + #[test] + fn map_inbound_status_maps_other_codes_to_server() { + let status = tonic::Status::internal("boom"); + + assert!(matches!( + map_inbound_status(&status), + StorageClientError::Server(_) + )); + } } From 00c3196cdd486e9993e44bfd89ea9f936dfaf68f Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Fri, 26 Jun 2026 16:31:38 -0400 Subject: [PATCH 02/21] Implement the missing storage gRPC services --- components/spider-proto-rust/src/lib.rs | 1 + components/spider-proto-rust/src/scheduler.rs | 40 + .../spider-proto-rust/src/unpack/storage.rs | 100 +++ components/spider-scheduler/src/error.rs | 7 - components/spider-storage/src/grpc.rs | 764 ++++++++++++++++-- components/spider-storage/src/state.rs | 2 +- .../spider-storage/src/state/test_utils.rs | 28 +- 7 files changed, 876 insertions(+), 66 deletions(-) create mode 100644 components/spider-proto-rust/src/scheduler.rs diff --git a/components/spider-proto-rust/src/lib.rs b/components/spider-proto-rust/src/lib.rs index c894df6a9..80a7be199 100644 --- a/components/spider-proto-rust/src/lib.rs +++ b/components/spider-proto-rust/src/lib.rs @@ -4,6 +4,7 @@ pub mod error; pub mod id; pub mod io; pub mod job; +pub mod scheduler; pub mod unpack; #[allow(clippy::all, clippy::nursery, clippy::pedantic)] diff --git a/components/spider-proto-rust/src/scheduler.rs b/components/spider-proto-rust/src/scheduler.rs new file mode 100644 index 000000000..e01403e0b --- /dev/null +++ b/components/spider-proto-rust/src/scheduler.rs @@ -0,0 +1,40 @@ +//! Conversions between protobuf scheduler messages and their Spider core representations. + +use spider_core::types::scheduler::RegisteredScheduler; + +use crate::storage; + +impl From for storage::Scheduler { + fn from(scheduler: RegisteredScheduler) -> Self { + Self { + scheduler_id: scheduler.id.get(), + ip_address: scheduler.ip_address.to_string(), + port: u32::from(scheduler.port), + } + } +} + +#[cfg(test)] +mod tests { + use std::net::IpAddr; + + use spider_core::types::{id::SchedulerId, scheduler::RegisteredScheduler}; + + use crate::storage; + + #[test] + fn registered_scheduler_to_protocol_carries_id_ip_and_port() { + const SCHEDULER_ID: SchedulerId = SchedulerId::from(42); + const PORT: u16 = 5678; + + let scheduler = storage::Scheduler::from(RegisteredScheduler { + id: SCHEDULER_ID, + ip_address: IpAddr::V4("127.0.0.1".parse().expect("valid IP")), + port: PORT, + }); + + assert_eq!(scheduler.scheduler_id, SCHEDULER_ID.get()); + assert_eq!(scheduler.ip_address, "127.0.0.1"); + assert_eq!(scheduler.port, u32::from(PORT)); + } +} diff --git a/components/spider-proto-rust/src/unpack/storage.rs b/components/spider-proto-rust/src/unpack/storage.rs index e6967727a..5b0843467 100644 --- a/components/spider-proto-rust/src/unpack/storage.rs +++ b/components/spider-proto-rust/src/unpack/storage.rs @@ -1,5 +1,7 @@ //! [`RequestUnpack`] implementations for `storage.proto` requests. +use std::{net::IpAddr, time::Duration}; + use spider_core::types::id::{ ExecutionManagerId, JobId, @@ -13,11 +15,17 @@ use tonic::Code; use crate::{ storage::{ self, + AddResourceGroupRequest, + ExecutionManagerIdRequest, JobIdRequest, + PollReadyTasksRequest, + RegisterExecutionManagerRequest, RegisterJobRequest, + RegisterSchedulerRequest, RegisterTaskInstanceRequest, ReportTaskFailureRequest, ReportTaskSuccessRequest, + VerifyResourceGroupRequest, }, unpack::{RequestUnpack, UnpackError}, }; @@ -137,6 +145,98 @@ impl RequestUnpack for ReportTaskFailureRequest { } } +/// Unpacks [`AddResourceGroupRequest`] into a tuple containing: +/// +/// * The external resource group ID. +/// * The password. +impl RequestUnpack for AddResourceGroupRequest { + type Unpacked = (String, Vec); + + fn unpack(self) -> Result { + Ok((self.external_resource_group_id, self.password)) + } +} + +/// Unpacks [`VerifyResourceGroupRequest`] into a tuple containing: +/// +/// * The resource group ID. +/// * The password. +impl RequestUnpack for VerifyResourceGroupRequest { + type Unpacked = (ResourceGroupId, Vec); + + fn unpack(self) -> Result { + Ok((ResourceGroupId::from(self.resource_group_id), self.password)) + } +} + +/// Unpacks [`RegisterExecutionManagerRequest`] into the execution manager's IP address. +impl RequestUnpack for RegisterExecutionManagerRequest { + type Unpacked = IpAddr; + + fn unpack(self) -> Result { + self.ip_address + .parse::() + .map_err(|error| invalid_argument(format!("invalid IP address: {error}"))) + } +} + +/// Unpacks [`ExecutionManagerIdRequest`] into an [`ExecutionManagerId`]. +impl RequestUnpack for ExecutionManagerIdRequest { + type Unpacked = ExecutionManagerId; + + fn unpack(self) -> Result { + Ok(ExecutionManagerId::from(self.execution_manager_id)) + } +} + +/// Unpacks [`RegisterSchedulerRequest`] into a tuple containing: +/// +/// * The scheduler IP address. +/// * The scheduler port. +impl RequestUnpack for RegisterSchedulerRequest { + type Unpacked = (IpAddr, u16); + + fn unpack(self) -> Result { + let ip_address = self + .ip_address + .parse::() + .map_err(|error| invalid_argument(format!("invalid IP address: {error}")))?; + let port = u16::try_from(self.port) + .map_err(|_| invalid_argument(format!("port does not fit in `u16`: {}", self.port)))?; + Ok((ip_address, port)) + } +} + +/// Unpacks [`PollReadyTasksRequest`] into a tuple containing: +/// +/// * The maximum number of entries to return. +/// * The maximum duration to block waiting for entries. +impl RequestUnpack for PollReadyTasksRequest { + type Unpacked = (usize, Duration); + + fn unpack(self) -> Result { + let max_items = usize::try_from(self.max_items).map_err(|_| { + invalid_argument(format!( + "max_items does not fit in `usize`: {}", + self.max_items + )) + })?; + Ok((max_items, Duration::from_millis(self.wait_ms))) + } +} + +/// Builds an [`UnpackError`] carrying [`Code::InvalidArgument`] and the given message. +/// +/// # Returns +/// +/// An [`UnpackError`] whose [`Code`] is [`Code::InvalidArgument`] and whose message is `message`. +const fn invalid_argument(message: String) -> UnpackError { + UnpackError { + code: Code::InvalidArgument, + message, + } +} + /// Converts a protobuf [`storage::TaskId`] into a core [`TaskId`]. /// /// # Returns diff --git a/components/spider-scheduler/src/error.rs b/components/spider-scheduler/src/error.rs index df7c52789..af72a8e4c 100644 --- a/components/spider-scheduler/src/error.rs +++ b/components/spider-scheduler/src/error.rs @@ -13,13 +13,6 @@ pub enum StorageClientError { #[error("job not found: {0:?}")] JobNotFound(JobId), - /// The scheduler's storage session is stale. - #[error("stale storage session: {storage_session:?}")] - StaleSession { - /// Storage's current session ID. - storage_session: SessionId, - }, - /// The storage server returned an invalid input error. #[error("invalid storage request: {0}")] InvalidInput(String), diff --git a/components/spider-storage/src/grpc.rs b/components/spider-storage/src/grpc.rs index bbfbba229..ee9ec4a80 100644 --- a/components/spider-storage/src/grpc.rs +++ b/components/spider-storage/src/grpc.rs @@ -1,7 +1,10 @@ //! gRPC service adapters for the storage runtime. use async_trait::async_trait; -use spider_core::types::{id::TaskId, io::SerializedTaskOutputs}; +use spider_core::types::{ + id::{SessionId, TaskId}, + io::SerializedTaskOutputs, +}; use spider_proto_rust::{ storage::{ self, @@ -19,9 +22,9 @@ use tokio_util::sync::CancellationToken; use tonic::{Request, Response, Status}; use crate::{ - cache::error::CacheError, + cache::error::{CacheError, InternalError}, db::{DbError, DbStorage}, - ready_queue::ReadyQueueSender, + ready_queue::{CleanupTaskMarker, CommitTaskMarker, ReadyQueueEntry, ReadyQueueSender}, state::{ServiceState, StorageServerError}, task_instance_pool::TaskInstancePoolConnector, }; @@ -75,7 +78,7 @@ impl< /// The [`Status`] to send to the client: /// /// * `UNAVAILABLE` for a fatal cache-internal error (the service will restart). - /// * `UNAUTHENTICATED` for an unknown or unauthorized resource group. + /// * `UNAUTHENTICATED` for an unknown or unauthorized resource group, or a wrong password. /// * `NOT_FOUND` for a missing job. /// * `FAILED_PRECONDITION` for operations on an invalid job state. /// * `INVALID_ARGUMENT` for a malformed task graph, inputs, or request. @@ -88,14 +91,7 @@ impl< const SERVICE_NAME: &str = "JobOrchestration"; match error { StorageServerError::Cache(CacheError::Internal(e)) => { - tracing::error!( - error = % e, - service = SERVICE_NAME, - tag, - "Internal error in the cache layer. Cancelling service." - ); - self.cancellation_token.cancel(); - Status::unavailable("storage service unavailable") + self.fatal_unavailable_status(SERVICE_NAME, tag, &e) } StorageServerError::Db(db_error) => match &db_error { @@ -157,15 +153,7 @@ impl< Status::invalid_argument(error.to_string()) } - _ => { - tracing::error!( - error = % error, - service = SERVICE_NAME, - tag, - "Unexpected internal error." - ); - Status::internal("internal error") - } + _ => self.unexpected_internal_status(SERVICE_NAME, tag, &error), } } @@ -192,14 +180,7 @@ impl< const SERVICE_NAME: &str = "TaskInstanceManagement"; match error { StorageServerError::Cache(CacheError::Internal(e)) => { - tracing::error!( - error = % e, - service = SERVICE_NAME, - tag, - "Internal error in the cache layer. Cancelling service." - ); - self.cancellation_token.cancel(); - Status::internal("storage service unavailable") + self.fatal_internal_status(SERVICE_NAME, tag, &e) } StorageServerError::StaleSession(storage_session) => { @@ -245,18 +226,251 @@ impl< Status::invalid_argument(error.to_string()) } - _ => { - tracing::error!( + _ => self.unexpected_internal_status(SERVICE_NAME, tag, &error), + } + } + + /// Error handler for inbound queue service errors. + /// + /// This function maps the given [`StorageServerError`] to a [`Status`] that can be sent to the + /// client. The errors are logged for observability. + /// + /// # Returns + /// + /// The [`Status`] to send to the client: + /// + /// * `UNAVAILABLE` when the ready-queue channel is closed (the inbound queue can no longer + /// yield entries). + /// * `INTERNAL` for a fatal cache-internal error (the service will restart) or any other + /// unexpected failure. + pub fn inbound_queue_service_error_handler( + &self, + error: StorageServerError, + tag: &'static str, + ) -> Status { + const SERVICE_NAME: &str = "InboundQueue"; + match error { + StorageServerError::Cache(CacheError::Internal( + InternalError::ReadyQueueChannelClosed, + )) => { + tracing::warn!( + service = SERVICE_NAME, + tag, + "Inbound queue channel is closed." + ); + Status::unavailable("inbound queue is closed") + } + + StorageServerError::Cache(CacheError::Internal(e)) => { + self.fatal_internal_status(SERVICE_NAME, tag, &e) + } + + error => self.unexpected_internal_status(SERVICE_NAME, tag, &error), + } + } + + /// Error handler for resource group management service errors. + /// + /// This function maps the given [`StorageServerError`] to a [`Status`] that can be sent to the + /// client. The errors are logged for observability. + /// + /// # Returns + /// + /// The [`Status`] to send to the client: + /// + /// * `UNAUTHENTICATED` for an unknown resource group or a wrong password. + /// * `ALREADY_EXISTS` for a duplicate external resource group ID. + /// * `UNAVAILABLE` for a fatal cache-internal error (the service will restart). + /// * `INTERNAL` for any other unexpected failure. + pub fn resource_group_management_service_error_handler( + &self, + error: StorageServerError, + tag: &'static str, + ) -> Status { + const SERVICE_NAME: &str = "ResourceGroupManagement"; + match error { + error @ StorageServerError::Db( + DbError::ResourceGroupNotFound(_) | DbError::InvalidPassword(_), + ) => { + tracing::warn!( + error = % error, + service = SERVICE_NAME, + tag, + "Invalid resource group." + ); + Status::unauthenticated("invalid resource group") + } + + error @ StorageServerError::Db(DbError::ResourceGroupAlreadyExists(_)) => { + tracing::warn!( + error = % error, + service = SERVICE_NAME, + tag, + "Resource group already exists." + ); + Status::already_exists(error.to_string()) + } + + StorageServerError::Cache(CacheError::Internal(e)) => { + self.fatal_unavailable_status(SERVICE_NAME, tag, &e) + } + + error => self.unexpected_internal_status(SERVICE_NAME, tag, &error), + } + } + + /// Error handler for execution manager liveness service errors. + /// + /// This function maps the given [`StorageServerError`] to a [`Status`] that can be sent to the + /// client. The errors are logged for observability. + /// + /// # Returns + /// + /// The [`Status`] to send to the client: + /// + /// * `FAILED_PRECONDITION` when the execution manager has already been reaped. + /// * `INVALID_ARGUMENT` for an illegal execution manager ID. + /// * `UNAVAILABLE` for a fatal cache-internal error (the service will restart). + /// * `INTERNAL` for any other unexpected failure. + pub fn execution_manager_liveness_service_error_handler( + &self, + error: StorageServerError, + tag: &'static str, + ) -> Status { + const SERVICE_NAME: &str = "ExecutionManagerLiveness"; + match error { + error @ StorageServerError::Db(DbError::ExecutionManagerAlreadyDead(_)) => { + tracing::warn!( error = % error, service = SERVICE_NAME, tag, - "Unexpected internal error. Cancelling service to avoid cache corruption." + "Execution manager already marked dead." ); - self.cancellation_token.cancel(); - Status::internal("internal error") + Status::failed_precondition(error.to_string()) } + + error @ StorageServerError::Db(DbError::IllegalExecutionManagerId(_)) => { + tracing::warn!( + error = % error, + service = SERVICE_NAME, + tag, + "Illegal execution manager id." + ); + Status::invalid_argument(error.to_string()) + } + + StorageServerError::Cache(CacheError::Internal(e)) => { + self.fatal_unavailable_status(SERVICE_NAME, tag, &e) + } + + error => self.unexpected_internal_status(SERVICE_NAME, tag, &error), } } + + /// Error handler for scheduler registration service errors. + /// + /// This function maps the given [`StorageServerError`] to a [`Status`] that can be sent to the + /// client. The errors are logged for observability. + /// + /// # Returns + /// + /// The [`Status`] to send to the client: + /// + /// * `UNAVAILABLE` for a fatal cache-internal error (the service will restart). + /// * `INTERNAL` for any other failure; scheduler registration currently has no caller-visible + /// error classification beyond a generic server error. + #[must_use] + pub fn scheduler_registration_service_error_handler( + &self, + error: StorageServerError, + tag: &'static str, + ) -> Status { + const SERVICE_NAME: &str = "SchedulerRegistration"; + match error { + StorageServerError::Cache(CacheError::Internal(e)) => { + self.fatal_unavailable_status(SERVICE_NAME, tag, &e) + } + + error => self.unexpected_internal_status(SERVICE_NAME, tag, &error), + } + } + + /// Logs a fatal cache-internal error, cancels the service, and returns an `UNAVAILABLE` status. + /// + /// Shared by the service error handlers whose `Cache(CacheError::Internal)` arm must read as a + /// transient "service restarting" signal (job orchestration, resource-group management, + /// liveness, and scheduler registration). The error is unrecoverable, so the whole storage + /// service is cancelled to avoid cache corruption. + /// + /// # Returns + /// + /// `Status::unavailable("storage service unavailable")`. + fn fatal_unavailable_status( + &self, + service_name: &'static str, + tag: &'static str, + error: &InternalError, + ) -> Status { + tracing::error!( + error = % error, + service = service_name, + tag, + "Internal error in the cache layer. Cancelling service." + ); + self.cancellation_token.cancel(); + Status::unavailable("storage service unavailable") + } + + /// Logs a fatal cache-internal error, cancels the service, and returns an `INTERNAL` status. + /// + /// Shared by the service error handlers whose `UNAVAILABLE` code is already reserved for a + /// caller-facing semantic (task-instance `StaleSession`, inbound `InboundClosed`), so the fatal + /// cache-internal arm falls back to `INTERNAL`. The error is unrecoverable, so the whole + /// storage service is cancelled to avoid cache corruption. + /// + /// # Returns + /// + /// `Status::internal("storage service unavailable")`. + fn fatal_internal_status( + &self, + service_name: &'static str, + tag: &'static str, + error: &InternalError, + ) -> Status { + tracing::error!( + error = % error, + service = service_name, + tag, + "Internal error in the cache layer. Cancelling service." + ); + self.cancellation_token.cancel(); + Status::internal("storage service unavailable") + } + + /// Logs an unexpected error, cancels the service, and returns an `INTERNAL` status. + /// + /// Shared by every service error handler for the catch-all fallback arm. An unmapped error is + /// treated as unrecoverable, so the whole storage service is cancelled to avoid cache + /// corruption. + /// + /// # Returns + /// + /// `Status::internal("internal error")`. + fn unexpected_internal_status( + &self, + service_name: &'static str, + tag: &'static str, + error: &StorageServerError, + ) -> Status { + tracing::error!( + error = % error, + service = service_name, + tag, + "Unexpected internal error. Cancelling service to avoid cache corruption." + ); + self.cancellation_token.cancel(); + Status::internal("internal error") + } } /// Implementation of [`JobOrchestrationService`]. @@ -493,23 +707,62 @@ impl< { async fn poll_ready_tasks( &self, - _request: Request, + request: Request, ) -> Result, Status> { - todo!("Not implemented") + let (max_items, wait) = request.into_inner().unpack()?; + tracing::debug!(max_items, ?wait, "Poll ready tasks request received."); + let entries = self + .inner + .poll_ready_tasks(max_items, wait) + .await + .map_err(|error| self.inbound_queue_service_error_handler(error, "poll_ready_tasks"))?; + Ok(Response::new(storage::PollReadyTasksResponse { + tasks: Some(ready_tasks(self.inner.session_id(), entries)), + })) } async fn poll_ready_commit_tasks( &self, - _request: Request, + request: Request, ) -> Result, Status> { - todo!("Not implemented") + let (max_items, wait) = request.into_inner().unpack()?; + tracing::debug!( + max_items, + ?wait, + "Poll ready commit tasks request received." + ); + let entries = self + .inner + .poll_commit_ready_tasks(max_items, wait) + .await + .map_err(|error| { + self.inbound_queue_service_error_handler(error, "poll_ready_commit_tasks") + })?; + Ok(Response::new(storage::PollReadyTasksResponse { + tasks: Some(ready_tasks(self.inner.session_id(), entries)), + })) } async fn poll_ready_cleanup_tasks( &self, - _request: Request, + request: Request, ) -> Result, Status> { - todo!("Not implemented") + let (max_items, wait) = request.into_inner().unpack()?; + tracing::debug!( + max_items, + ?wait, + "Poll ready cleanup tasks request received." + ); + let entries = self + .inner + .poll_cleanup_ready_tasks(max_items, wait) + .await + .map_err(|error| { + self.inbound_queue_service_error_handler(error, "poll_ready_cleanup_tasks") + })?; + Ok(Response::new(storage::PollReadyTasksResponse { + tasks: Some(ready_tasks(self.inner.session_id(), entries)), + })) } } @@ -523,16 +776,38 @@ impl< { async fn add_resource_group( &self, - _request: Request, + request: Request, ) -> Result, Status> { - todo!("Not implemented") + let (external_id, password) = request.into_inner().unpack()?; + tracing::info!(external_id = % external_id, "Add resource group request received."); + let rg_id = self + .inner + .add_resource_group(external_id, password) + .await + .map_err(|error| { + self.resource_group_management_service_error_handler(error, "add_resource_group") + })?; + Ok(Response::new(storage::ResourceGroupIdResponse { + resource_group_id: rg_id.get(), + })) } async fn verify_resource_group( &self, - _request: Request, + request: Request, ) -> Result, Status> { - todo!("Not implemented") + let (rg_id, password) = request.into_inner().unpack()?; + tracing::info!( + rg_id = rg_id.get(), + "Verify resource group request received." + ); + self.inner + .verify_resource_group(rg_id, &password) + .await + .map_err(|error| { + self.resource_group_management_service_error_handler(error, "verify_resource_group") + })?; + Ok(Response::new(storage::ResourceGroupOperationResponse {})) } } @@ -546,16 +821,51 @@ impl< { async fn register_execution_manager( &self, - _request: Request, + request: Request, ) -> Result, Status> { - todo!("Not implemented") + let ip_address = request.into_inner().unpack()?; + tracing::info!(% ip_address, "Execution manager registration request received."); + let em_id = self + .inner + .register_execution_manager(ip_address) + .await + .map_err(|error| { + self.execution_manager_liveness_service_error_handler( + error, + "register_execution_manager", + ) + })?; + Ok(Response::new(storage::RegisterExecutionManagerResponse { + registration: Some(storage::ExecutionManagerRegistration { + execution_manager_id: em_id.get(), + session_id: self.inner.session_id(), + }), + })) } async fn update_execution_manager_heartbeat( &self, - _request: Request, + request: Request, ) -> Result, Status> { - todo!("Not implemented") + let em_id = request.into_inner().unpack()?; + tracing::info!( + em_id = em_id.get(), + "Execution manager heartbeat request received." + ); + self.inner + .update_execution_manager_heartbeat(em_id) + .await + .map_err(|error| { + self.execution_manager_liveness_service_error_handler( + error, + "update_execution_manager_heartbeat", + ) + })?; + Ok(Response::new( + storage::UpdateExecutionManagerHeartbeatResponse { + session_id: self.inner.session_id(), + }, + )) } } @@ -569,16 +879,41 @@ impl< { async fn register_scheduler( &self, - _request: Request, + request: Request, ) -> Result, Status> { - todo!("Not implemented") + let (ip_address, port) = request.into_inner().unpack()?; + tracing::info!(% ip_address, port, "Scheduler registration request received."); + let scheduler_id = self + .inner + .register_scheduler(ip_address, port) + .await + .map_err(|error| { + self.scheduler_registration_service_error_handler(error, "register_scheduler") + })?; + Ok(Response::new(storage::RegisterSchedulerResponse { + registration: Some(storage::SchedulerRegistration { + scheduler_id: scheduler_id.get(), + session_id: self.inner.session_id(), + }), + })) } async fn get_schedulers( &self, _request: Request, ) -> Result, Status> { - todo!("Not implemented") + tracing::info!("Get schedulers request received."); + let schedulers = self.inner.get_schedulers().await.map_err(|error| { + self.scheduler_registration_service_error_handler(error, "get_schedulers") + })?; + Ok(Response::new(storage::GetSchedulersResponse { + schedulers: Some(storage::SchedulerRegistrations { + schedulers: schedulers + .into_iter() + .map(storage::Scheduler::from) + .collect(), + }), + })) } } @@ -594,10 +929,68 @@ impl< &self, _request: Request, ) -> Result, Status> { - todo!("Not implemented") + let session_id = self.inner.session_id(); + tracing::info!(session_id, "Get session request received."); + Ok(Response::new(storage::GetSessionResponse { session_id })) } } +/// Converts a ready-queue task kind into its protobuf [`storage::TaskId`] form. +trait ToProtoTaskId { + /// # Returns + /// + /// The protobuf task ID for this ready-queue lane marker. + fn to_proto_task_id(self) -> storage::TaskId; +} + +impl ToProtoTaskId for spider_core::task::TaskIndex { + fn to_proto_task_id(self) -> storage::TaskId { + storage::TaskId::from(TaskId::Index(self)) + } +} + +impl ToProtoTaskId for CommitTaskMarker { + fn to_proto_task_id(self) -> storage::TaskId { + storage::TaskId::from(TaskId::Commit) + } +} + +impl ToProtoTaskId for CleanupTaskMarker { + fn to_proto_task_id(self) -> storage::TaskId { + storage::TaskId::from(TaskId::Cleanup) + } +} + +/// Builds a [`storage::ReadyTasks`] message from a batch of ready-queue entries. +/// +/// # Type Parameters +/// +/// * `TaskKindType` - The kind of ready-queue task (`ReadyTask`, `CommitTask`, or `CleanupTask`) +/// carried by each entry; must be convertible to a protobuf task ID. +/// +/// # Returns +/// +/// A [`storage::ReadyTasks`] carrying the storage session and the flattened ready tasks. +fn ready_tasks( + session_id: SessionId, + entries: Vec>, +) -> storage::ReadyTasks { + let tasks = entries + .into_iter() + .map(|entry| { + let resource_group_id = entry.resource_group_id.get(); + let job_id = entry.job_id.get(); + let task_id = entry.task_kind.to_proto_task_id(); + storage::ReadyTask { + resource_group_id, + job_id, + task_id: Some(task_id), + } + }) + .collect(); + storage::ReadyTasks { session_id, tasks } +} + /// # Returns /// /// A [`storage::JobStateResponse`] carrying the given job state. @@ -608,3 +1001,270 @@ fn make_job_state_response( state: storage::JobState::from(state).into(), }) } + +#[cfg(test)] +mod tests { + use spider_core::types::id::{ExecutionManagerId, JobId, ResourceGroupId, SessionId}; + use tokio_util::sync::CancellationToken; + use tonic::{Code, Request}; + + use super::*; + use crate::{ + ready_queue::{ReadyQueueConfig, ReadyQueueSenderHandle, create_ready_queue}, + state::{ + JobCache, + JobCacheGcHandle, + test_utils::{MockDbConnector, MockReadyQueueSender, MockTaskInstancePoolConnector}, + }, + }; + + type TestGrpcState = + GrpcServiceState; + + type TestGrpcStateWithReadyQueue = + GrpcServiceState; + + const TEST_SESSION_ID: SessionId = 1; + + /// # Returns + /// + /// A [`TestGrpcState`] backed by a default mock DB connector. + fn create_grpc_service() -> TestGrpcState { + create_grpc_service_with_db(MockDbConnector::default()) + } + + /// # Returns + /// + /// A [`TestGrpcState`] backed by `db` and a [`MockReadyQueueSender`]. + fn create_grpc_service_with_db(db: MockDbConnector) -> TestGrpcState { + let (_sender, receiver) = + create_ready_queue(&ReadyQueueConfig::default()).expect("ready queue creation"); + let service = ServiceState::new( + db, + TEST_SESSION_ID, + JobCache::new(), + MockReadyQueueSender, + receiver, + MockTaskInstancePoolConnector, + JobCacheGcHandle::new(tokio::sync::mpsc::unbounded_channel().0), + ); + GrpcServiceState::new(service, CancellationToken::new()) + } + + /// # Returns + /// + /// A [`TestGrpcStateWithReadyQueue`] wired to a real ready queue, plus the queue's sender + /// handle so tests can enqueue entries. + fn create_grpc_service_with_ready_queue( + db: MockDbConnector, + ) -> (TestGrpcStateWithReadyQueue, ReadyQueueSenderHandle) { + let (sender, receiver) = + create_ready_queue(&ReadyQueueConfig::default()).expect("ready queue creation"); + let service = ServiceState::new( + db, + TEST_SESSION_ID, + JobCache::new(), + sender.clone(), + receiver, + MockTaskInstancePoolConnector, + JobCacheGcHandle::new(tokio::sync::mpsc::unbounded_channel().0), + ); + ( + GrpcServiceState::new(service, CancellationToken::new()), + sender, + ) + } + + #[tokio::test] + async fn get_session_returns_service_session_id() -> anyhow::Result<()> { + let service = create_grpc_service(); + let response = service + .get_session(Request::new(storage::Void {})) + .await? + .into_inner(); + assert_eq!(response.session_id, TEST_SESSION_ID); + Ok(()) + } + + #[tokio::test] + async fn verify_resource_group_rejects_wrong_password_as_unauthenticated() -> anyhow::Result<()> + { + let service = create_grpc_service(); + let password = b"secret".to_vec(); + let rg_id = service + .add_resource_group(Request::new(storage::AddResourceGroupRequest { + external_resource_group_id: "external-rg".to_owned(), + password: password.clone(), + })) + .await? + .into_inner() + .resource_group_id; + + let result = service + .verify_resource_group(Request::new(storage::VerifyResourceGroupRequest { + resource_group_id: rg_id, + password: b"wrong".to_vec(), + })) + .await; + let status = result.expect_err("a wrong password should be rejected"); + assert_eq!(status.code(), Code::Unauthenticated); + Ok(()) + } + + #[test] + fn job_orchestration_maps_unknown_resource_group_to_unauthenticated() { + let service = create_grpc_service(); + let status = service.job_orchestration_service_error_handler( + StorageServerError::Db(DbError::ResourceGroupNotFound(ResourceGroupId::from(7))), + "test", + ); + assert_eq!(status.code(), Code::Unauthenticated); + } + + #[test] + fn job_orchestration_maps_wrong_password_to_unauthenticated() { + let service = create_grpc_service(); + let status = service.job_orchestration_service_error_handler( + StorageServerError::Db(DbError::InvalidPassword(ResourceGroupId::from(7))), + "test", + ); + assert_eq!(status.code(), Code::Unauthenticated); + } + + #[test] + fn job_orchestration_maps_fatal_cache_internal_to_unavailable() { + let service = create_grpc_service(); + let status = service.job_orchestration_service_error_handler( + StorageServerError::Cache(CacheError::Internal(InternalError::TaskNotRunning)), + "test", + ); + assert_eq!(status.code(), Code::Unavailable); + } + + #[test] + fn resource_group_management_maps_unknown_resource_group_to_unauthenticated() { + let service = create_grpc_service(); + let status = service.resource_group_management_service_error_handler( + StorageServerError::Db(DbError::ResourceGroupNotFound(ResourceGroupId::from(7))), + "test", + ); + assert_eq!(status.code(), Code::Unauthenticated); + } + + #[test] + fn resource_group_management_maps_fatal_cache_internal_to_unavailable() { + let service = create_grpc_service(); + let status = service.resource_group_management_service_error_handler( + StorageServerError::Cache(CacheError::Internal(InternalError::TaskNotRunning)), + "test", + ); + assert_eq!(status.code(), Code::Unavailable); + } + + #[tokio::test] + async fn register_execution_manager_returns_id_and_session() -> anyhow::Result<()> { + let service = create_grpc_service(); + let response = service + .register_execution_manager(Request::new(storage::RegisterExecutionManagerRequest { + ip_address: "127.0.0.1".to_owned(), + })) + .await? + .into_inner(); + let registration = response + .registration + .expect("registration should be present"); + assert_eq!(registration.session_id, TEST_SESSION_ID); + assert_ne!(registration.execution_manager_id, 0); + Ok(()) + } + + #[tokio::test] + async fn heartbeat_returns_session_for_registered_em() -> anyhow::Result<()> { + let service = create_grpc_service(); + let em_id = service + .register_execution_manager(Request::new(storage::RegisterExecutionManagerRequest { + ip_address: "127.0.0.1".to_owned(), + })) + .await? + .into_inner() + .registration + .expect("registration should be present") + .execution_manager_id; + + let response = service + .update_execution_manager_heartbeat(Request::new(storage::ExecutionManagerIdRequest { + execution_manager_id: em_id, + })) + .await? + .into_inner(); + assert_eq!(response.session_id, TEST_SESSION_ID); + Ok(()) + } + + #[tokio::test] + async fn heartbeat_rejects_unknown_em() -> anyhow::Result<()> { + let service = create_grpc_service(); + let result = service + .update_execution_manager_heartbeat(Request::new(storage::ExecutionManagerIdRequest { + execution_manager_id: ExecutionManagerId::from(999).get(), + })) + .await; + let status = result.expect_err("an unknown em id should be rejected"); + assert_eq!(status.code(), Code::InvalidArgument); + Ok(()) + } + + #[tokio::test] + async fn register_and_get_schedulers() -> anyhow::Result<()> { + let service = create_grpc_service(); + let register_response = service + .register_scheduler(Request::new(storage::RegisterSchedulerRequest { + ip_address: "127.0.0.1".to_owned(), + port: 5678, + })) + .await? + .into_inner(); + let registration = register_response + .registration + .expect("registration should be present"); + let scheduler_id = registration.scheduler_id; + assert_eq!(registration.session_id, TEST_SESSION_ID); + + let get_response = service + .get_schedulers(Request::new(storage::Void {})) + .await? + .into_inner(); + let schedulers = get_response + .schedulers + .expect("schedulers should be present"); + assert_eq!(schedulers.schedulers.len(), 1); + assert_eq!(schedulers.schedulers[0].scheduler_id, scheduler_id); + Ok(()) + } + + #[tokio::test] + async fn poll_ready_tasks_returns_entries() -> anyhow::Result<()> { + const TASK_INDEX: usize = 3; + let (service, sender) = create_grpc_service_with_ready_queue(MockDbConnector::default()); + let rg_id = ResourceGroupId::from(7); + let job_id = JobId::from(11); + sender + .send_task_ready(rg_id, job_id, vec![TASK_INDEX]) + .await + .expect("send_task_ready should succeed"); + + let response = service + .poll_ready_tasks(Request::new(storage::PollReadyTasksRequest { + max_items: 10, + wait_ms: 100, + })) + .await? + .into_inner(); + let tasks = response.tasks.expect("ready tasks should be present"); + assert_eq!(tasks.session_id, TEST_SESSION_ID); + assert_eq!(tasks.tasks.len(), 1); + assert_eq!(tasks.tasks[0].resource_group_id, rg_id.get()); + assert_eq!(tasks.tasks[0].job_id, job_id.get()); + Ok(()) + } +} diff --git a/components/spider-storage/src/state.rs b/components/spider-storage/src/state.rs index 4b76ec055..3ea42a3b5 100644 --- a/components/spider-storage/src/state.rs +++ b/components/spider-storage/src/state.rs @@ -11,4 +11,4 @@ pub use runtime::{Runtime, create_runtime}; pub use service::ServiceState; #[cfg(test)] -mod test_utils; +pub(crate) mod test_utils; diff --git a/components/spider-storage/src/state/test_utils.rs b/components/spider-storage/src/state/test_utils.rs index 491ec6be0..360c2863e 100644 --- a/components/spider-storage/src/state/test_utils.rs +++ b/components/spider-storage/src/state/test_utils.rs @@ -79,6 +79,8 @@ pub struct MockDbConnector { pub next_resource_group_id: Arc, pub execution_managers: Arc>, pub next_execution_manager_id: Arc, + pub schedulers: Arc>, + pub next_scheduler_id: Arc, pub session_id: SessionId, } @@ -92,6 +94,8 @@ impl Default for MockDbConnector { next_resource_group_id: Arc::new(AtomicUsize::new(1)), execution_managers: Arc::new(DashMap::new()), next_execution_manager_id: Arc::new(AtomicUsize::new(1)), + schedulers: Arc::new(DashMap::new()), + next_scheduler_id: Arc::new(AtomicUsize::new(1)), session_id: 0, } } @@ -257,18 +261,30 @@ impl ExecutionManagerLivenessManagement for MockDbConnector { impl SchedulerRegistrationManagement for MockDbConnector { async fn register_scheduler( &self, - _ip_address: IpAddr, - _port: u16, + ip_address: IpAddr, + port: u16, ) -> Result { - unreachable!("not implemented for mock connector") + // Mirror the production semantics: only one scheduler is registered at a time. + self.schedulers.clear(); + let counter = self.next_scheduler_id.fetch_add(1, Ordering::Relaxed); + let scheduler_id = SchedulerId::from(counter as u64); + self.schedulers.insert( + scheduler_id, + RegisteredScheduler { + id: scheduler_id, + ip_address, + port, + }, + ); + Ok(scheduler_id) } async fn get_schedulers(&self) -> Result, DbError> { - unreachable!("not implemented for mock connector") + Ok(self.schedulers.iter().map(|entry| *entry).collect()) } - async fn is_scheduler_registered(&self, _scheduler_id: SchedulerId) -> Result { - unreachable!("not implemented for mock connector") + async fn is_scheduler_registered(&self, scheduler_id: SchedulerId) -> Result { + Ok(self.schedulers.contains_key(&scheduler_id)) } } From b41f226ff840d55e4a6c1783035dc75091468032 Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Fri, 26 Jun 2026 16:31:38 -0400 Subject: [PATCH 03/21] Refactor ready-task builder to take a task-id closure instead of a trait --- components/spider-storage/src/grpc.rs | 81 +++++-------------- .../spider-storage/src/state/test_utils.rs | 28 ++----- 2 files changed, 26 insertions(+), 83 deletions(-) diff --git a/components/spider-storage/src/grpc.rs b/components/spider-storage/src/grpc.rs index ee9ec4a80..22c61cf90 100644 --- a/components/spider-storage/src/grpc.rs +++ b/components/spider-storage/src/grpc.rs @@ -24,7 +24,7 @@ use tonic::{Request, Response, Status}; use crate::{ cache::error::{CacheError, InternalError}, db::{DbError, DbStorage}, - ready_queue::{CleanupTaskMarker, CommitTaskMarker, ReadyQueueEntry, ReadyQueueSender}, + ready_queue::{ReadyQueueEntry, ReadyQueueSender}, state::{ServiceState, StorageServerError}, task_instance_pool::TaskInstancePoolConnector, }; @@ -717,7 +717,11 @@ impl< .await .map_err(|error| self.inbound_queue_service_error_handler(error, "poll_ready_tasks"))?; Ok(Response::new(storage::PollReadyTasksResponse { - tasks: Some(ready_tasks(self.inner.session_id(), entries)), + tasks: Some(build_ready_tasks( + self.inner.session_id(), + entries, + |task_index| storage::TaskId::from(TaskId::Index(task_index)), + )), })) } @@ -739,7 +743,9 @@ impl< self.inbound_queue_service_error_handler(error, "poll_ready_commit_tasks") })?; Ok(Response::new(storage::PollReadyTasksResponse { - tasks: Some(ready_tasks(self.inner.session_id(), entries)), + tasks: Some(build_ready_tasks(self.inner.session_id(), entries, |_| { + storage::TaskId::from(TaskId::Commit) + })), })) } @@ -761,7 +767,9 @@ impl< self.inbound_queue_service_error_handler(error, "poll_ready_cleanup_tasks") })?; Ok(Response::new(storage::PollReadyTasksResponse { - tasks: Some(ready_tasks(self.inner.session_id(), entries)), + tasks: Some(build_ready_tasks(self.inner.session_id(), entries, |_| { + storage::TaskId::from(TaskId::Cleanup) + })), })) } } @@ -935,52 +943,31 @@ impl< } } -/// Converts a ready-queue task kind into its protobuf [`storage::TaskId`] form. -trait ToProtoTaskId { - /// # Returns - /// - /// The protobuf task ID for this ready-queue lane marker. - fn to_proto_task_id(self) -> storage::TaskId; -} - -impl ToProtoTaskId for spider_core::task::TaskIndex { - fn to_proto_task_id(self) -> storage::TaskId { - storage::TaskId::from(TaskId::Index(self)) - } -} - -impl ToProtoTaskId for CommitTaskMarker { - fn to_proto_task_id(self) -> storage::TaskId { - storage::TaskId::from(TaskId::Commit) - } -} - -impl ToProtoTaskId for CleanupTaskMarker { - fn to_proto_task_id(self) -> storage::TaskId { - storage::TaskId::from(TaskId::Cleanup) - } -} - /// Builds a [`storage::ReadyTasks`] message from a batch of ready-queue entries. /// /// # Type Parameters /// /// * `TaskKindType` - The kind of ready-queue task (`ReadyTask`, `CommitTask`, or `CleanupTask`) -/// carried by each entry; must be convertible to a protobuf task ID. +/// carried by each entry. +/// +/// # Arguments +/// +/// * `to_task_id` - Converts each entry's lane-specific task kind into its protobuf task ID. /// /// # Returns /// /// A [`storage::ReadyTasks`] carrying the storage session and the flattened ready tasks. -fn ready_tasks( +fn build_ready_tasks( session_id: SessionId, entries: Vec>, + to_task_id: impl Fn(TaskKindType) -> storage::TaskId, ) -> storage::ReadyTasks { let tasks = entries .into_iter() .map(|entry| { let resource_group_id = entry.resource_group_id.get(); let job_id = entry.job_id.get(); - let task_id = entry.task_kind.to_proto_task_id(); + let task_id = to_task_id(entry.task_kind); storage::ReadyTask { resource_group_id, job_id, @@ -1214,34 +1201,6 @@ mod tests { Ok(()) } - #[tokio::test] - async fn register_and_get_schedulers() -> anyhow::Result<()> { - let service = create_grpc_service(); - let register_response = service - .register_scheduler(Request::new(storage::RegisterSchedulerRequest { - ip_address: "127.0.0.1".to_owned(), - port: 5678, - })) - .await? - .into_inner(); - let registration = register_response - .registration - .expect("registration should be present"); - let scheduler_id = registration.scheduler_id; - assert_eq!(registration.session_id, TEST_SESSION_ID); - - let get_response = service - .get_schedulers(Request::new(storage::Void {})) - .await? - .into_inner(); - let schedulers = get_response - .schedulers - .expect("schedulers should be present"); - assert_eq!(schedulers.schedulers.len(), 1); - assert_eq!(schedulers.schedulers[0].scheduler_id, scheduler_id); - Ok(()) - } - #[tokio::test] async fn poll_ready_tasks_returns_entries() -> anyhow::Result<()> { const TASK_INDEX: usize = 3; diff --git a/components/spider-storage/src/state/test_utils.rs b/components/spider-storage/src/state/test_utils.rs index 360c2863e..491ec6be0 100644 --- a/components/spider-storage/src/state/test_utils.rs +++ b/components/spider-storage/src/state/test_utils.rs @@ -79,8 +79,6 @@ pub struct MockDbConnector { pub next_resource_group_id: Arc, pub execution_managers: Arc>, pub next_execution_manager_id: Arc, - pub schedulers: Arc>, - pub next_scheduler_id: Arc, pub session_id: SessionId, } @@ -94,8 +92,6 @@ impl Default for MockDbConnector { next_resource_group_id: Arc::new(AtomicUsize::new(1)), execution_managers: Arc::new(DashMap::new()), next_execution_manager_id: Arc::new(AtomicUsize::new(1)), - schedulers: Arc::new(DashMap::new()), - next_scheduler_id: Arc::new(AtomicUsize::new(1)), session_id: 0, } } @@ -261,30 +257,18 @@ impl ExecutionManagerLivenessManagement for MockDbConnector { impl SchedulerRegistrationManagement for MockDbConnector { async fn register_scheduler( &self, - ip_address: IpAddr, - port: u16, + _ip_address: IpAddr, + _port: u16, ) -> Result { - // Mirror the production semantics: only one scheduler is registered at a time. - self.schedulers.clear(); - let counter = self.next_scheduler_id.fetch_add(1, Ordering::Relaxed); - let scheduler_id = SchedulerId::from(counter as u64); - self.schedulers.insert( - scheduler_id, - RegisteredScheduler { - id: scheduler_id, - ip_address, - port, - }, - ); - Ok(scheduler_id) + unreachable!("not implemented for mock connector") } async fn get_schedulers(&self) -> Result, DbError> { - Ok(self.schedulers.iter().map(|entry| *entry).collect()) + unreachable!("not implemented for mock connector") } - async fn is_scheduler_registered(&self, scheduler_id: SchedulerId) -> Result { - Ok(self.schedulers.contains_key(&scheduler_id)) + async fn is_scheduler_registered(&self, _scheduler_id: SchedulerId) -> Result { + unreachable!("not implemented for mock connector") } } From b98b08506e23dd0a7e3ee6d1f262f68248e3193d Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Fri, 26 Jun 2026 17:50:01 -0400 Subject: [PATCH 04/21] Rename map_inbound_status to inbound_status_to_error to match main's status_to_error convention --- .../src/storage_client/grpc.rs | 32 ++++++++++--------- 1 file changed, 17 insertions(+), 15 deletions(-) diff --git a/components/spider-scheduler/src/storage_client/grpc.rs b/components/spider-scheduler/src/storage_client/grpc.rs index 41e96bac1..f75de362c 100644 --- a/components/spider-scheduler/src/storage_client/grpc.rs +++ b/components/spider-scheduler/src/storage_client/grpc.rs @@ -14,6 +14,7 @@ use spider_proto_rust::storage::{ }; use tonic::{ Code, + Status, transport::{Channel, Endpoint}, }; @@ -65,7 +66,7 @@ impl SchedulerStorageClient for GrpcSchedulerStorageClient { .clone() .poll_ready_tasks(request) .await - .map_err(|status| map_inbound_status(&status))? + .map_err(|status| inbound_status_to_error(&status))? .into_inner(); poll_ready_tasks_response_to_result(response) } @@ -81,7 +82,7 @@ impl SchedulerStorageClient for GrpcSchedulerStorageClient { .clone() .poll_ready_commit_tasks(request) .await - .map_err(|status| map_inbound_status(&status))? + .map_err(|status| inbound_status_to_error(&status))? .into_inner(); poll_ready_tasks_response_to_result(response) } @@ -97,7 +98,7 @@ impl SchedulerStorageClient for GrpcSchedulerStorageClient { .clone() .poll_ready_cleanup_tasks(request) .await - .map_err(|status| map_inbound_status(&status))? + .map_err(|status| inbound_status_to_error(&status))? .into_inner(); poll_ready_tasks_response_to_result(response) } @@ -120,15 +121,16 @@ impl SchedulerStorageClient for GrpcSchedulerStorageClient { } } -/// Maps a [`tonic::Status`] returned by an inbound-queue RPC into a [`StorageClientError`]. +/// Maps an inbound-queue gRPC [`Status`] to a [`StorageClientError`]. /// /// # Returns /// -/// * [`StorageClientError::InboundClosed`] when the inbound queue is closed, signalled by -/// `UNAVAILABLE`. -/// * [`StorageClientError::InvalidInput`] for a malformed request, signalled by `INVALID_ARGUMENT`. -/// * [`StorageClientError::Server`] for any other failure. -fn map_inbound_status(status: &tonic::Status) -> StorageClientError { +/// The [`StorageClientError`] for `status`'s code: +/// +/// * [`StorageClientError::InboundClosed`] for `UNAVAILABLE`. +/// * [`StorageClientError::InvalidInput`] for `INVALID_ARGUMENT`. +/// * [`StorageClientError::Server`] for any other code. +fn inbound_status_to_error(status: &Status) -> StorageClientError { match status.code() { Code::Unavailable => StorageClientError::InboundClosed, Code::InvalidArgument => to_invalid_input_error(status.message()), @@ -321,32 +323,32 @@ mod tests { } #[test] - fn map_inbound_status_maps_unavailable_to_inbound_closed() { + fn inbound_status_maps_unavailable_to_inbound_closed() { let status = tonic::Status::unavailable("inbound queue is closed"); assert!(matches!( - map_inbound_status(&status), + inbound_status_to_error(&status), StorageClientError::InboundClosed )); } #[test] - fn map_inbound_status_maps_invalid_argument_to_invalid_input() { + fn inbound_status_maps_invalid_argument_to_invalid_input() { const MESSAGE: &str = "bad max_items"; let status = tonic::Status::invalid_argument(MESSAGE); - match map_inbound_status(&status) { + match inbound_status_to_error(&status) { StorageClientError::InvalidInput(message) => assert_eq!(message, MESSAGE), error => panic!("unexpected inbound status mapping: {error:?}"), } } #[test] - fn map_inbound_status_maps_other_codes_to_server() { + fn inbound_status_maps_other_codes_to_server() { let status = tonic::Status::internal("boom"); assert!(matches!( - map_inbound_status(&status), + inbound_status_to_error(&status), StorageClientError::Server(_) )); } From b241c34f99b9133a9a4c01200843db02363fbbe3 Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Fri, 26 Jun 2026 17:50:01 -0400 Subject: [PATCH 05/21] Rename map_inbound_status and map_liveness_status to match main's status_to_error convention --- .../src/client/grpc/liveness.rs | 32 +++++++++---------- .../src/storage_client/grpc.rs | 32 ++++++++++--------- 2 files changed, 33 insertions(+), 31 deletions(-) diff --git a/components/spider-execution-manager/src/client/grpc/liveness.rs b/components/spider-execution-manager/src/client/grpc/liveness.rs index 9eeecb172..feae5040b 100644 --- a/components/spider-execution-manager/src/client/grpc/liveness.rs +++ b/components/spider-execution-manager/src/client/grpc/liveness.rs @@ -10,6 +10,7 @@ use spider_proto_rust::storage::{ }; use tonic::{ Code, + Status, transport::{Channel, Endpoint}, }; @@ -52,7 +53,7 @@ impl LivenessClient for GrpcLivenessClient { .clone() .register_execution_manager(request) .await - .map_err(|status| map_liveness_status(&status))? + .map_err(|status| status_to_error(&status))? .into_inner(); register_response_to_result(response) @@ -70,24 +71,23 @@ impl LivenessClient for GrpcLivenessClient { .clone() .update_execution_manager_heartbeat(request) .await - .map_err(|status| map_liveness_status(&status))? + .map_err(|status| status_to_error(&status))? .into_inner(); heartbeat_response_to_result(response) } } -/// Maps a [`tonic::Status`] returned by an execution-manager-liveness RPC into a -/// [`LivenessResponseError`]. +/// Maps an execution-manager-liveness gRPC [`Status`] to a [`LivenessResponseError`]. /// /// # Returns /// -/// * [`LivenessResponseError::MarkedDead`] when storage has already reaped the execution manager, -/// signalled by `FAILED_PRECONDITION`. -/// * [`LivenessResponseError::IllegalId`] when storage rejects the execution manager id, signalled -/// by `INVALID_ARGUMENT`. -/// * [`LivenessResponseError::Transport`] for any other failure. -fn map_liveness_status(status: &tonic::Status) -> LivenessResponseError { +/// The [`LivenessResponseError`] for `status`'s code: +/// +/// * [`LivenessResponseError::MarkedDead`] for `FAILED_PRECONDITION`. +/// * [`LivenessResponseError::IllegalId`] for `INVALID_ARGUMENT`. +/// * [`LivenessResponseError::Transport`] for any other code. +fn status_to_error(status: &Status) -> LivenessResponseError { match status.code() { Code::FailedPrecondition => LivenessResponseError::MarkedDead, Code::InvalidArgument => LivenessResponseError::IllegalId(status.message().to_owned()), @@ -226,32 +226,32 @@ mod tests { } #[test] - fn map_liveness_status_maps_failed_precondition_to_marked_dead() { + fn status_maps_failed_precondition_to_marked_dead() { let status = tonic::Status::failed_precondition("already dead"); assert!(matches!( - map_liveness_status(&status), + status_to_error(&status), LivenessResponseError::MarkedDead )); } #[test] - fn map_liveness_status_maps_invalid_argument_to_illegal_id() { + fn status_maps_invalid_argument_to_illegal_id() { const ERROR_MSG: &str = "bad em id"; let status = tonic::Status::invalid_argument(ERROR_MSG); - match map_liveness_status(&status) { + match status_to_error(&status) { LivenessResponseError::IllegalId(message) => assert_eq!(message, ERROR_MSG), error => panic!("unexpected liveness status mapping: {error:?}"), } } #[test] - fn map_liveness_status_maps_other_codes_to_transport() { + fn status_maps_other_codes_to_transport() { let status = tonic::Status::internal("boom"); assert!(matches!( - map_liveness_status(&status), + status_to_error(&status), LivenessResponseError::Transport(_) )); } diff --git a/components/spider-scheduler/src/storage_client/grpc.rs b/components/spider-scheduler/src/storage_client/grpc.rs index 41e96bac1..f75de362c 100644 --- a/components/spider-scheduler/src/storage_client/grpc.rs +++ b/components/spider-scheduler/src/storage_client/grpc.rs @@ -14,6 +14,7 @@ use spider_proto_rust::storage::{ }; use tonic::{ Code, + Status, transport::{Channel, Endpoint}, }; @@ -65,7 +66,7 @@ impl SchedulerStorageClient for GrpcSchedulerStorageClient { .clone() .poll_ready_tasks(request) .await - .map_err(|status| map_inbound_status(&status))? + .map_err(|status| inbound_status_to_error(&status))? .into_inner(); poll_ready_tasks_response_to_result(response) } @@ -81,7 +82,7 @@ impl SchedulerStorageClient for GrpcSchedulerStorageClient { .clone() .poll_ready_commit_tasks(request) .await - .map_err(|status| map_inbound_status(&status))? + .map_err(|status| inbound_status_to_error(&status))? .into_inner(); poll_ready_tasks_response_to_result(response) } @@ -97,7 +98,7 @@ impl SchedulerStorageClient for GrpcSchedulerStorageClient { .clone() .poll_ready_cleanup_tasks(request) .await - .map_err(|status| map_inbound_status(&status))? + .map_err(|status| inbound_status_to_error(&status))? .into_inner(); poll_ready_tasks_response_to_result(response) } @@ -120,15 +121,16 @@ impl SchedulerStorageClient for GrpcSchedulerStorageClient { } } -/// Maps a [`tonic::Status`] returned by an inbound-queue RPC into a [`StorageClientError`]. +/// Maps an inbound-queue gRPC [`Status`] to a [`StorageClientError`]. /// /// # Returns /// -/// * [`StorageClientError::InboundClosed`] when the inbound queue is closed, signalled by -/// `UNAVAILABLE`. -/// * [`StorageClientError::InvalidInput`] for a malformed request, signalled by `INVALID_ARGUMENT`. -/// * [`StorageClientError::Server`] for any other failure. -fn map_inbound_status(status: &tonic::Status) -> StorageClientError { +/// The [`StorageClientError`] for `status`'s code: +/// +/// * [`StorageClientError::InboundClosed`] for `UNAVAILABLE`. +/// * [`StorageClientError::InvalidInput`] for `INVALID_ARGUMENT`. +/// * [`StorageClientError::Server`] for any other code. +fn inbound_status_to_error(status: &Status) -> StorageClientError { match status.code() { Code::Unavailable => StorageClientError::InboundClosed, Code::InvalidArgument => to_invalid_input_error(status.message()), @@ -321,32 +323,32 @@ mod tests { } #[test] - fn map_inbound_status_maps_unavailable_to_inbound_closed() { + fn inbound_status_maps_unavailable_to_inbound_closed() { let status = tonic::Status::unavailable("inbound queue is closed"); assert!(matches!( - map_inbound_status(&status), + inbound_status_to_error(&status), StorageClientError::InboundClosed )); } #[test] - fn map_inbound_status_maps_invalid_argument_to_invalid_input() { + fn inbound_status_maps_invalid_argument_to_invalid_input() { const MESSAGE: &str = "bad max_items"; let status = tonic::Status::invalid_argument(MESSAGE); - match map_inbound_status(&status) { + match inbound_status_to_error(&status) { StorageClientError::InvalidInput(message) => assert_eq!(message, MESSAGE), error => panic!("unexpected inbound status mapping: {error:?}"), } } #[test] - fn map_inbound_status_maps_other_codes_to_server() { + fn inbound_status_maps_other_codes_to_server() { let status = tonic::Status::internal("boom"); assert!(matches!( - map_inbound_status(&status), + inbound_status_to_error(&status), StorageClientError::Server(_) )); } From c84d32429b6585e17db3020dc08ddb776a8c57d9 Mon Sep 17 00:00:00 2001 From: LinZhihao-723 Date: Sun, 28 Jun 2026 14:52:36 -0400 Subject: [PATCH 06/21] Update proto. --- .../src/generated/storage.rs | 22 ++++++++----------- components/spider-proto/storage/storage.proto | 10 +++------ components/spider-storage/src/grpc.rs | 10 ++++----- 3 files changed, 17 insertions(+), 25 deletions(-) diff --git a/components/spider-proto-rust/src/generated/storage.rs b/components/spider-proto-rust/src/generated/storage.rs index 2a1b15cf5..c017f19f1 100644 --- a/components/spider-proto-rust/src/generated/storage.rs +++ b/components/spider-proto-rust/src/generated/storage.rs @@ -221,10 +221,6 @@ pub struct GetSessionResponse { #[prost(uint64, tag = "1")] pub session_id: u64, } -#[derive(Clone, Copy, PartialEq, ::prost::Message)] -pub struct TaskInstanceOperationResponse {} -#[derive(Clone, Copy, PartialEq, ::prost::Message)] -pub struct ResourceGroupOperationResponse {} #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] #[repr(i32)] pub enum JobState { @@ -642,7 +638,7 @@ pub mod task_instance_management_service_client { &mut self, request: impl tonic::IntoRequest, ) -> std::result::Result< - tonic::Response, + tonic::Response, tonic::Status, > { self.inner @@ -671,7 +667,7 @@ pub mod task_instance_management_service_client { &mut self, request: impl tonic::IntoRequest, ) -> std::result::Result< - tonic::Response, + tonic::Response, tonic::Status, > { self.inner @@ -1001,7 +997,7 @@ pub mod resource_group_management_service_client { &mut self, request: impl tonic::IntoRequest, ) -> std::result::Result< - tonic::Response, + tonic::Response, tonic::Status, > { self.inner @@ -1938,14 +1934,14 @@ pub mod task_instance_management_service_server { &self, request: tonic::Request, ) -> std::result::Result< - tonic::Response, + tonic::Response, tonic::Status, >; async fn report_task_failure( &self, request: tonic::Request, ) -> std::result::Result< - tonic::Response, + tonic::Response, tonic::Status, >; } @@ -2086,7 +2082,7 @@ pub mod task_instance_management_service_server { T: TaskInstanceManagementService, > tonic::server::UnaryService for ReportTaskSuccessSvc { - type Response = super::TaskInstanceOperationResponse; + type Response = super::super::common::Void; type Future = BoxFuture< tonic::Response, tonic::Status, @@ -2137,7 +2133,7 @@ pub mod task_instance_management_service_server { T: TaskInstanceManagementService, > tonic::server::UnaryService for ReportTaskFailureSvc { - type Response = super::TaskInstanceOperationResponse; + type Response = super::super::common::Void; type Future = BoxFuture< tonic::Response, tonic::Status, @@ -2537,7 +2533,7 @@ pub mod resource_group_management_service_server { &self, request: tonic::Request, ) -> std::result::Result< - tonic::Response, + tonic::Response, tonic::Status, >; } @@ -2678,7 +2674,7 @@ pub mod resource_group_management_service_server { T: ResourceGroupManagementService, > tonic::server::UnaryService for VerifyResourceGroupSvc { - type Response = super::ResourceGroupOperationResponse; + type Response = super::super::common::Void; type Future = BoxFuture< tonic::Response, tonic::Status, diff --git a/components/spider-proto/storage/storage.proto b/components/spider-proto/storage/storage.proto index ace49c437..11caf6e3f 100644 --- a/components/spider-proto/storage/storage.proto +++ b/components/spider-proto/storage/storage.proto @@ -15,8 +15,8 @@ service JobOrchestrationService { service TaskInstanceManagementService { rpc RegisterTaskInstance(RegisterTaskInstanceRequest) returns (RegisterTaskInstanceResponse); - rpc ReportTaskSuccess(ReportTaskSuccessRequest) returns (TaskInstanceOperationResponse); - rpc ReportTaskFailure(ReportTaskFailureRequest) returns (TaskInstanceOperationResponse); + rpc ReportTaskSuccess(ReportTaskSuccessRequest) returns (common.Void); + rpc ReportTaskFailure(ReportTaskFailureRequest) returns (common.Void); } service InboundQueueService { @@ -27,7 +27,7 @@ service InboundQueueService { service ResourceGroupManagementService { rpc AddResourceGroup(AddResourceGroupRequest) returns (ResourceGroupIdResponse); - rpc VerifyResourceGroup(VerifyResourceGroupRequest) returns (ResourceGroupOperationResponse); + rpc VerifyResourceGroup(VerifyResourceGroupRequest) returns (common.Void); } service ExecutionManagerLivenessService { @@ -215,7 +215,3 @@ enum JobState { FAILED = 6; CANCELLED = 7; } - -message TaskInstanceOperationResponse {} - -message ResourceGroupOperationResponse {} diff --git a/components/spider-storage/src/grpc.rs b/components/spider-storage/src/grpc.rs index d1cc012b4..ea0a145ed 100644 --- a/components/spider-storage/src/grpc.rs +++ b/components/spider-storage/src/grpc.rs @@ -419,7 +419,7 @@ impl< async fn report_task_success( &self, request: Request, - ) -> Result, Status> { + ) -> Result, Status> { let (session_id, job_id, task_id, task_instance_id, serialized_outputs) = request.into_inner().unpack()?; tracing::info!( @@ -457,13 +457,13 @@ impl< self.task_instance_management_service_error_handler(error, "report_task_success") })?; - Ok(Response::new(storage::TaskInstanceOperationResponse {})) + Ok(Response::new(common::Void {})) } async fn report_task_failure( &self, request: Request, - ) -> Result, Status> { + ) -> Result, Status> { let (session_id, job_id, task_id, task_instance_id, error_message) = request.into_inner().unpack()?; tracing::info!( @@ -481,7 +481,7 @@ impl< .map_err(|error| { self.task_instance_management_service_error_handler(error, "report_task_failure") })?; - Ok(Response::new(storage::TaskInstanceOperationResponse {})) + Ok(Response::new(common::Void {})) } } @@ -533,7 +533,7 @@ impl< async fn verify_resource_group( &self, _request: Request, - ) -> Result, Status> { + ) -> Result, Status> { todo!("Not implemented") } } From 896eef6c0afd0ed8e1c502b7f711734c365e9130 Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Sun, 28 Jun 2026 15:52:14 -0400 Subject: [PATCH 07/21] Remove session id 0 check --- .../src/client/grpc/liveness.rs | 60 +++++++------------ 1 file changed, 21 insertions(+), 39 deletions(-) diff --git a/components/spider-execution-manager/src/client/grpc/liveness.rs b/components/spider-execution-manager/src/client/grpc/liveness.rs index 72462c229..612d6255b 100644 --- a/components/spider-execution-manager/src/client/grpc/liveness.rs +++ b/components/spider-execution-manager/src/client/grpc/liveness.rs @@ -81,7 +81,7 @@ impl LivenessClient for GrpcLivenessClient { .map_err(|status| status_to_error(&status))? .into_inner(); - heartbeat_response_to_result(response) + Ok(heartbeat_response_to_result(response)) } } @@ -109,38 +109,24 @@ fn status_to_error(status: &Status) -> LivenessResponseError { fn register_response_to_result( response: storage::RegisterExecutionManagerResponse, ) -> Result { - match response.registration { - Some(registration) => { - if registration.session_id == 0 { - return Err(LivenessResponseError::Transport( - "register execution manager response carried a zero session id".to_owned(), - )); - } - Ok(RegistrationResponse { - em_id: ExecutionManagerId::from(registration.execution_manager_id), - session_id: registration.session_id, - }) - } - None => Err(LivenessResponseError::Transport( + let registration = response.registration.ok_or_else(|| { + LivenessResponseError::Transport( "register execution manager response missing registration".to_owned(), - )), - } + ) + })?; + Ok(RegistrationResponse { + em_id: ExecutionManagerId::from(registration.execution_manager_id), + session_id: registration.session_id, + }) } /// # Returns /// -/// [`storage::UpdateExecutionManagerHeartbeatResponse`] converted into -/// [`Result`]. -fn heartbeat_response_to_result( +/// The [`SessionId`] carried by `response`. +const fn heartbeat_response_to_result( response: storage::UpdateExecutionManagerHeartbeatResponse, -) -> Result { - let session_id = response.session_id; - if session_id == 0 { - return Err(LivenessResponseError::Transport( - "update execution manager heartbeat response carried a zero session id".to_owned(), - )); - } - Ok(session_id) +) -> SessionId { + response.session_id } /// Converts a displayable transport-layer error into [`LivenessResponseError::Transport`]. @@ -194,7 +180,7 @@ mod tests { } #[test] - fn register_response_to_result_rejects_zero_session_id() { + fn register_response_to_result_accepts_zero_session_id() { let response = storage::RegisterExecutionManagerResponse { registration: Some(storage::ExecutionManagerRegistration { execution_manager_id: 5, @@ -202,10 +188,10 @@ mod tests { }), }; - assert!(matches!( - register_response_to_result(response), - Err(LivenessResponseError::Transport(_)) - )); + let registration = register_response_to_result(response) + .expect("registration response conversion should succeed"); + + assert_eq!(registration.session_id, 0); } #[test] @@ -216,20 +202,16 @@ mod tests { session_id: SESSION_ID, }; - let session_id = heartbeat_response_to_result(response) - .expect("heartbeat response conversion should succeed"); + let session_id = heartbeat_response_to_result(response); assert_eq!(session_id, SESSION_ID); } #[test] - fn heartbeat_response_to_result_rejects_zero_session_id() { + fn heartbeat_response_to_result_accepts_zero_session_id() { let response = storage::UpdateExecutionManagerHeartbeatResponse { session_id: 0 }; - assert!(matches!( - heartbeat_response_to_result(response), - Err(LivenessResponseError::Transport(_)) - )); + assert_eq!(heartbeat_response_to_result(response), 0); } #[test] From c137a3b999cf5e7e14b636674d474b44ade30bb8 Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Sun, 28 Jun 2026 16:31:51 -0400 Subject: [PATCH 08/21] Map inbound queue closed to internal error --- components/spider-scheduler/src/error.rs | 4 ---- .../src/storage_client/grpc.rs | 20 ++++++++++--------- .../src/storage_client/mod.rs | 18 +++++++++++------ 3 files changed, 23 insertions(+), 19 deletions(-) diff --git a/components/spider-scheduler/src/error.rs b/components/spider-scheduler/src/error.rs index df7c52789..a05a85494 100644 --- a/components/spider-scheduler/src/error.rs +++ b/components/spider-scheduler/src/error.rs @@ -5,10 +5,6 @@ use spider_core::types::id::{JobId, SessionId}; /// Errors returned by [`crate::storage_client::SchedulerStorageClient`] operations. #[derive(Debug, thiserror::Error)] pub enum StorageClientError { - /// The inbound queue is closed and can no longer yield ready entries. - #[error("inbound queue is closed")] - InboundClosed, - /// No job with the requested identifier exists. #[error("job not found: {0:?}")] JobNotFound(JobId), diff --git a/components/spider-scheduler/src/storage_client/grpc.rs b/components/spider-scheduler/src/storage_client/grpc.rs index e3eff0623..6e9e598a4 100644 --- a/components/spider-scheduler/src/storage_client/grpc.rs +++ b/components/spider-scheduler/src/storage_client/grpc.rs @@ -142,12 +142,13 @@ impl SchedulerStorageClient for GrpcSchedulerStorageClient { /// /// The [`StorageClientError`] for `status`'s code: /// -/// * [`StorageClientError::InboundClosed`] for `UNAVAILABLE`. +/// * [`StorageClientError::Transport`] for `UNAVAILABLE` (a network-level failure). /// * [`StorageClientError::InvalidInput`] for `INVALID_ARGUMENT`. -/// * [`StorageClientError::Server`] for any other code. +/// * [`StorageClientError::Server`] for any other code (including `INTERNAL`, which the storage +/// server sends when the inbound queue is closed). fn inbound_status_to_error(status: &Status) -> StorageClientError { match status.code() { - Code::Unavailable => StorageClientError::InboundClosed, + Code::Unavailable => to_transport_error(status.message()), Code::InvalidArgument => to_invalid_input_error(status.message()), _ => StorageClientError::Server(status.message().to_owned()), } @@ -338,13 +339,14 @@ mod tests { } #[test] - fn inbound_status_maps_unavailable_to_inbound_closed() { - let status = tonic::Status::unavailable("inbound queue is closed"); + fn inbound_status_maps_unavailable_to_transport() { + const MESSAGE: &str = "inbound queue is closed"; + let status = tonic::Status::unavailable(MESSAGE); - assert!(matches!( - inbound_status_to_error(&status), - StorageClientError::InboundClosed - )); + match inbound_status_to_error(&status) { + StorageClientError::Transport(message) => assert_eq!(message, MESSAGE), + error => panic!("unexpected inbound status mapping: {error:?}"), + } } #[test] diff --git a/components/spider-scheduler/src/storage_client/mod.rs b/components/spider-scheduler/src/storage_client/mod.rs index 808048e0f..15d1e91b9 100644 --- a/components/spider-scheduler/src/storage_client/mod.rs +++ b/components/spider-scheduler/src/storage_client/mod.rs @@ -39,8 +39,10 @@ pub trait SchedulerStorageClient: Send + Sync + Clone { /// /// Returns an error if: /// - /// * [`StorageClientError::InboundClosed`] if the regular-task lane is closed and can no longer - /// yield entries. + /// * [`StorageClientError::Server`] if the regular-task lane is closed and can no longer yield + /// entries, or the storage server returns another error. + /// * [`StorageClientError::Transport`] if the storage transport fails or returns malformed + /// data. async fn poll_ready( &self, max_items: usize, @@ -65,8 +67,10 @@ pub trait SchedulerStorageClient: Send + Sync + Clone { /// /// Returns an error if: /// - /// * [`StorageClientError::InboundClosed`] if the commit-task lane is closed and can no longer - /// yield entries. + /// * [`StorageClientError::Server`] if the commit-task lane is closed and can no longer yield + /// entries, or the storage server returns another error. + /// * [`StorageClientError::Transport`] if the storage transport fails or returns malformed + /// data. async fn poll_commit_ready( &self, max_items: usize, @@ -91,8 +95,10 @@ pub trait SchedulerStorageClient: Send + Sync + Clone { /// /// Returns an error if: /// - /// * [`StorageClientError::InboundClosed`] if the cleanup-task lane is closed and can no longer - /// yield entries. + /// * [`StorageClientError::Server`] if the cleanup-task lane is closed and can no longer yield + /// entries, or the storage server returns another error. + /// * [`StorageClientError::Transport`] if the storage transport fails or returns malformed + /// data. async fn poll_cleanup_ready( &self, max_items: usize, From 2aa8d8e6a92adb75f3dfc7cb18318b64cb3bec17 Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Sun, 28 Jun 2026 16:35:06 -0400 Subject: [PATCH 09/21] Map inbound queue closed to internal error on the server --- components/spider-storage/src/grpc.rs | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/components/spider-storage/src/grpc.rs b/components/spider-storage/src/grpc.rs index 738d3c756..42c31993b 100644 --- a/components/spider-storage/src/grpc.rs +++ b/components/spider-storage/src/grpc.rs @@ -241,9 +241,8 @@ impl< /// /// The [`Status`] to send to the client: /// - /// * `UNAVAILABLE` when the ready-queue channel is closed (the inbound queue can no longer - /// yield entries). - /// * `INTERNAL` for a fatal cache-internal error (the service will restart) or any other + /// * `INTERNAL` when the ready-queue channel is closed (the inbound queue can no longer yield + /// entries), for a fatal cache-internal error (the service will restart), or for any other /// unexpected failure. pub fn inbound_queue_service_error_handler( &self, @@ -260,7 +259,7 @@ impl< tag, "Inbound queue channel is closed." ); - Status::unavailable("inbound queue is closed") + Status::internal("inbound queue is closed") } StorageServerError::Cache(CacheError::Internal(e)) => { @@ -405,7 +404,7 @@ impl< /// Shared by every service error handler's `Cache(CacheError::Internal)` arm. A fatal /// cache-internal error is unrecoverable, so the whole storage service is cancelled to avoid /// cache corruption. It is reported as `INTERNAL` rather than `UNAVAILABLE`, which is reserved - /// for transport-level unavailability such as a dropped connection or a closed inbound queue. + /// for transport-level unavailability such as a dropped connection. /// /// # Returns /// From 3ec8e32cb4474eeb7c21972e6fd77a352d8dd9ee Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Sun, 28 Jun 2026 22:41:03 -0400 Subject: [PATCH 10/21] Fix build_ready_tasks docstring to reference actual lane marker types --- components/spider-storage/src/grpc.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/components/spider-storage/src/grpc.rs b/components/spider-storage/src/grpc.rs index 42c31993b..4e9daa742 100644 --- a/components/spider-storage/src/grpc.rs +++ b/components/spider-storage/src/grpc.rs @@ -925,8 +925,10 @@ impl< /// /// # Type Parameters /// -/// * `TaskKindType` - The kind of ready-queue task (`ReadyTask`, `CommitTask`, or `CleanupTask`) -/// carried by each entry. +/// * `TaskKindType` - The kind of ready-queue task carried by each entry +/// ([`spider_core::task::TaskIndex`] for the regular lane, +/// [`crate::ready_queue::CommitTaskMarker`] for the commit lane, or +/// [`crate::ready_queue::CleanupTaskMarker`] for the cleanup lane). /// /// # Arguments /// From 8a379bbc264782a8298f6dd0190c617497edab4b Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Fri, 3 Jul 2026 00:57:12 -0400 Subject: [PATCH 11/21] Fix logging --- components/spider-storage/src/grpc.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/components/spider-storage/src/grpc.rs b/components/spider-storage/src/grpc.rs index 4e9daa742..347a29805 100644 --- a/components/spider-storage/src/grpc.rs +++ b/components/spider-storage/src/grpc.rs @@ -688,7 +688,7 @@ impl< request: Request, ) -> Result, Status> { let (max_items, wait) = request.into_inner().unpack()?; - tracing::debug!(max_items, ?wait, "Poll ready tasks request received."); + tracing::info!(max_items, ? wait, "Poll ready tasks request received."); let entries = self .inner .poll_ready_tasks(max_items, wait) @@ -708,9 +708,9 @@ impl< request: Request, ) -> Result, Status> { let (max_items, wait) = request.into_inner().unpack()?; - tracing::debug!( + tracing::info!( max_items, - ?wait, + ? wait, "Poll ready commit tasks request received." ); let entries = self @@ -732,9 +732,9 @@ impl< request: Request, ) -> Result, Status> { let (max_items, wait) = request.into_inner().unpack()?; - tracing::debug!( + tracing::info!( max_items, - ?wait, + ? wait, "Poll ready cleanup tasks request received." ); let entries = self From 74e467fe573200274bc79c10f356944d813ef1f9 Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Fri, 3 Jul 2026 01:00:26 -0400 Subject: [PATCH 12/21] Fix docstring --- components/spider-storage/src/grpc.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/components/spider-storage/src/grpc.rs b/components/spider-storage/src/grpc.rs index 347a29805..cce9685b7 100644 --- a/components/spider-storage/src/grpc.rs +++ b/components/spider-storage/src/grpc.rs @@ -408,7 +408,7 @@ impl< /// /// # Returns /// - /// `Status::internal("storage service unavailable")`. + /// An `INTERNAL` [`Status`] carrying the message `"storage service unavailable"`. fn fatal_internal_status( &self, service_name: &'static str, @@ -433,7 +433,7 @@ impl< /// /// # Returns /// - /// `Status::internal("internal error")`. + /// An `INTERNAL` [`Status`] carrying the message `"internal error"`. fn unexpected_internal_status( &self, service_name: &'static str, From 7a6f42d3433f3c558c26e0e0c7b751b2ea1d7b74 Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Fri, 3 Jul 2026 01:01:12 -0400 Subject: [PATCH 13/21] Remove test --- components/spider-storage/src/grpc.rs | 239 -------------------------- 1 file changed, 239 deletions(-) diff --git a/components/spider-storage/src/grpc.rs b/components/spider-storage/src/grpc.rs index cce9685b7..94931cdbd 100644 --- a/components/spider-storage/src/grpc.rs +++ b/components/spider-storage/src/grpc.rs @@ -968,242 +968,3 @@ fn make_job_state_response( state: storage::JobState::from(state).into(), }) } - -#[cfg(test)] -mod tests { - use spider_core::types::id::{ExecutionManagerId, JobId, ResourceGroupId, SessionId}; - use tokio_util::sync::CancellationToken; - use tonic::{Code, Request}; - - use super::*; - use crate::{ - ready_queue::{ReadyQueueConfig, ReadyQueueSenderHandle, create_ready_queue}, - state::{ - JobCache, - JobCacheGcHandle, - test_utils::{MockDbConnector, MockReadyQueueSender, MockTaskInstancePoolConnector}, - }, - }; - - type TestGrpcState = - GrpcServiceState; - - type TestGrpcStateWithReadyQueue = - GrpcServiceState; - - const TEST_SESSION_ID: SessionId = 1; - - /// # Returns - /// - /// A [`TestGrpcState`] backed by a default mock DB connector. - fn create_grpc_service() -> TestGrpcState { - create_grpc_service_with_db(MockDbConnector::default()) - } - - /// # Returns - /// - /// A [`TestGrpcState`] backed by `db` and a [`MockReadyQueueSender`]. - fn create_grpc_service_with_db(db: MockDbConnector) -> TestGrpcState { - let (_sender, receiver) = - create_ready_queue(&ReadyQueueConfig::default()).expect("ready queue creation"); - let service = ServiceState::new( - db, - TEST_SESSION_ID, - JobCache::new(), - MockReadyQueueSender, - receiver, - MockTaskInstancePoolConnector, - JobCacheGcHandle::new(tokio::sync::mpsc::unbounded_channel().0), - ); - GrpcServiceState::new(service, CancellationToken::new()) - } - - /// # Returns - /// - /// A [`TestGrpcStateWithReadyQueue`] wired to a real ready queue, plus the queue's sender - /// handle so tests can enqueue entries. - fn create_grpc_service_with_ready_queue( - db: MockDbConnector, - ) -> (TestGrpcStateWithReadyQueue, ReadyQueueSenderHandle) { - let (sender, receiver) = - create_ready_queue(&ReadyQueueConfig::default()).expect("ready queue creation"); - let service = ServiceState::new( - db, - TEST_SESSION_ID, - JobCache::new(), - sender.clone(), - receiver, - MockTaskInstancePoolConnector, - JobCacheGcHandle::new(tokio::sync::mpsc::unbounded_channel().0), - ); - ( - GrpcServiceState::new(service, CancellationToken::new()), - sender, - ) - } - - #[tokio::test] - async fn get_session_returns_service_session_id() -> anyhow::Result<()> { - let service = create_grpc_service(); - let response = service - .get_session(Request::new(common::Void {})) - .await? - .into_inner(); - assert_eq!(response.session_id, TEST_SESSION_ID); - Ok(()) - } - - #[tokio::test] - async fn verify_resource_group_rejects_wrong_password_as_unauthenticated() -> anyhow::Result<()> - { - let service = create_grpc_service(); - let password = b"secret".to_vec(); - let rg_id = service - .add_resource_group(Request::new(storage::AddResourceGroupRequest { - external_resource_group_id: "external-rg".to_owned(), - password: password.clone(), - })) - .await? - .into_inner() - .resource_group_id; - - let result = service - .verify_resource_group(Request::new(storage::VerifyResourceGroupRequest { - resource_group_id: rg_id, - password: b"wrong".to_vec(), - })) - .await; - let status = result.expect_err("a wrong password should be rejected"); - assert_eq!(status.code(), Code::Unauthenticated); - Ok(()) - } - - #[test] - fn job_orchestration_maps_unknown_resource_group_to_unauthenticated() { - let service = create_grpc_service(); - let status = service.job_orchestration_service_error_handler( - StorageServerError::Db(DbError::ResourceGroupNotFound(ResourceGroupId::from(7))), - "test", - ); - assert_eq!(status.code(), Code::Unauthenticated); - } - - #[test] - fn job_orchestration_maps_wrong_password_to_unauthenticated() { - let service = create_grpc_service(); - let status = service.job_orchestration_service_error_handler( - StorageServerError::Db(DbError::InvalidPassword(ResourceGroupId::from(7))), - "test", - ); - assert_eq!(status.code(), Code::Unauthenticated); - } - - #[test] - fn job_orchestration_maps_fatal_cache_internal_to_internal() { - let service = create_grpc_service(); - let status = service.job_orchestration_service_error_handler( - StorageServerError::Cache(CacheError::Internal(InternalError::TaskNotRunning)), - "test", - ); - assert_eq!(status.code(), Code::Internal); - } - - #[test] - fn resource_group_management_maps_unknown_resource_group_to_unauthenticated() { - let service = create_grpc_service(); - let status = service.resource_group_management_service_error_handler( - StorageServerError::Db(DbError::ResourceGroupNotFound(ResourceGroupId::from(7))), - "test", - ); - assert_eq!(status.code(), Code::Unauthenticated); - } - - #[test] - fn resource_group_management_maps_fatal_cache_internal_to_internal() { - let service = create_grpc_service(); - let status = service.resource_group_management_service_error_handler( - StorageServerError::Cache(CacheError::Internal(InternalError::TaskNotRunning)), - "test", - ); - assert_eq!(status.code(), Code::Internal); - } - - #[tokio::test] - async fn register_execution_manager_returns_id_and_session() -> anyhow::Result<()> { - let service = create_grpc_service(); - let response = service - .register_execution_manager(Request::new(storage::RegisterExecutionManagerRequest { - ip_address: "127.0.0.1".to_owned(), - })) - .await? - .into_inner(); - let registration = response - .registration - .expect("registration should be present"); - assert_eq!(registration.session_id, TEST_SESSION_ID); - assert_ne!(registration.execution_manager_id, 0); - Ok(()) - } - - #[tokio::test] - async fn heartbeat_returns_session_for_registered_em() -> anyhow::Result<()> { - let service = create_grpc_service(); - let em_id = service - .register_execution_manager(Request::new(storage::RegisterExecutionManagerRequest { - ip_address: "127.0.0.1".to_owned(), - })) - .await? - .into_inner() - .registration - .expect("registration should be present") - .execution_manager_id; - - let response = service - .update_execution_manager_heartbeat(Request::new(storage::ExecutionManagerIdRequest { - execution_manager_id: em_id, - })) - .await? - .into_inner(); - assert_eq!(response.session_id, TEST_SESSION_ID); - Ok(()) - } - - #[tokio::test] - async fn heartbeat_rejects_unknown_em() -> anyhow::Result<()> { - let service = create_grpc_service(); - let result = service - .update_execution_manager_heartbeat(Request::new(storage::ExecutionManagerIdRequest { - execution_manager_id: ExecutionManagerId::from(999).get(), - })) - .await; - let status = result.expect_err("an unknown em id should be rejected"); - assert_eq!(status.code(), Code::InvalidArgument); - Ok(()) - } - - #[tokio::test] - async fn poll_ready_tasks_returns_entries() -> anyhow::Result<()> { - const TASK_INDEX: usize = 3; - let (service, sender) = create_grpc_service_with_ready_queue(MockDbConnector::default()); - let rg_id = ResourceGroupId::from(7); - let job_id = JobId::from(11); - sender - .send_task_ready(rg_id, job_id, vec![TASK_INDEX]) - .await - .expect("send_task_ready should succeed"); - - let response = service - .poll_ready_tasks(Request::new(storage::PollReadyTasksRequest { - max_items: 10, - wait_ms: 100, - })) - .await? - .into_inner(); - let tasks = response.tasks.expect("ready tasks should be present"); - assert_eq!(tasks.session_id, TEST_SESSION_ID); - assert_eq!(tasks.tasks.len(), 1); - assert_eq!(tasks.tasks[0].resource_group_id, rg_id.get()); - assert_eq!(tasks.tasks[0].job_id, job_id.get()); - Ok(()) - } -} From 4c7f9e575fbbae9c6eb027f666eef5c1bfa0eaec Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Fri, 3 Jul 2026 11:53:35 -0400 Subject: [PATCH 14/21] Add resend ready task proto --- .../src/generated/storage.rs | 82 +++++++++++++++++++ components/spider-proto/storage/storage.proto | 1 + 2 files changed, 83 insertions(+) diff --git a/components/spider-proto-rust/src/generated/storage.rs b/components/spider-proto-rust/src/generated/storage.rs index b4d4fcc5e..1c38d945a 100644 --- a/components/spider-proto-rust/src/generated/storage.rs +++ b/components/spider-proto-rust/src/generated/storage.rs @@ -1635,6 +1635,32 @@ pub mod inbound_queue_service_client { ); self.inner.unary(req, path, codec).await } + pub async fn resend_ready_tasks( + &mut self, + request: impl tonic::IntoRequest, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + > { + self.inner + .ready() + .await + .map_err(|e| { + tonic::Status::unknown( + format!("Service was not ready: {}", e.into()), + ) + })?; + let codec = tonic_prost::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/storage.InboundQueueService/ResendReadyTasks", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert( + GrpcMethod::new("storage.InboundQueueService", "ResendReadyTasks"), + ); + self.inner.unary(req, path, codec).await + } } } /// Generated server implementations. @@ -1671,6 +1697,13 @@ pub mod inbound_queue_service_server { tonic::Response, tonic::Status, >; + async fn resend_ready_tasks( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; } #[derive(Debug)] pub struct InboundQueueServiceServer { @@ -1895,6 +1928,55 @@ pub mod inbound_queue_service_server { }; Box::pin(fut) } + "/storage.InboundQueueService/ResendReadyTasks" => { + #[allow(non_camel_case_types)] + struct ResendReadyTasksSvc(pub Arc); + impl< + T: InboundQueueService, + > tonic::server::UnaryService + for ResendReadyTasksSvc { + type Response = super::super::common::Void; + type Future = BoxFuture< + tonic::Response, + tonic::Status, + >; + fn call( + &mut self, + request: tonic::Request, + ) -> Self::Future { + let inner = Arc::clone(&self.0); + let fut = async move { + ::resend_ready_tasks( + &inner, + request, + ) + .await + }; + Box::pin(fut) + } + } + let accept_compression_encodings = self.accept_compression_encodings; + let send_compression_encodings = self.send_compression_encodings; + let max_decoding_message_size = self.max_decoding_message_size; + let max_encoding_message_size = self.max_encoding_message_size; + let inner = self.inner.clone(); + let fut = async move { + let method = ResendReadyTasksSvc(inner); + let codec = tonic_prost::ProstCodec::default(); + let mut grpc = tonic::server::Grpc::new(codec) + .apply_compression_config( + accept_compression_encodings, + send_compression_encodings, + ) + .apply_max_message_size_config( + max_decoding_message_size, + max_encoding_message_size, + ); + let res = grpc.unary(method, req).await; + Ok(res) + }; + Box::pin(fut) + } _ => { Box::pin(async move { let mut response = http::Response::new( diff --git a/components/spider-proto/storage/storage.proto b/components/spider-proto/storage/storage.proto index 11caf6e3f..1913b1237 100644 --- a/components/spider-proto/storage/storage.proto +++ b/components/spider-proto/storage/storage.proto @@ -23,6 +23,7 @@ service InboundQueueService { rpc PollReadyTasks(PollReadyTasksRequest) returns (PollReadyTasksResponse); rpc PollReadyCommitTasks(PollReadyTasksRequest) returns (PollReadyTasksResponse); rpc PollReadyCleanupTasks(PollReadyTasksRequest) returns (PollReadyTasksResponse); + rpc ResendReadyTasks(common.Void) returns (common.Void); } service ResourceGroupManagementService { From 3d7b80cce7a1be8a3450361fd5b1b717a88d34f7 Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Fri, 3 Jul 2026 12:00:22 -0400 Subject: [PATCH 15/21] Add resend ready task impl in storage and use in scheduler --- .../src/core_impl/round_robin/tests.rs | 4 ++++ components/spider-scheduler/src/runtime.rs | 4 ++++ .../src/storage_client/grpc.rs | 22 ++++++++++++++----- .../src/storage_client/mod.rs | 11 ++++++++++ components/spider-storage/src/grpc.rs | 11 ++++++++++ 5 files changed, 47 insertions(+), 5 deletions(-) diff --git a/components/spider-scheduler/src/core_impl/round_robin/tests.rs b/components/spider-scheduler/src/core_impl/round_robin/tests.rs index a98228013..b8a259442 100644 --- a/components/spider-scheduler/src/core_impl/round_robin/tests.rs +++ b/components/spider-scheduler/src/core_impl/round_robin/tests.rs @@ -178,6 +178,10 @@ impl SchedulerStorageClient for MockStorageClient { async fn job_state(&self, _job_id: JobId) -> Result { Ok(JobState::Running) } + + async fn resend_ready_tasks(&self) -> Result<(), StorageClientError> { + Ok(()) + } } /// # Returns diff --git a/components/spider-scheduler/src/runtime.rs b/components/spider-scheduler/src/runtime.rs index d514498a7..4cc677c7a 100644 --- a/components/spider-scheduler/src/runtime.rs +++ b/components/spider-scheduler/src/runtime.rs @@ -235,6 +235,10 @@ mod tests { async fn job_state(&self, _job_id: JobId) -> Result { Ok(JobState::Running) } + + async fn resend_ready_tasks(&self) -> Result<(), StorageClientError> { + Ok(()) + } } /// # Returns diff --git a/components/spider-scheduler/src/storage_client/grpc.rs b/components/spider-scheduler/src/storage_client/grpc.rs index 7c5239698..c9ca8c2a0 100644 --- a/components/spider-scheduler/src/storage_client/grpc.rs +++ b/components/spider-scheduler/src/storage_client/grpc.rs @@ -7,11 +7,14 @@ use spider_core::{ job::JobState, types::id::{JobId, ResourceGroupId, SchedulerId, SessionId, TaskId}, }; -use spider_proto_rust::storage::{ - self, - inbound_queue_service_client::InboundQueueServiceClient, - job_orchestration_service_client::JobOrchestrationServiceClient, - scheduler_registration_service_client::SchedulerRegistrationServiceClient, +use spider_proto_rust::{ + common, + storage::{ + self, + inbound_queue_service_client::InboundQueueServiceClient, + job_orchestration_service_client::JobOrchestrationServiceClient, + scheduler_registration_service_client::SchedulerRegistrationServiceClient, + }, }; use spider_utils::grpc::client::ConnectionPool; use tonic::{ @@ -167,6 +170,15 @@ impl SchedulerStorageClient for GrpcSchedulerStorageClient { .into_inner(); job_state_response_to_result(response) } + + async fn resend_ready_tasks(&self) -> Result<(), StorageClientError> { + self.inbound_queue + .get_client() + .resend_ready_tasks(common::Void {}) + .await + .map_err(|status| inbound_status_to_error(&status))?; + Ok(()) + } } /// Maps an inbound-queue gRPC [`Status`] to a [`StorageClientError`]. diff --git a/components/spider-scheduler/src/storage_client/mod.rs b/components/spider-scheduler/src/storage_client/mod.rs index 3f3f75246..73666fd52 100644 --- a/components/spider-scheduler/src/storage_client/mod.rs +++ b/components/spider-scheduler/src/storage_client/mod.rs @@ -145,4 +145,15 @@ pub trait SchedulerStorageClient: Send + Sync + Clone { /// * [`StorageClientError::Server`] if the storage server returns an error. /// * [`StorageClientError::Transport`] if the storage server returns malformed data. async fn job_state(&self, job_id: JobId) -> Result; + + /// Resends ready tasks for all jobs in the cache to the ready queue. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`StorageClientError::Server`] if the storage server returns an error. + /// * [`StorageClientError::Transport`] if the storage transport fails or returns malformed + /// data. + async fn resend_ready_tasks(&self) -> Result<(), StorageClientError>; } diff --git a/components/spider-storage/src/grpc.rs b/components/spider-storage/src/grpc.rs index 94931cdbd..9fd90e00d 100644 --- a/components/spider-storage/src/grpc.rs +++ b/components/spider-storage/src/grpc.rs @@ -750,6 +750,17 @@ impl< })), })) } + + async fn resend_ready_tasks( + &self, + _request: Request, + ) -> Result, Status> { + tracing::info!("Resend ready tasks request received."); + self.inner.resend_ready_tasks().await.map_err(|error| { + self.inbound_queue_service_error_handler(error, "resend_ready_tasks") + })?; + Ok(Response::new(common::Void {})) + } } #[async_trait] From ef3d891ec740b5d254d119b3155f6db7c9bf4128 Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Fri, 3 Jul 2026 12:11:10 -0400 Subject: [PATCH 16/21] Add call to resend ready task on runtime creation --- components/spider-scheduler/src/runtime.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/components/spider-scheduler/src/runtime.rs b/components/spider-scheduler/src/runtime.rs index 4cc677c7a..aa8c42bb0 100644 --- a/components/spider-scheduler/src/runtime.rs +++ b/components/spider-scheduler/src/runtime.rs @@ -125,6 +125,8 @@ pub async fn create_runtime Date: Fri, 3 Jul 2026 12:16:03 -0400 Subject: [PATCH 17/21] Fix docstring --- components/spider-scheduler/src/storage_client/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/components/spider-scheduler/src/storage_client/mod.rs b/components/spider-scheduler/src/storage_client/mod.rs index 73666fd52..b9a21477c 100644 --- a/components/spider-scheduler/src/storage_client/mod.rs +++ b/components/spider-scheduler/src/storage_client/mod.rs @@ -146,7 +146,7 @@ pub trait SchedulerStorageClient: Send + Sync + Clone { /// * [`StorageClientError::Transport`] if the storage server returns malformed data. async fn job_state(&self, job_id: JobId) -> Result; - /// Resends ready tasks for all jobs in the cache to the ready queue. + /// Ask storage to resends ready tasks for all jobs in the cache to the ready queue. /// /// # Errors /// From ac5e5d6e4bda131b3539bbe45acdbc54a528eded Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Fri, 3 Jul 2026 12:16:22 -0400 Subject: [PATCH 18/21] Fix docstring --- components/spider-scheduler/src/storage_client/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/components/spider-scheduler/src/storage_client/mod.rs b/components/spider-scheduler/src/storage_client/mod.rs index b9a21477c..97f8f3745 100644 --- a/components/spider-scheduler/src/storage_client/mod.rs +++ b/components/spider-scheduler/src/storage_client/mod.rs @@ -146,7 +146,7 @@ pub trait SchedulerStorageClient: Send + Sync + Clone { /// * [`StorageClientError::Transport`] if the storage server returns malformed data. async fn job_state(&self, job_id: JobId) -> Result; - /// Ask storage to resends ready tasks for all jobs in the cache to the ready queue. + /// Asks storage to resends ready tasks for all jobs in the cache to the ready queue. /// /// # Errors /// From f3830ccbfc7e8ffe1f552ef4fdc6c1dbc7a065ed Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Fri, 3 Jul 2026 12:32:36 -0400 Subject: [PATCH 19/21] Fix docstring --- components/spider-scheduler/src/storage_client/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/components/spider-scheduler/src/storage_client/mod.rs b/components/spider-scheduler/src/storage_client/mod.rs index 97f8f3745..c4c66ad23 100644 --- a/components/spider-scheduler/src/storage_client/mod.rs +++ b/components/spider-scheduler/src/storage_client/mod.rs @@ -146,7 +146,7 @@ pub trait SchedulerStorageClient: Send + Sync + Clone { /// * [`StorageClientError::Transport`] if the storage server returns malformed data. async fn job_state(&self, job_id: JobId) -> Result; - /// Asks storage to resends ready tasks for all jobs in the cache to the ready queue. + /// Asks storage to resend ready tasks for all jobs in the cache to the ready queue. /// /// # Errors /// From edcb47f29fd0f6187a2aab52cd8efcd97ad19c74 Mon Sep 17 00:00:00 2001 From: LinZhihao-723 Date: Fri, 3 Jul 2026 23:37:08 -0400 Subject: [PATCH 20/21] Check off . --- components/spider-storage/src/grpc.rs | 143 +++++++++----------------- 1 file changed, 50 insertions(+), 93 deletions(-) diff --git a/components/spider-storage/src/grpc.rs b/components/spider-storage/src/grpc.rs index 94931cdbd..d46d500a0 100644 --- a/components/spider-storage/src/grpc.rs +++ b/components/spider-storage/src/grpc.rs @@ -1,10 +1,7 @@ //! gRPC service adapters for the storage runtime. use async_trait::async_trait; -use spider_core::types::{ - id::{SessionId, TaskId}, - io::SerializedTaskOutputs, -}; +use spider_core::types::{id::TaskId, io::SerializedTaskOutputs}; use spider_proto_rust::{ common, storage::{ @@ -241,33 +238,17 @@ impl< /// /// The [`Status`] to send to the client: /// - /// * `INTERNAL` when the ready-queue channel is closed (the inbound queue can no longer yield - /// entries), for a fatal cache-internal error (the service will restart), or for any other - /// unexpected failure. + /// * `INTERNAL` for any failure happened on the server side. This method should never fail + /// under the service's assumption. + #[must_use] + #[allow(clippy::needless_pass_by_value)] pub fn inbound_queue_service_error_handler( &self, error: StorageServerError, tag: &'static str, ) -> Status { const SERVICE_NAME: &str = "InboundQueue"; - match error { - StorageServerError::Cache(CacheError::Internal( - InternalError::ReadyQueueChannelClosed, - )) => { - tracing::warn!( - service = SERVICE_NAME, - tag, - "Inbound queue channel is closed." - ); - Status::internal("inbound queue is closed") - } - - StorageServerError::Cache(CacheError::Internal(e)) => { - self.fatal_internal_status(SERVICE_NAME, tag, &e) - } - - error => self.unexpected_internal_status(SERVICE_NAME, tag, &error), - } + self.unexpected_internal_status(SERVICE_NAME, tag, &error) } /// Error handler for resource group management service errors. @@ -401,11 +382,6 @@ impl< /// Logs a fatal cache-internal error, cancels the service, and returns an `INTERNAL` status. /// - /// Shared by every service error handler's `Cache(CacheError::Internal)` arm. A fatal - /// cache-internal error is unrecoverable, so the whole storage service is cancelled to avoid - /// cache corruption. It is reported as `INTERNAL` rather than `UNAVAILABLE`, which is reserved - /// for transport-level unavailability such as a dropped connection. - /// /// # Returns /// /// An `INTERNAL` [`Status`] carrying the message `"storage service unavailable"`. @@ -427,10 +403,6 @@ impl< /// Logs an unexpected error, cancels the service, and returns an `INTERNAL` status. /// - /// Shared by every service error handler for the catch-all fallback arm. An unmapped error is - /// treated as unrecoverable, so the whole storage service is cancelled to avoid cache - /// corruption. - /// /// # Returns /// /// An `INTERNAL` [`Status`] carrying the message `"internal error"`. @@ -449,6 +421,42 @@ impl< self.cancellation_token.cancel(); Status::internal("internal error") } + + /// Builds a [`storage::ReadyTasks`] message from a batch of ready-queue entries. + /// + /// # Type Parameters + /// + /// * `TaskKindType` - The kind of ready-queue task carried by each entry: + /// * [`spider_core::task::TaskIndex`] for the regular lane. + /// * [`crate::ready_queue::CommitTaskMarker`] for the commit lane. + /// * [`crate::ready_queue::CleanupTaskMarker`] for the cleanup lane. + /// + /// # Returns + /// + /// A [`storage::ReadyTasks`] carrying the storage session and the flattened ready tasks. + fn build_ready_tasks( + &self, + entries: Vec>, + to_task_id: impl Fn(TaskKindType) -> common::TaskId, + ) -> storage::ReadyTasks { + let tasks = entries + .into_iter() + .map(|entry| { + let resource_group_id = entry.resource_group_id.get(); + let job_id = entry.job_id.get(); + let task_id = to_task_id(entry.task_kind); + storage::ReadyTask { + resource_group_id, + job_id, + task_id: Some(task_id), + } + }) + .collect(); + storage::ReadyTasks { + session_id: self.inner.session_id(), + tasks, + } + } } /// Implementation of [`JobOrchestrationService`]. @@ -688,18 +696,16 @@ impl< request: Request, ) -> Result, Status> { let (max_items, wait) = request.into_inner().unpack()?; - tracing::info!(max_items, ? wait, "Poll ready tasks request received."); + tracing::info!(max_items, "Poll ready tasks request received."); let entries = self .inner .poll_ready_tasks(max_items, wait) .await .map_err(|error| self.inbound_queue_service_error_handler(error, "poll_ready_tasks"))?; Ok(Response::new(storage::PollReadyTasksResponse { - tasks: Some(build_ready_tasks( - self.inner.session_id(), - entries, - |task_index| common::TaskId::from(TaskId::Index(task_index)), - )), + tasks: Some(self.build_ready_tasks(entries, |task_index| { + common::TaskId::from(TaskId::Index(task_index)) + })), })) } @@ -708,11 +714,7 @@ impl< request: Request, ) -> Result, Status> { let (max_items, wait) = request.into_inner().unpack()?; - tracing::info!( - max_items, - ? wait, - "Poll ready commit tasks request received." - ); + tracing::info!(max_items, "Poll ready commit tasks request received."); let entries = self .inner .poll_commit_ready_tasks(max_items, wait) @@ -721,9 +723,7 @@ impl< self.inbound_queue_service_error_handler(error, "poll_ready_commit_tasks") })?; Ok(Response::new(storage::PollReadyTasksResponse { - tasks: Some(build_ready_tasks(self.inner.session_id(), entries, |_| { - common::TaskId::from(TaskId::Commit) - })), + tasks: Some(self.build_ready_tasks(entries, |_| common::TaskId::from(TaskId::Commit))), })) } @@ -732,11 +732,7 @@ impl< request: Request, ) -> Result, Status> { let (max_items, wait) = request.into_inner().unpack()?; - tracing::info!( - max_items, - ? wait, - "Poll ready cleanup tasks request received." - ); + tracing::info!(max_items, "Poll ready cleanup tasks request received."); let entries = self .inner .poll_cleanup_ready_tasks(max_items, wait) @@ -745,9 +741,7 @@ impl< self.inbound_queue_service_error_handler(error, "poll_ready_cleanup_tasks") })?; Ok(Response::new(storage::PollReadyTasksResponse { - tasks: Some(build_ready_tasks(self.inner.session_id(), entries, |_| { - common::TaskId::from(TaskId::Cleanup) - })), + tasks: Some(self.build_ready_tasks(entries, |_| common::TaskId::from(TaskId::Cleanup))), })) } } @@ -921,43 +915,6 @@ impl< } } -/// Builds a [`storage::ReadyTasks`] message from a batch of ready-queue entries. -/// -/// # Type Parameters -/// -/// * `TaskKindType` - The kind of ready-queue task carried by each entry -/// ([`spider_core::task::TaskIndex`] for the regular lane, -/// [`crate::ready_queue::CommitTaskMarker`] for the commit lane, or -/// [`crate::ready_queue::CleanupTaskMarker`] for the cleanup lane). -/// -/// # Arguments -/// -/// * `to_task_id` - Converts each entry's lane-specific task kind into its protobuf task ID. -/// -/// # Returns -/// -/// A [`storage::ReadyTasks`] carrying the storage session and the flattened ready tasks. -fn build_ready_tasks( - session_id: SessionId, - entries: Vec>, - to_task_id: impl Fn(TaskKindType) -> common::TaskId, -) -> storage::ReadyTasks { - let tasks = entries - .into_iter() - .map(|entry| { - let resource_group_id = entry.resource_group_id.get(); - let job_id = entry.job_id.get(); - let task_id = to_task_id(entry.task_kind); - storage::ReadyTask { - resource_group_id, - job_id, - task_id: Some(task_id), - } - }) - .collect(); - storage::ReadyTasks { session_id, tasks } -} - /// # Returns /// /// A [`storage::JobStateResponse`] carrying the given job state. From d943a943b6d0899b5f369e38e1c0948ebb723cb8 Mon Sep 17 00:00:00 2001 From: LinZhihao-723 Date: Sat, 4 Jul 2026 00:13:15 -0400 Subject: [PATCH 21/21] Add a flag to control whether to fail the service in . --- components/spider-storage/src/grpc.rs | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/components/spider-storage/src/grpc.rs b/components/spider-storage/src/grpc.rs index d46d500a0..f40731e72 100644 --- a/components/spider-storage/src/grpc.rs +++ b/components/spider-storage/src/grpc.rs @@ -152,7 +152,7 @@ impl< Status::invalid_argument(error.to_string()) } - _ => self.unexpected_internal_status(SERVICE_NAME, tag, &error), + _ => self.unexpected_internal_status(SERVICE_NAME, tag, &error, false), } } @@ -225,7 +225,7 @@ impl< Status::invalid_argument(error.to_string()) } - _ => self.unexpected_internal_status(SERVICE_NAME, tag, &error), + _ => self.unexpected_internal_status(SERVICE_NAME, tag, &error, true), } } @@ -248,7 +248,7 @@ impl< tag: &'static str, ) -> Status { const SERVICE_NAME: &str = "InboundQueue"; - self.unexpected_internal_status(SERVICE_NAME, tag, &error) + self.unexpected_internal_status(SERVICE_NAME, tag, &error, false) } /// Error handler for resource group management service errors. @@ -298,7 +298,7 @@ impl< self.fatal_internal_status(SERVICE_NAME, tag, &e) } - error => self.unexpected_internal_status(SERVICE_NAME, tag, &error), + error => self.unexpected_internal_status(SERVICE_NAME, tag, &error, false), } } @@ -347,7 +347,7 @@ impl< self.fatal_internal_status(SERVICE_NAME, tag, &e) } - error => self.unexpected_internal_status(SERVICE_NAME, tag, &error), + error => self.unexpected_internal_status(SERVICE_NAME, tag, &error, false), } } @@ -376,7 +376,7 @@ impl< self.fatal_internal_status(SERVICE_NAME, tag, &e) } - error => self.unexpected_internal_status(SERVICE_NAME, tag, &error), + error => self.unexpected_internal_status(SERVICE_NAME, tag, &error, false), } } @@ -403,6 +403,8 @@ impl< /// Logs an unexpected error, cancels the service, and returns an `INTERNAL` status. /// + /// If `is_fatal` flag is raised, the service cancellation token will be fired. + /// /// # Returns /// /// An `INTERNAL` [`Status`] carrying the message `"internal error"`. @@ -411,6 +413,7 @@ impl< service_name: &'static str, tag: &'static str, error: &StorageServerError, + is_fatal: bool, ) -> Status { tracing::error!( error = % error, @@ -418,7 +421,9 @@ impl< tag, "Unexpected internal error. Cancelling service to avoid cache corruption." ); - self.cancellation_token.cancel(); + if is_fatal { + self.cancellation_token.cancel(); + } Status::internal("internal error") }