From b8b43ebdcba6f7464ca58fe2d8560becb86fbc34 Mon Sep 17 00:00:00 2001 From: LinZhihao-723 Date: Tue, 23 Jun 2026 18:35:01 -0400 Subject: [PATCH] Done. --- components/spider-core/src/types/id.rs | 10 + .../src/client/grpc/storage.rs | 137 ++++--------- .../src/client/storage.rs | 4 +- .../spider-execution-manager/src/runtime.rs | 4 +- components/spider-proto-rust/src/error.rs | 8 + .../src/generated/storage.rs | 111 +++------- components/spider-proto-rust/src/io.rs | 82 ++++++++ components/spider-proto-rust/src/lib.rs | 1 + .../spider-proto-rust/src/unpack/storage.rs | 149 +++++++++++++- components/spider-proto/storage/storage.proto | 43 ++-- components/spider-storage/src/cache/error.rs | 11 +- components/spider-storage/src/grpc.rs | 189 +++++++++++++++++- components/spider-storage/src/state/error.rs | 9 +- .../spider-storage/src/state/service.rs | 8 +- .../huntsman/em-runtime/tests/test_runtime.rs | 10 +- 15 files changed, 536 insertions(+), 240 deletions(-) create mode 100644 components/spider-proto-rust/src/io.rs diff --git a/components/spider-core/src/types/id.rs b/components/spider-core/src/types/id.rs index 462ab4580..e500c2ce7 100644 --- a/components/spider-core/src/types/id.rs +++ b/components/spider-core/src/types/id.rs @@ -138,6 +138,16 @@ pub enum TaskId { Cleanup, } +impl Display for TaskId { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Index(index) => write!(formatter, "{index}"), + Self::Commit => write!(formatter, "commit"), + Self::Cleanup => write!(formatter, "cleanup"), + } + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum JobIdMarker {} pub type JobId = Id; diff --git a/components/spider-execution-manager/src/client/grpc/storage.rs b/components/spider-execution-manager/src/client/grpc/storage.rs index 4e0a4ef1e..ee2827aee 100644 --- a/components/spider-execution-manager/src/client/grpc/storage.rs +++ b/components/spider-execution-manager/src/client/grpc/storage.rs @@ -10,12 +10,13 @@ use spider_core::types::{ }; use spider_proto_rust::storage::{ self, - register_task_instance_response, - task_instance_management_error, task_instance_management_service_client::TaskInstanceManagementServiceClient, - task_instance_operation_response, }; -use tonic::transport::{Channel, Endpoint}; +use tonic::{ + Code, + Status, + transport::{Channel, Endpoint}, +}; use crate::client::storage::{StorageClient, StorageResponseError}; @@ -65,22 +66,16 @@ impl StorageClient for GrpcStorageClient { .clone() .register_task_instance(request) .await - .map_err(to_transport_error)? + .map_err(|status| status_to_error(&status))? .into_inner(); - match response.result { - Some(register_task_instance_response::Result::ExecutionContext(bytes)) => { - bincode::deserialize(&bytes).map_err(|error| { - StorageResponseError::Transport(format!( - "failed to decode execution context: {error}" - )) - }) - } - Some(register_task_instance_response::Result::Error(error)) => Err(error.into()), - None => Err(StorageResponseError::Transport( - "register task instance response missing result".to_owned(), - )), - } + let execution_context = response.execution_context.ok_or_else(|| { + StorageResponseError::Transport( + "register task instance response missing execution context".to_owned(), + ) + })?; + ExecutionContext::try_from(execution_context) + .map_err(|error| StorageResponseError::Transport(error.to_string())) } async fn report_task_success( @@ -100,15 +95,12 @@ impl StorageClient for GrpcStorageClient { serialized_outputs: serialized_outputs.unwrap_or_default(), task_instance_id, }; - let response = self - .client + self.client .clone() .report_task_success(request) .await - .map_err(to_transport_error)? - .into_inner(); - - storage_operation_response_to_result(response) + .map_err(|status| status_to_error(&status))?; + Ok(()) } async fn report_task_failure( @@ -128,52 +120,31 @@ impl StorageClient for GrpcStorageClient { error_message, task_instance_id, }; - let response = self - .client + self.client .clone() .report_task_failure(request) .await - .map_err(to_transport_error)? - .into_inner(); - - storage_operation_response_to_result(response) - } -} - -impl From for StorageResponseError { - fn from(error: storage::TaskInstanceManagementError) -> Self { - match task_instance_management_error::ErrCode::try_from(error.err_code) { - Ok(task_instance_management_error::ErrCode::StaleSession) => Self::StaleSession { - storage_session: error.storage_session, - }, - Ok(task_instance_management_error::ErrCode::CacheStale) => { - Self::CacheStale(error.message) - } - Ok( - task_instance_management_error::ErrCode::Server - | task_instance_management_error::ErrCode::Unspecified, - ) => Self::Server(error.message), - Ok(task_instance_management_error::ErrCode::InvalidInput) => { - Self::InvalidInput(error.message) - } - Err(error) => Self::Transport(format!("unknown task instance error kind: {error}")), - } + .map_err(|status| status_to_error(&status))?; + Ok(()) } } +/// Maps a task-instance management gRPC [`Status`] to a [`StorageResponseError`]. +/// /// # Returns /// -/// [`storage::TaskInstanceOperationResponse`] converted into -/// [`Result<(), StorageResponseError>`]. -fn storage_operation_response_to_result( - response: storage::TaskInstanceOperationResponse, -) -> Result<(), StorageResponseError> { - match response.result { - Some(task_instance_operation_response::Result::Ok(_)) => Ok(()), - Some(task_instance_operation_response::Result::Error(error)) => Err(error.into()), - None => Err(StorageResponseError::Transport( - "storage operation response missing `result` message".to_owned(), - )), +/// The [`StorageResponseError`] for `status`'s code: +/// +/// * [`StorageResponseError::StaleSession`] for `UNAVAILABLE`. +/// * [`StorageResponseError::CacheStale`] for `FAILED_PRECONDITION`. +/// * [`StorageResponseError::InvalidInput`] for `INVALID_ARGUMENT`. +/// * [`StorageResponseError::Server`] for any other code. +fn status_to_error(status: &Status) -> StorageResponseError { + match status.code() { + Code::Unavailable => StorageResponseError::StaleSession(status.message().to_owned()), + Code::FailedPrecondition => StorageResponseError::CacheStale(status.message().to_owned()), + Code::InvalidArgument => StorageResponseError::InvalidInput(status.message().to_owned()), + _ => StorageResponseError::Server(status.message().to_owned()), } } @@ -191,44 +162,18 @@ mod tests { use super::*; #[test] - fn storage_error_maps_stale_session() { - let error = storage::TaskInstanceManagementError { - err_code: task_instance_management_error::ErrCode::StaleSession.into(), - message: "stale".to_owned(), - storage_session: 7, - }; - - match StorageResponseError::from(error) { - StorageResponseError::StaleSession { storage_session } => { - assert_eq!(7, storage_session); - } - error => panic!("unexpected storage response error: {error:?}"), - } - } - - #[test] - fn storage_error_maps_unknown_kind_to_transport_error() { - let error = storage::TaskInstanceManagementError { - err_code: 99, - message: "unknown".to_owned(), - storage_session: 0, - }; - - match StorageResponseError::from(error) { - StorageResponseError::Transport(message) => { - assert!(message.contains("unknown task instance error kind")); - } - error => panic!("unexpected storage response error: {error:?}"), + fn status_maps_unavailable_to_stale_session() { + match status_to_error(&Status::unavailable("storage session is 9")) { + StorageResponseError::StaleSession(message) => assert!(message.contains('9')), + error => panic!("unexpected error: {error:?}"), } } #[test] - fn missing_storage_operation_result_is_transport_error() { - match storage_operation_response_to_result(storage::TaskInstanceOperationResponse { - result: None, - }) { - Err(StorageResponseError::Transport(_)) => {} - result => panic!("unexpected storage operation result: {result:?}"), + fn status_maps_invalid_argument_to_invalid_input() { + match status_to_error(&Status::invalid_argument("bad task id")) { + StorageResponseError::InvalidInput(message) => assert!(message.contains("bad task id")), + error => panic!("unexpected error: {error:?}"), } } } diff --git a/components/spider-execution-manager/src/client/storage.rs b/components/spider-execution-manager/src/client/storage.rs index b8e13a4d7..aa900bdc0 100644 --- a/components/spider-execution-manager/src/client/storage.rs +++ b/components/spider-execution-manager/src/client/storage.rs @@ -17,8 +17,8 @@ use spider_core::types::{ #[derive(Debug, thiserror::Error)] pub enum StorageResponseError { /// The `session_id` carried with the request does not match storage's current session. - #[error("stale session (storage now at {storage_session})")] - StaleSession { storage_session: SessionId }, + #[error("stale session: {0}")] + StaleSession(String), /// Storage's job cache rejected the operation as stale (e.g. the task or its job has already /// terminated). diff --git a/components/spider-execution-manager/src/runtime.rs b/components/spider-execution-manager/src/runtime.rs index 91d011395..8b2a9f3fa 100644 --- a/components/spider-execution-manager/src/runtime.rs +++ b/components/spider-execution-manager/src/runtime.rs @@ -398,10 +398,10 @@ impl< Ok(Some(execution_context)) } Err(err) => match &err { - StorageResponseError::StaleSession { storage_session } => { + StorageResponseError::StaleSession(message) => { tracing::warn!( bundle_session = response.session_id, - storage_session = storage_session, + error = % message, job_id = ? response.task_assignment.job_id, task_id = ? response.task_assignment.task_id, "Storage rejected task registration as stale. Dropping the assignment." diff --git a/components/spider-proto-rust/src/error.rs b/components/spider-proto-rust/src/error.rs index f9014e7b7..4c1db18b2 100644 --- a/components/spider-proto-rust/src/error.rs +++ b/components/spider-proto-rust/src/error.rs @@ -14,4 +14,12 @@ pub enum Error { /// A protobuf [`crate::storage::JobState`] was left unspecified. #[error("job state is unspecified")] JobStateUnspecified, + + /// A protobuf [`crate::storage::TdlContext`] was missing. + #[error("TDL context is missing")] + TdlContextMissing, + + /// A protobuf [`crate::storage::TimeoutPolicy`] was missing. + #[error("timeout policy is missing")] + TimeoutPolicyMissing, } diff --git a/components/spider-proto-rust/src/generated/storage.rs b/components/spider-proto-rust/src/generated/storage.rs index 3266f9525..a0b5ec5c9 100644 --- a/components/spider-proto-rust/src/generated/storage.rs +++ b/components/spider-proto-rust/src/generated/storage.rs @@ -89,18 +89,33 @@ pub struct RegisterTaskInstanceRequest { } #[derive(Clone, PartialEq, ::prost::Message)] pub struct RegisterTaskInstanceResponse { - #[prost(oneof = "register_task_instance_response::Result", tags = "1, 2")] - pub result: ::core::option::Option, + #[prost(message, optional, tag = "1")] + pub execution_context: ::core::option::Option, } -/// Nested message and enum types in `RegisterTaskInstanceResponse`. -pub mod register_task_instance_response { - #[derive(Clone, PartialEq, ::prost::Oneof)] - pub enum Result { - #[prost(bytes, tag = "1")] - ExecutionContext(::prost::alloc::vec::Vec), - #[prost(message, tag = "2")] - Error(super::TaskInstanceManagementError), - } +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ExecutionContext { + #[prost(uint64, tag = "1")] + pub task_instance_id: u64, + #[prost(message, optional, tag = "2")] + pub tdl_context: ::core::option::Option, + #[prost(message, optional, tag = "3")] + pub timeout_policy: ::core::option::Option, + #[prost(bytes = "vec", tag = "4")] + pub serialized_inputs: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct TdlContext { + #[prost(string, tag = "1")] + pub package: ::prost::alloc::string::String, + #[prost(string, tag = "2")] + pub task_func: ::prost::alloc::string::String, +} +#[derive(Clone, Copy, PartialEq, ::prost::Message)] +pub struct TimeoutPolicy { + #[prost(uint64, tag = "1")] + pub soft_timeout_ms: u64, + #[prost(uint64, tag = "2")] + pub hard_timeout_ms: u64, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct ReportTaskSuccessRequest { @@ -297,21 +312,8 @@ pub mod task_id { Cleanup(super::Void), } } -#[derive(Clone, PartialEq, ::prost::Message)] -pub struct TaskInstanceOperationResponse { - #[prost(oneof = "task_instance_operation_response::Result", tags = "1, 2")] - pub result: ::core::option::Option, -} -/// Nested message and enum types in `TaskInstanceOperationResponse`. -pub mod task_instance_operation_response { - #[derive(Clone, PartialEq, ::prost::Oneof)] - pub enum Result { - #[prost(message, tag = "1")] - Ok(super::Void), - #[prost(message, tag = "2")] - Error(super::TaskInstanceManagementError), - } -} +#[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")] @@ -330,63 +332,6 @@ pub mod resource_group_operation_response { #[derive(Clone, Copy, PartialEq, ::prost::Message)] pub struct Void {} #[derive(Clone, PartialEq, ::prost::Message)] -pub struct TaskInstanceManagementError { - #[prost(enumeration = "task_instance_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 `TaskInstanceManagementError`. -pub mod task_instance_management_error { - #[derive( - Clone, - Copy, - Debug, - PartialEq, - Eq, - Hash, - PartialOrd, - Ord, - ::prost::Enumeration - )] - #[repr(i32)] - pub enum ErrCode { - Unspecified = 0, - StaleSession = 1, - CacheStale = 2, - Server = 3, - InvalidInput = 4, - } - 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::CacheStale => "CACHE_STALE", - 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), - "CACHE_STALE" => Some(Self::CacheStale), - "SERVER" => Some(Self::Server), - "INVALID_INPUT" => Some(Self::InvalidInput), - _ => None, - } - } - } -} -#[derive(Clone, PartialEq, ::prost::Message)] pub struct InboundQueueResponseError { #[prost(enumeration = "inbound_queue_response_error::ErrCode", tag = "1")] pub err_code: i32, diff --git a/components/spider-proto-rust/src/io.rs b/components/spider-proto-rust/src/io.rs new file mode 100644 index 000000000..0a2a8d154 --- /dev/null +++ b/components/spider-proto-rust/src/io.rs @@ -0,0 +1,82 @@ +//! Conversions between protobuf I/O messages and their Spider core representations. + +use spider_core::{ + task::{TdlContext, TimeoutPolicy}, + types::io::ExecutionContext, +}; + +use crate::{error::Error, storage}; + +impl TryFrom for ExecutionContext { + type Error = Error; + + fn try_from(execution_context: storage::ExecutionContext) -> Result { + let tdl_context = execution_context + .tdl_context + .ok_or(Error::TdlContextMissing)?; + let timeout_policy = execution_context + .timeout_policy + .ok_or(Error::TimeoutPolicyMissing)?; + Ok(Self { + task_instance_id: execution_context.task_instance_id, + tdl_context: TdlContext { + package: tdl_context.package, + task_func: tdl_context.task_func, + }, + timeout_policy: TimeoutPolicy { + soft_timeout_ms: timeout_policy.soft_timeout_ms, + hard_timeout_ms: timeout_policy.hard_timeout_ms, + }, + serialized_inputs: execution_context.serialized_inputs, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn execution_context_converts_from_proto() { + let proto = storage::ExecutionContext { + task_instance_id: 7, + tdl_context: Some(storage::TdlContext { + package: "pkg".to_owned(), + task_func: "func".to_owned(), + }), + timeout_policy: Some(storage::TimeoutPolicy { + soft_timeout_ms: 100, + hard_timeout_ms: 200, + }), + serialized_inputs: vec![1, 2, 3], + }; + + let execution_context = + ExecutionContext::try_from(proto).expect("conversion should succeed"); + + assert_eq!(execution_context.task_instance_id, 7); + assert_eq!(execution_context.tdl_context.package, "pkg"); + assert_eq!(execution_context.tdl_context.task_func, "func"); + assert_eq!(execution_context.timeout_policy.soft_timeout_ms, 100); + assert_eq!(execution_context.timeout_policy.hard_timeout_ms, 200); + assert_eq!(execution_context.serialized_inputs, vec![1, 2, 3]); + } + + #[test] + fn execution_context_rejects_missing_tdl_context() { + let proto = storage::ExecutionContext { + task_instance_id: 7, + tdl_context: None, + timeout_policy: Some(storage::TimeoutPolicy { + soft_timeout_ms: 100, + hard_timeout_ms: 200, + }), + serialized_inputs: Vec::new(), + }; + + assert!(matches!( + ExecutionContext::try_from(proto), + Err(Error::TdlContextMissing) + )); + } +} diff --git a/components/spider-proto-rust/src/lib.rs b/components/spider-proto-rust/src/lib.rs index 2061cd673..c894df6a9 100644 --- a/components/spider-proto-rust/src/lib.rs +++ b/components/spider-proto-rust/src/lib.rs @@ -2,6 +2,7 @@ pub mod error; pub mod id; +pub mod io; pub mod job; pub mod unpack; diff --git a/components/spider-proto-rust/src/unpack/storage.rs b/components/spider-proto-rust/src/unpack/storage.rs index 89671cd05..170f52dc0 100644 --- a/components/spider-proto-rust/src/unpack/storage.rs +++ b/components/spider-proto-rust/src/unpack/storage.rs @@ -1,10 +1,24 @@ //! [`RequestUnpack`] implementations for `storage.proto` requests. -use spider_core::types::id::{JobId, ResourceGroupId}; +use spider_core::types::id::{ + ExecutionManagerId, + JobId, + ResourceGroupId, + SessionId, + TaskId, + TaskInstanceId, +}; use tonic::Code; use crate::{ - storage::{JobIdRequest, RegisterJobRequest}, + storage::{ + self, + JobIdRequest, + RegisterJobRequest, + RegisterTaskInstanceRequest, + ReportTaskFailureRequest, + ReportTaskSuccessRequest, + }, unpack::{RequestUnpack, UnpackError}, }; @@ -17,13 +31,18 @@ impl RequestUnpack for RegisterJobRequest { type Unpacked = (ResourceGroupId, String, Vec); fn unpack(self) -> Result { - let serialized_task_graph = - String::from_utf8(self.serialized_task_graph).map_err(|error| { - tracing::error!(error = % error, "Invalid UTF-8 in serialized task graph."); - UnpackError { - code: Code::InvalidArgument, - message: format!("invalid UTF-8 in serialized task graph: {error}"), - } + let serialized_task_graph = String::from_utf8(self.serialized_task_graph) + .map_err(|error| UnpackError { + code: Code::InvalidArgument, + message: format!("invalid UTF-8 in serialized task graph: {error}"), + }) + .inspect_err(|error| { + tracing::error!( + error = % error.message, + request = "RegisterJob", + resource_group_id = self.resource_group_id, + "Failed to unpack request." + ); })?; Ok(( ResourceGroupId::from(self.resource_group_id), @@ -41,3 +60,115 @@ impl RequestUnpack for JobIdRequest { Ok(JobId::from(self.job_id)) } } + +/// Unpacks [`RegisterTaskInstanceRequest`] into a tuple containing: +/// +/// * The session ID. +/// * The job ID. +/// * The task ID. +/// * The execution manager ID. +impl RequestUnpack for RegisterTaskInstanceRequest { + type Unpacked = (SessionId, JobId, TaskId, ExecutionManagerId); + + fn unpack(self) -> Result { + let task_id = unpack_task_id(self.task_id).inspect_err(|error| { + tracing::error!( + error = % error.message, + request = "RegisterTaskInstance", + em_id = self.execution_manager_id, + "Failed to unpack request." + ); + })?; + Ok(( + self.session_id, + JobId::from(self.job_id), + task_id, + ExecutionManagerId::from(self.execution_manager_id), + )) + } +} + +/// Unpacks [`ReportTaskSuccessRequest`] into a tuple containing: +/// +/// * The session ID. +/// * The job ID. +/// * The task ID. +/// * The task instance ID. +/// * The serialized task outputs. +impl RequestUnpack for ReportTaskSuccessRequest { + type Unpacked = (SessionId, JobId, TaskId, TaskInstanceId, Vec); + + fn unpack(self) -> Result { + let task_id = unpack_task_id(self.task_id).inspect_err(|error| { + tracing::error!( + error = % error.message, + request = "ReportTaskSuccess", + em_id = self.execution_manager_id, + task_instance_id = self.task_instance_id, + "Failed to unpack request." + ); + })?; + + Ok(( + self.session_id, + JobId::from(self.job_id), + task_id, + self.task_instance_id, + self.serialized_outputs, + )) + } +} + +/// Unpacks [`ReportTaskFailureRequest`] into a tuple containing: +/// +/// * The session ID. +/// * The job ID. +/// * The task ID. +/// * The task instance ID. +/// * The error message. +impl RequestUnpack for ReportTaskFailureRequest { + type Unpacked = (SessionId, JobId, TaskId, TaskInstanceId, String); + + fn unpack(self) -> Result { + let task_id = unpack_task_id(self.task_id).inspect_err(|error| { + tracing::error!( + error = % error.message, + request = "ReportTaskFailure", + em_id = self.execution_manager_id, + task_instance_id = self.task_instance_id, + "Failed to unpack request." + ); + })?; + + Ok(( + self.session_id, + JobId::from(self.job_id), + task_id, + self.task_instance_id, + self.error_message, + )) + } +} + +/// Converts a protobuf [`storage::TaskId`] into a core [`TaskId`]. +/// +/// # Returns +/// +/// The core [`TaskId`] on success. +/// +/// # Errors +/// +/// Returns an error if: +/// +/// * [`Code::InvalidArgument`] (as [`UnpackError`]) if the task ID is absent or carries an index +/// that cannot be represented. +fn unpack_task_id(task_id: Option) -> Result { + let task_id = task_id.ok_or_else(|| UnpackError { + code: Code::InvalidArgument, + message: "task ID is missing".to_owned(), + })?; + TaskId::try_from(task_id).map_err(|error| UnpackError { + code: Code::InvalidArgument, + message: error.to_string(), + }) +} diff --git a/components/spider-proto/storage/storage.proto b/components/spider-proto/storage/storage.proto index 97f126a45..a686bb2ab 100644 --- a/components/spider-proto/storage/storage.proto +++ b/components/spider-proto/storage/storage.proto @@ -105,10 +105,24 @@ message RegisterTaskInstanceRequest { } message RegisterTaskInstanceResponse { - oneof result { - bytes execution_context = 1; - TaskInstanceManagementError error = 2; - } + ExecutionContext execution_context = 1; +} + +message ExecutionContext { + uint64 task_instance_id = 1; + TdlContext tdl_context = 2; + TimeoutPolicy timeout_policy = 3; + bytes serialized_inputs = 4; +} + +message TdlContext { + string package = 1; + string task_func = 2; +} + +message TimeoutPolicy { + uint64 soft_timeout_ms = 1; + uint64 hard_timeout_ms = 2; } message ReportTaskSuccessRequest { @@ -232,12 +246,7 @@ enum JobState { CANCELLED = 7; } -message TaskInstanceOperationResponse { - oneof result { - Void ok = 1; - TaskInstanceManagementError error = 2; - } -} +message TaskInstanceOperationResponse {} message ResourceGroupOperationResponse { oneof result { @@ -248,20 +257,6 @@ message ResourceGroupOperationResponse { message Void {} -message TaskInstanceManagementError { - enum ErrCode { - ERR_CODE_UNSPECIFIED = 0; - STALE_SESSION = 1; - CACHE_STALE = 2; - SERVER = 3; - INVALID_INPUT = 4; - } - - ErrCode err_code = 1; - string message = 2; - uint64 storage_session = 3; -} - message InboundQueueResponseError { enum ErrCode { ERR_CODE_UNSPECIFIED = 0; diff --git a/components/spider-storage/src/cache/error.rs b/components/spider-storage/src/cache/error.rs index c033ac927..9a2eb4689 100644 --- a/components/spider-storage/src/cache/error.rs +++ b/components/spider-storage/src/cache/error.rs @@ -1,6 +1,8 @@ use spider_core::{job::JobState, task::TaskState, types::id::JobId}; use spider_tdl::wire::WireError; +use crate::db::DbError; + /// Enums for all possible errors that can occur in a cache operation. #[derive(thiserror::Error, Debug)] pub enum CacheError { @@ -9,9 +11,12 @@ pub enum CacheError { #[error(transparent)] StaleState(#[from] StaleStateError), +} - #[error(transparent)] - Db(#[from] crate::db::DbError), +impl From for CacheError { + fn from(err: DbError) -> Self { + Self::Internal(InternalError::Db(err)) + } } /// Enums for all internal errors. @@ -128,7 +133,7 @@ pub enum StaleStateError { #[error("job no longer in the cleanup-ready state")] JobNoLongerCleanupReady, - #[error("job already terminated")] + #[error("job already terminated in state {0:?}")] JobAlreadyTerminated(JobState), #[error("job already requested for cancellation")] diff --git a/components/spider-storage/src/grpc.rs b/components/spider-storage/src/grpc.rs index e808bebd0..6f42d6ba8 100644 --- a/components/spider-storage/src/grpc.rs +++ b/components/spider-storage/src/grpc.rs @@ -1,6 +1,7 @@ //! gRPC service adapters for the storage runtime. use async_trait::async_trait; +use spider_core::types::id::TaskId; use spider_proto_rust::{ storage::{ self, @@ -97,8 +98,7 @@ impl< Status::unavailable("storage service unavailable") } - StorageServerError::Db(db_error) - | StorageServerError::Cache(CacheError::Db(db_error)) => match &db_error { + StorageServerError::Db(db_error) => match &db_error { DbError::ResourceGroupNotFound(_) | DbError::InvalidPassword(_) => { tracing::warn!( error = % db_error, @@ -170,6 +170,95 @@ impl< } } } + + /// Error handler for task instance 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: + /// + /// * `INTERNAL` for: + /// * A fatal cache-internal error (the service will restart). + /// * Any other unexpected error (the service will restart). + /// * `UNAVAILABLE` for a request issued from a stale session. + /// * `FAILED_PRECONDITION` for a request issued against a stale cache state. + /// * `INVALID_ARGUMENT` for malformed inputs or a malformed request. + pub fn task_instance_management_service_error_handler( + &self, + error: StorageServerError, + tag: &'static str, + ) -> Status { + 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") + } + + StorageServerError::StaleSession(storage_session) => { + tracing::warn!( + storage_session, + service = SERVICE_NAME, + tag, + "The request was issued from a stale session." + ); + Status::unavailable(format!( + "stale session; current storage session is {storage_session}" + )) + } + + StorageServerError::Cache(CacheError::StaleState(_)) => { + tracing::warn!( + error = % error, + service = SERVICE_NAME, + tag, + "The request was issued from a stale cache state." + ); + Status::failed_precondition(format!("cache stale: {error}")) + } + + StorageServerError::JobNotFound(job_id) => { + tracing::warn!( + job_id = job_id.get(), + service = SERVICE_NAME, + tag, + "The request attempts to access a job that does not exist in the cache." + ); + // The absence of the job is considered a stale cache state. + Status::failed_precondition(format!("cache stale: {error}")) + } + + error @ (StorageServerError::Tdl(_) | StorageServerError::BadRequest(_)) => { + tracing::warn!( + error = % error, + service = SERVICE_NAME, + tag, + "Invalid argument." + ); + Status::invalid_argument(error.to_string()) + } + + _ => { + 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`]. @@ -283,23 +372,107 @@ impl< { async fn register_task_instance( &self, - _request: Request, + request: Request, ) -> Result, Status> { - todo!("Not implemented") + let (session_id, job_id, task_id, em_id) = request.into_inner().unpack()?; + tracing::info!( + session_id = session_id, + job_id = job_id.get(), + task_id = % task_id, + em_id = em_id.get(), + "Task instance registration request received." + ); + + let execution_context = self + .inner + .create_task_instance(session_id, job_id, task_id, em_id) + .await + .map_err(|error| { + self.task_instance_management_service_error_handler(error, "register_task_instance") + })?; + + Ok(Response::new(storage::RegisterTaskInstanceResponse { + execution_context: Some(storage::ExecutionContext { + task_instance_id: execution_context.task_instance_id, + tdl_context: Some(storage::TdlContext { + package: execution_context.tdl_context.package, + task_func: execution_context.tdl_context.task_func, + }), + timeout_policy: Some(storage::TimeoutPolicy { + soft_timeout_ms: execution_context.timeout_policy.soft_timeout_ms, + hard_timeout_ms: execution_context.timeout_policy.hard_timeout_ms, + }), + serialized_inputs: execution_context.serialized_inputs, + }), + })) } async fn report_task_success( &self, - _request: Request, + request: Request, ) -> Result, Status> { - todo!("Not implemented") + let (session_id, job_id, task_id, task_instance_id, serialized_outputs) = + request.into_inner().unpack()?; + tracing::info!( + session_id = session_id, + job_id = job_id.get(), + task_id = % task_id, + task_instance_id, + "Task instance completion request (success) received." + ); + + let _job_state = match task_id { + TaskId::Index(task_index) => { + self.inner + .succeed_task_instance( + session_id, + job_id, + task_instance_id, + task_index, + serialized_outputs, + ) + .await + } + TaskId::Commit => { + self.inner + .succeed_commit_task_instance(session_id, job_id, task_instance_id) + .await + } + TaskId::Cleanup => { + self.inner + .succeed_cleanup_task_instance(session_id, job_id, task_instance_id) + .await + } + } + .map_err(|error| { + self.task_instance_management_service_error_handler(error, "report_task_success") + })?; + + Ok(Response::new(storage::TaskInstanceOperationResponse {})) } async fn report_task_failure( &self, - _request: Request, + request: Request, ) -> Result, Status> { - todo!("Not implemented") + let (session_id, job_id, task_id, task_instance_id, error_message) = + request.into_inner().unpack()?; + tracing::info!( + session_id = session_id, + job_id = job_id.get(), + task_id = % task_id, + task_instance_id, + "Task instance completion request (failure) received." + ); + + let _job_state = self + .inner + .fail_task_instance(session_id, job_id, task_instance_id, task_id, error_message) + .await + .map_err(|error| { + self.task_instance_management_service_error_handler(error, "report_task_failure") + })?; + Ok(Response::new(storage::TaskInstanceOperationResponse {})) } } diff --git a/components/spider-storage/src/state/error.rs b/components/spider-storage/src/state/error.rs index e8a13670b..4f2ef6d6d 100644 --- a/components/spider-storage/src/state/error.rs +++ b/components/spider-storage/src/state/error.rs @@ -1,4 +1,7 @@ -use spider_core::{task, types::id::JobId}; +use spider_core::{ + task, + types::id::{JobId, SessionId}, +}; use spider_tdl::error::TdlError; use crate::{cache::error::CacheError, db::DbError}; @@ -18,8 +21,8 @@ pub enum StorageServerError { #[error(transparent)] Tdl(#[from] TdlError), - #[error("stale session")] - StaleSession, + #[error("current storage session is {0}")] + StaleSession(SessionId), #[error("server is shutting down: {0}")] Stopping(String), diff --git a/components/spider-storage/src/state/service.rs b/components/spider-storage/src/state/service.rs index a72d615d6..e403e5118 100644 --- a/components/spider-storage/src/state/service.rs +++ b/components/spider-storage/src/state/service.rs @@ -718,7 +718,7 @@ impl< /// Returns [`StorageServerError::StaleSession`] if the session IDs don't match. fn validate_session(&self, session_id: SessionId) -> Result<(), StorageServerError> { if session_id != self.inner.session_id { - return Err(StorageServerError::StaleSession); + return Err(StorageServerError::StaleSession(self.inner.session_id)); } Ok(()) } @@ -1516,7 +1516,7 @@ mod tests { ) .await; assert!( - matches!(result, Err(StorageServerError::StaleSession)), + matches!(result, Err(StorageServerError::StaleSession(_))), "create_task_instance should return StaleSession on session mismatch" ); } @@ -1532,7 +1532,7 @@ mod tests { ) .await; assert!( - matches!(result, Err(StorageServerError::StaleSession)), + matches!(result, Err(StorageServerError::StaleSession(_))), "succeed_task_instance should return StaleSession on session mismatch" ); } @@ -1548,7 +1548,7 @@ mod tests { ) .await; assert!( - matches!(result, Err(StorageServerError::StaleSession)), + matches!(result, Err(StorageServerError::StaleSession(_))), "fail_task_instance should return StaleSession on session mismatch" ); } diff --git a/tests/huntsman/em-runtime/tests/test_runtime.rs b/tests/huntsman/em-runtime/tests/test_runtime.rs index 93f268cdc..5840d890d 100644 --- a/tests/huntsman/em-runtime/tests/test_runtime.rs +++ b/tests/huntsman/em-runtime/tests/test_runtime.rs @@ -391,9 +391,9 @@ async fn stale_session_drops_assignment_and_refreshes() -> anyhow::Result<()> { let baseline = liveness.heartbeat_count(); scheduler.push(Ok(assignment_with_session(STALE_SESSION))); - storage.push_register_response(Err(StorageResponseError::StaleSession { - storage_session: CURRENT_SESSION, - })); + storage.push_register_response(Err(StorageResponseError::StaleSession(format!( + "storage now at {CURRENT_SESSION}" + )))); let join = tokio::spawn(runtime.run()); // Stale-session response triggers liveness refresh and drops the assignment. @@ -437,9 +437,7 @@ async fn recoverable_storage_errors_drop_assignment() -> anyhow::Result<()> { // assignment and poll the scheduler again. let recoverable_errors = [ StorageResponseError::CacheStale("stale cache".to_owned()), - StorageResponseError::StaleSession { - storage_session: SESSION_ID + 1, - }, + StorageResponseError::StaleSession(format!("storage now at {}", SESSION_ID + 1)), ]; let num_errors = recoverable_errors.len() as u64; for err in recoverable_errors {