feat(spider-huntsman): Replace custom error code with gRPC status.

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<storage::ExecutionManagerLivenessError> 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<storage::ExecutionManagerLivenessError> for LivenessResponseError {
 fn register_response_to_result(
     response: storage::RegisterExecutionManagerResponse,
 ) -> Result<RegistrationResponse, LivenessResponseError> {
-    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<SessionId, LivenessResponseError> {
-    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(_)
+        ));
+    }
 }