feat(spider-storage): Implement the task instance management gRPC service and migrate it to Status-based errors.

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<JobIdMarker>;

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<storage::TaskInstanceManagementError> 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:?}"),
         }
     }
 }

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).

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."

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,
 }

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<register_task_instance_response::Result>,
+    #[prost(message, optional, tag = "1")]
+    pub execution_context: ::core::option::Option<ExecutionContext>,
 }
-/// 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<u8>),
-        #[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<TdlContext>,
+    #[prost(message, optional, tag = "3")]
+    pub timeout_policy: ::core::option::Option<TimeoutPolicy>,
+    #[prost(bytes = "vec", tag = "4")]
+    pub serialized_inputs: ::prost::alloc::vec::Vec<u8>,
+}
+#[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<task_instance_operation_response::Result>,
-}
-/// 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<Self> {
-            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,