From ade82fbbeb1665b6c2a0ab4f2f714a5915749b9f Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Wed, 10 Jun 2026 11:58:08 -0400 Subject: [PATCH 01/16] Add storage proto --- components/spider-proto/storage/storage.proto | 177 ++++++++++++++++++ 1 file changed, 177 insertions(+) diff --git a/components/spider-proto/storage/storage.proto b/components/spider-proto/storage/storage.proto index 3d6f5483..d7802d82 100644 --- a/components/spider-proto/storage/storage.proto +++ b/components/spider-proto/storage/storage.proto @@ -2,12 +2,110 @@ syntax = "proto3"; package storage; +service JobManagementService { + rpc SubmitJob(SubmitJobRequest) returns (SubmitJobResponse); + rpc StartJob(JobIdRequest) returns (JobStateResponse); + rpc CancelJob(JobIdRequest) returns (JobStateResponse); + rpc GetJobState(JobIdRequest) returns (JobStateResponse); + rpc GetJobOutputs(JobIdRequest) returns (JobOutputsResponse); + rpc GetJobError(JobIdRequest) returns (JobErrorResponse); + rpc DeleteExpiredTerminatedJobs(DeleteExpiredTerminatedJobsRequest) + returns (DeleteExpiredTerminatedJobsResponse); + rpc ResendReadyTasks(ResendReadyTasksRequest) returns (StorageOperationResponse); +} + service TaskInstanceManagementService { rpc RegisterTaskInstance(RegisterTaskInstanceRequest) returns (RegisterTaskInstanceResponse); rpc ReportTaskSuccess(ReportTaskSuccessRequest) returns (StorageOperationResponse); rpc ReportTaskFailure(ReportTaskFailureRequest) returns (StorageOperationResponse); } +service ResourceGroupManagementService { + rpc AddResourceGroup(AddResourceGroupRequest) returns (ResourceGroupIdResponse); + rpc VerifyResourceGroup(VerifyResourceGroupRequest) returns (StorageOperationResponse); + rpc DeleteResourceGroup(ResourceGroupIdRequest) returns (StorageOperationResponse); +} + +service ExecutionManagerLivenessService { + rpc RegisterExecutionManager(RegisterExecutionManagerRequest) + returns (ExecutionManagerIdResponse); + rpc UpdateExecutionManagerHeartbeat(ExecutionManagerIdRequest) + returns (StorageOperationResponse); + rpc IsExecutionManagerAlive(ExecutionManagerIdRequest) + returns (IsExecutionManagerAliveResponse); + rpc GetDeadExecutionManagers(GetDeadExecutionManagersRequest) + returns (GetDeadExecutionManagersResponse); +} + +service SessionManagementService { + rpc GetSession(Void) returns (GetSessionResponse); +} + +message SubmitJobRequest { + uint64 resource_group_id = 1; + bytes serialized_task_graph = 2; + repeated bytes serialized_inputs = 3; + uint64 session_id = 4; +} + +message SubmitJobResponse { + oneof result { + uint64 job_id = 1; + StorageError error = 2; + } +} + +message JobIdRequest { + uint64 job_id = 1; + uint64 session_id = 2; +} + +message JobStateResponse { + oneof result { + JobState state = 1; + StorageError error = 2; + } +} + +message JobOutputsResponse { + oneof result { + JobOutputs outputs = 1; + StorageError error = 2; + } +} + +message JobOutputs { + repeated bytes outputs = 1; +} + +message JobErrorResponse { + oneof result { + string error_message = 1; + StorageError error = 2; + } +} + +message DeleteExpiredTerminatedJobsRequest { + uint64 expire_after_sec = 1; + uint64 session_id = 2; +} + +message DeleteExpiredTerminatedJobsResponse { + oneof result { + DeletedJobs deleted_jobs = 1; + StorageError error = 2; + } +} + +message DeletedJobs { + repeated uint64 job_ids = 1; +} + +message ResendReadyTasksRequest { + optional uint64 job_id = 1; + uint64 session_id = 2; +} + message RegisterTaskInstanceRequest { uint64 job_id = 1; TaskId task_id = 2; @@ -38,6 +136,74 @@ message ReportTaskFailureRequest { string error_message = 5; } +message AddResourceGroupRequest { + string external_resource_group_id = 1; + bytes password = 2; + uint64 session_id = 3; +} + +message ResourceGroupIdRequest { + uint64 resource_group_id = 1; + uint64 session_id = 2; +} + +message ResourceGroupIdResponse { + oneof result { + uint64 resource_group_id = 1; + StorageError error = 2; + } +} + +message VerifyResourceGroupRequest { + uint64 resource_group_id = 1; + bytes password = 2; + uint64 session_id = 3; +} + +message RegisterExecutionManagerRequest { + string ip_address = 1; + uint64 session_id = 2; +} + +message ExecutionManagerIdRequest { + uint64 execution_manager_id = 1; + uint64 session_id = 2; +} + +message ExecutionManagerIdResponse { + oneof result { + uint64 execution_manager_id = 1; + StorageError error = 2; + } +} + +message IsExecutionManagerAliveResponse { + oneof result { + bool alive = 1; + StorageError error = 2; + } +} + +message GetDeadExecutionManagersRequest { + uint64 stale_after_sec = 1; + uint64 session_id = 2; +} + +message GetDeadExecutionManagersResponse { + oneof result { + DeadExecutionManagers dead_execution_managers = 1; + StorageError error = 2; + } +} + +message DeadExecutionManagers { + repeated uint64 execution_manager_ids = 1; +} + +message GetSessionResponse { + uint64 session_id = 1; +} + message TaskId { oneof kind { uint64 index = 1; @@ -46,6 +212,17 @@ message TaskId { } } +enum JobState { + JOB_STATE_UNSPECIFIED = 0; + READY = 1; + RUNNING = 2; + COMMIT_READY = 3; + CLEANUP_READY = 4; + SUCCEEDED = 5; + FAILED = 6; + CANCELLED = 7; +} + message StorageOperationResponse { oneof result { Void ok = 1; From b1f5153f0db6e110fb75933e9372d641a12edd53 Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Wed, 10 Jun 2026 11:59:42 -0400 Subject: [PATCH 02/16] Add generated file --- .../src/generated/storage.rs | 2740 ++++++++++++++++- 1 file changed, 2626 insertions(+), 114 deletions(-) diff --git a/components/spider-proto-rust/src/generated/storage.rs b/components/spider-proto-rust/src/generated/storage.rs index 4a4cd353..8b6b2c00 100644 --- a/components/spider-proto-rust/src/generated/storage.rs +++ b/components/spider-proto-rust/src/generated/storage.rs @@ -1,4 +1,121 @@ // This file is @generated by prost-build. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct SubmitJobRequest { + #[prost(uint64, tag = "1")] + pub resource_group_id: u64, + #[prost(bytes = "vec", tag = "2")] + pub serialized_task_graph: ::prost::alloc::vec::Vec, + #[prost(bytes = "vec", repeated, tag = "3")] + pub serialized_inputs: ::prost::alloc::vec::Vec<::prost::alloc::vec::Vec>, + #[prost(uint64, tag = "4")] + pub session_id: u64, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct SubmitJobResponse { + #[prost(oneof = "submit_job_response::Result", tags = "1, 2")] + pub result: ::core::option::Option, +} +/// Nested message and enum types in `SubmitJobResponse`. +pub mod submit_job_response { + #[derive(Clone, PartialEq, ::prost::Oneof)] + pub enum Result { + #[prost(uint64, tag = "1")] + JobId(u64), + #[prost(message, tag = "2")] + Error(super::StorageError), + } +} +#[derive(Clone, Copy, PartialEq, ::prost::Message)] +pub struct JobIdRequest { + #[prost(uint64, tag = "1")] + pub job_id: u64, + #[prost(uint64, tag = "2")] + pub session_id: u64, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct JobStateResponse { + #[prost(oneof = "job_state_response::Result", tags = "1, 2")] + pub result: ::core::option::Option, +} +/// Nested message and enum types in `JobStateResponse`. +pub mod job_state_response { + #[derive(Clone, PartialEq, ::prost::Oneof)] + pub enum Result { + #[prost(enumeration = "super::JobState", tag = "1")] + State(i32), + #[prost(message, tag = "2")] + Error(super::StorageError), + } +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct JobOutputsResponse { + #[prost(oneof = "job_outputs_response::Result", tags = "1, 2")] + pub result: ::core::option::Option, +} +/// Nested message and enum types in `JobOutputsResponse`. +pub mod job_outputs_response { + #[derive(Clone, PartialEq, ::prost::Oneof)] + pub enum Result { + #[prost(message, tag = "1")] + Outputs(super::JobOutputs), + #[prost(message, tag = "2")] + Error(super::StorageError), + } +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct JobOutputs { + #[prost(bytes = "vec", repeated, tag = "1")] + pub outputs: ::prost::alloc::vec::Vec<::prost::alloc::vec::Vec>, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct JobErrorResponse { + #[prost(oneof = "job_error_response::Result", tags = "1, 2")] + pub result: ::core::option::Option, +} +/// Nested message and enum types in `JobErrorResponse`. +pub mod job_error_response { + #[derive(Clone, PartialEq, ::prost::Oneof)] + pub enum Result { + #[prost(string, tag = "1")] + ErrorMessage(::prost::alloc::string::String), + #[prost(message, tag = "2")] + Error(super::StorageError), + } +} +#[derive(Clone, Copy, PartialEq, ::prost::Message)] +pub struct DeleteExpiredTerminatedJobsRequest { + #[prost(uint64, tag = "1")] + pub expire_after_sec: u64, + #[prost(uint64, tag = "2")] + pub session_id: u64, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct DeleteExpiredTerminatedJobsResponse { + #[prost(oneof = "delete_expired_terminated_jobs_response::Result", tags = "1, 2")] + pub result: ::core::option::Option, +} +/// Nested message and enum types in `DeleteExpiredTerminatedJobsResponse`. +pub mod delete_expired_terminated_jobs_response { + #[derive(Clone, PartialEq, ::prost::Oneof)] + pub enum Result { + #[prost(message, tag = "1")] + DeletedJobs(super::DeletedJobs), + #[prost(message, tag = "2")] + Error(super::StorageError), + } +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct DeletedJobs { + #[prost(uint64, repeated, tag = "1")] + pub job_ids: ::prost::alloc::vec::Vec, +} +#[derive(Clone, Copy, PartialEq, ::prost::Message)] +pub struct ResendReadyTasksRequest { + #[prost(uint64, optional, tag = "1")] + pub job_id: ::core::option::Option, + #[prost(uint64, tag = "2")] + pub session_id: u64, +} #[derive(Clone, Copy, PartialEq, ::prost::Message)] pub struct RegisterTaskInstanceRequest { #[prost(uint64, tag = "1")] @@ -51,6 +168,122 @@ pub struct ReportTaskFailureRequest { #[prost(string, tag = "5")] pub error_message: ::prost::alloc::string::String, } +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct AddResourceGroupRequest { + #[prost(string, tag = "1")] + 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, Copy, PartialEq, ::prost::Message)] +pub struct ResourceGroupIdRequest { + #[prost(uint64, tag = "1")] + pub resource_group_id: u64, + #[prost(uint64, tag = "2")] + pub session_id: u64, +} +#[derive(Clone, 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::StorageError), + } +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct VerifyResourceGroupRequest { + #[prost(uint64, tag = "1")] + 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 { + #[prost(string, tag = "1")] + pub ip_address: ::prost::alloc::string::String, + #[prost(uint64, tag = "2")] + pub session_id: u64, +} +#[derive(Clone, Copy, PartialEq, ::prost::Message)] +pub struct ExecutionManagerIdRequest { + #[prost(uint64, tag = "1")] + pub execution_manager_id: u64, + #[prost(uint64, tag = "2")] + pub session_id: u64, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ExecutionManagerIdResponse { + #[prost(oneof = "execution_manager_id_response::Result", tags = "1, 2")] + pub result: ::core::option::Option, +} +/// Nested message and enum types in `ExecutionManagerIdResponse`. +pub mod execution_manager_id_response { + #[derive(Clone, PartialEq, ::prost::Oneof)] + pub enum Result { + #[prost(uint64, tag = "1")] + ExecutionManagerId(u64), + #[prost(message, tag = "2")] + Error(super::StorageError), + } +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct IsExecutionManagerAliveResponse { + #[prost(oneof = "is_execution_manager_alive_response::Result", tags = "1, 2")] + pub result: ::core::option::Option, +} +/// Nested message and enum types in `IsExecutionManagerAliveResponse`. +pub mod is_execution_manager_alive_response { + #[derive(Clone, PartialEq, ::prost::Oneof)] + pub enum Result { + #[prost(bool, tag = "1")] + Alive(bool), + #[prost(message, tag = "2")] + Error(super::StorageError), + } +} +#[derive(Clone, Copy, PartialEq, ::prost::Message)] +pub struct GetDeadExecutionManagersRequest { + #[prost(uint64, tag = "1")] + pub stale_after_sec: u64, + #[prost(uint64, tag = "2")] + pub session_id: u64, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct GetDeadExecutionManagersResponse { + #[prost(oneof = "get_dead_execution_managers_response::Result", tags = "1, 2")] + pub result: ::core::option::Option, +} +/// Nested message and enum types in `GetDeadExecutionManagersResponse`. +pub mod get_dead_execution_managers_response { + #[derive(Clone, PartialEq, ::prost::Oneof)] + pub enum Result { + #[prost(message, tag = "1")] + DeadExecutionManagers(super::DeadExecutionManagers), + #[prost(message, tag = "2")] + Error(super::StorageError), + } +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct DeadExecutionManagers { + #[prost(uint64, repeated, tag = "1")] + pub execution_manager_ids: ::prost::alloc::vec::Vec, +} +#[derive(Clone, Copy, PartialEq, ::prost::Message)] +pub struct GetSessionResponse { + #[prost(uint64, tag = "1")] + pub session_id: u64, +} #[derive(Clone, Copy, PartialEq, ::prost::Message)] pub struct TaskId { #[prost(oneof = "task_id::Kind", tags = "1, 2, 3")] @@ -145,8 +378,52 @@ pub mod storage_error { } } } +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] +#[repr(i32)] +pub enum JobState { + Unspecified = 0, + Ready = 1, + Running = 2, + CommitReady = 3, + CleanupReady = 4, + Succeeded = 5, + Failed = 6, + Cancelled = 7, +} +impl JobState { + /// 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 => "JOB_STATE_UNSPECIFIED", + Self::Ready => "READY", + Self::Running => "RUNNING", + Self::CommitReady => "COMMIT_READY", + Self::CleanupReady => "CLEANUP_READY", + Self::Succeeded => "SUCCEEDED", + Self::Failed => "FAILED", + Self::Cancelled => "CANCELLED", + } + } + /// Creates an enum from field names used in the ProtoBuf definition. + pub fn from_str_name(value: &str) -> ::core::option::Option { + match value { + "JOB_STATE_UNSPECIFIED" => Some(Self::Unspecified), + "READY" => Some(Self::Ready), + "RUNNING" => Some(Self::Running), + "COMMIT_READY" => Some(Self::CommitReady), + "CLEANUP_READY" => Some(Self::CleanupReady), + "SUCCEEDED" => Some(Self::Succeeded), + "FAILED" => Some(Self::Failed), + "CANCELLED" => Some(Self::Cancelled), + _ => None, + } + } +} /// Generated client implementations. -pub mod task_instance_management_service_client { +pub mod job_management_service_client { #![allow( unused_variables, dead_code, @@ -157,10 +434,10 @@ pub mod task_instance_management_service_client { use tonic::codegen::*; use tonic::codegen::http::Uri; #[derive(Debug, Clone)] - pub struct TaskInstanceManagementServiceClient { + pub struct JobManagementServiceClient { inner: tonic::client::Grpc, } - impl TaskInstanceManagementServiceClient { + impl JobManagementServiceClient { /// Attempt to create a new client by connecting to a given endpoint. pub async fn connect(dst: D) -> Result where @@ -171,7 +448,7 @@ pub mod task_instance_management_service_client { Ok(Self::new(conn)) } } - impl TaskInstanceManagementServiceClient + impl JobManagementServiceClient where T: tonic::client::GrpcService, T::Error: Into, @@ -189,7 +466,7 @@ pub mod task_instance_management_service_client { pub fn with_interceptor( inner: T, interceptor: F, - ) -> TaskInstanceManagementServiceClient> + ) -> JobManagementServiceClient> where F: tonic::service::Interceptor, T::ResponseBody: Default, @@ -203,9 +480,7 @@ pub mod task_instance_management_service_client { http::Request, >>::Error: Into + std::marker::Send + std::marker::Sync, { - TaskInstanceManagementServiceClient::new( - InterceptedService::new(inner, interceptor), - ) + JobManagementServiceClient::new(InterceptedService::new(inner, interceptor)) } /// Compress requests with the given encoding. /// @@ -238,11 +513,11 @@ pub mod task_instance_management_service_client { self.inner = self.inner.max_encoding_message_size(limit); self } - pub async fn register_task_instance( + pub async fn submit_job( &mut self, - request: impl tonic::IntoRequest, + request: impl tonic::IntoRequest, ) -> std::result::Result< - tonic::Response, + tonic::Response, tonic::Status, > { self.inner @@ -255,23 +530,18 @@ pub mod task_instance_management_service_client { })?; let codec = tonic::codec::ProstCodec::default(); let path = http::uri::PathAndQuery::from_static( - "/storage.TaskInstanceManagementService/RegisterTaskInstance", + "/storage.JobManagementService/SubmitJob", ); let mut req = request.into_request(); req.extensions_mut() - .insert( - GrpcMethod::new( - "storage.TaskInstanceManagementService", - "RegisterTaskInstance", - ), - ); + .insert(GrpcMethod::new("storage.JobManagementService", "SubmitJob")); self.inner.unary(req, path, codec).await } - pub async fn report_task_success( + pub async fn start_job( &mut self, - request: impl tonic::IntoRequest, + request: impl tonic::IntoRequest, ) -> std::result::Result< - tonic::Response, + tonic::Response, tonic::Status, > { self.inner @@ -284,23 +554,18 @@ pub mod task_instance_management_service_client { })?; let codec = tonic::codec::ProstCodec::default(); let path = http::uri::PathAndQuery::from_static( - "/storage.TaskInstanceManagementService/ReportTaskSuccess", + "/storage.JobManagementService/StartJob", ); let mut req = request.into_request(); req.extensions_mut() - .insert( - GrpcMethod::new( - "storage.TaskInstanceManagementService", - "ReportTaskSuccess", - ), - ); + .insert(GrpcMethod::new("storage.JobManagementService", "StartJob")); self.inner.unary(req, path, codec).await } - pub async fn report_task_failure( + pub async fn cancel_job( &mut self, - request: impl tonic::IntoRequest, + request: impl tonic::IntoRequest, ) -> std::result::Result< - tonic::Response, + tonic::Response, tonic::Status, > { self.inner @@ -313,69 +578,2069 @@ pub mod task_instance_management_service_client { })?; let codec = tonic::codec::ProstCodec::default(); let path = http::uri::PathAndQuery::from_static( - "/storage.TaskInstanceManagementService/ReportTaskFailure", + "/storage.JobManagementService/CancelJob", ); let mut req = request.into_request(); req.extensions_mut() - .insert( - GrpcMethod::new( - "storage.TaskInstanceManagementService", - "ReportTaskFailure", - ), - ); + .insert(GrpcMethod::new("storage.JobManagementService", "CancelJob")); self.inner.unary(req, path, codec).await } - } -} -/// Generated server implementations. -pub mod task_instance_management_service_server { - #![allow( - unused_variables, - dead_code, - missing_docs, - clippy::wildcard_imports, - clippy::let_unit_value, - )] - use tonic::codegen::*; - /// Generated trait containing gRPC methods that should be implemented for use with TaskInstanceManagementServiceServer. - #[async_trait] - pub trait TaskInstanceManagementService: std::marker::Send + std::marker::Sync + 'static { - async fn register_task_instance( - &self, - request: tonic::Request, + pub async fn get_job_state( + &mut self, + request: impl tonic::IntoRequest, ) -> std::result::Result< - tonic::Response, + tonic::Response, tonic::Status, - >; - async fn report_task_success( - &self, - request: tonic::Request, + > { + self.inner + .ready() + .await + .map_err(|e| { + tonic::Status::unknown( + format!("Service was not ready: {}", e.into()), + ) + })?; + let codec = tonic::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/storage.JobManagementService/GetJobState", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert(GrpcMethod::new("storage.JobManagementService", "GetJobState")); + self.inner.unary(req, path, codec).await + } + pub async fn get_job_outputs( + &mut self, + request: impl tonic::IntoRequest, ) -> std::result::Result< - tonic::Response, + tonic::Response, tonic::Status, - >; - async fn report_task_failure( - &self, - request: tonic::Request, + > { + self.inner + .ready() + .await + .map_err(|e| { + tonic::Status::unknown( + format!("Service was not ready: {}", e.into()), + ) + })?; + let codec = tonic::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/storage.JobManagementService/GetJobOutputs", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert( + GrpcMethod::new("storage.JobManagementService", "GetJobOutputs"), + ); + self.inner.unary(req, path, codec).await + } + pub async fn get_job_error( + &mut self, + request: impl tonic::IntoRequest, ) -> std::result::Result< - tonic::Response, + tonic::Response, tonic::Status, - >; - } - #[derive(Debug)] - pub struct TaskInstanceManagementServiceServer { - inner: Arc, - accept_compression_encodings: EnabledCompressionEncodings, - send_compression_encodings: EnabledCompressionEncodings, - max_decoding_message_size: Option, - max_encoding_message_size: Option, - } - impl TaskInstanceManagementServiceServer { - pub fn new(inner: T) -> Self { - Self::from_arc(Arc::new(inner)) - } - pub fn from_arc(inner: Arc) -> Self { - Self { + > { + self.inner + .ready() + .await + .map_err(|e| { + tonic::Status::unknown( + format!("Service was not ready: {}", e.into()), + ) + })?; + let codec = tonic::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/storage.JobManagementService/GetJobError", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert(GrpcMethod::new("storage.JobManagementService", "GetJobError")); + self.inner.unary(req, path, codec).await + } + pub async fn delete_expired_terminated_jobs( + &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::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/storage.JobManagementService/DeleteExpiredTerminatedJobs", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert( + GrpcMethod::new( + "storage.JobManagementService", + "DeleteExpiredTerminatedJobs", + ), + ); + 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::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/storage.JobManagementService/ResendReadyTasks", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert( + GrpcMethod::new("storage.JobManagementService", "ResendReadyTasks"), + ); + self.inner.unary(req, path, codec).await + } + } +} +/// Generated client implementations. +pub mod task_instance_management_service_client { + #![allow( + unused_variables, + dead_code, + missing_docs, + clippy::wildcard_imports, + clippy::let_unit_value, + )] + use tonic::codegen::*; + use tonic::codegen::http::Uri; + #[derive(Debug, Clone)] + pub struct TaskInstanceManagementServiceClient { + inner: tonic::client::Grpc, + } + impl TaskInstanceManagementServiceClient { + /// Attempt to create a new client by connecting to a given endpoint. + pub async fn connect(dst: D) -> Result + where + D: TryInto, + D::Error: Into, + { + let conn = tonic::transport::Endpoint::new(dst)?.connect().await?; + Ok(Self::new(conn)) + } + } + impl TaskInstanceManagementServiceClient + where + T: tonic::client::GrpcService, + T::Error: Into, + T::ResponseBody: Body + std::marker::Send + 'static, + ::Error: Into + std::marker::Send, + { + pub fn new(inner: T) -> Self { + let inner = tonic::client::Grpc::new(inner); + Self { inner } + } + pub fn with_origin(inner: T, origin: Uri) -> Self { + let inner = tonic::client::Grpc::with_origin(inner, origin); + Self { inner } + } + pub fn with_interceptor( + inner: T, + interceptor: F, + ) -> TaskInstanceManagementServiceClient> + where + F: tonic::service::Interceptor, + T::ResponseBody: Default, + T: tonic::codegen::Service< + http::Request, + Response = http::Response< + >::ResponseBody, + >, + >, + , + >>::Error: Into + std::marker::Send + std::marker::Sync, + { + TaskInstanceManagementServiceClient::new( + InterceptedService::new(inner, interceptor), + ) + } + /// Compress requests with the given encoding. + /// + /// This requires the server to support it otherwise it might respond with an + /// error. + #[must_use] + pub fn send_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.inner = self.inner.send_compressed(encoding); + self + } + /// Enable decompressing responses. + #[must_use] + pub fn accept_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.inner = self.inner.accept_compressed(encoding); + self + } + /// Limits the maximum size of a decoded message. + /// + /// Default: `4MB` + #[must_use] + pub fn max_decoding_message_size(mut self, limit: usize) -> Self { + self.inner = self.inner.max_decoding_message_size(limit); + self + } + /// Limits the maximum size of an encoded message. + /// + /// Default: `usize::MAX` + #[must_use] + pub fn max_encoding_message_size(mut self, limit: usize) -> Self { + self.inner = self.inner.max_encoding_message_size(limit); + self + } + pub async fn register_task_instance( + &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::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/storage.TaskInstanceManagementService/RegisterTaskInstance", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert( + GrpcMethod::new( + "storage.TaskInstanceManagementService", + "RegisterTaskInstance", + ), + ); + self.inner.unary(req, path, codec).await + } + pub async fn report_task_success( + &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::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/storage.TaskInstanceManagementService/ReportTaskSuccess", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert( + GrpcMethod::new( + "storage.TaskInstanceManagementService", + "ReportTaskSuccess", + ), + ); + self.inner.unary(req, path, codec).await + } + pub async fn report_task_failure( + &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::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/storage.TaskInstanceManagementService/ReportTaskFailure", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert( + GrpcMethod::new( + "storage.TaskInstanceManagementService", + "ReportTaskFailure", + ), + ); + self.inner.unary(req, path, codec).await + } + } +} +/// Generated client implementations. +pub mod resource_group_management_service_client { + #![allow( + unused_variables, + dead_code, + missing_docs, + clippy::wildcard_imports, + clippy::let_unit_value, + )] + use tonic::codegen::*; + use tonic::codegen::http::Uri; + #[derive(Debug, Clone)] + pub struct ResourceGroupManagementServiceClient { + inner: tonic::client::Grpc, + } + impl ResourceGroupManagementServiceClient { + /// Attempt to create a new client by connecting to a given endpoint. + pub async fn connect(dst: D) -> Result + where + D: TryInto, + D::Error: Into, + { + let conn = tonic::transport::Endpoint::new(dst)?.connect().await?; + Ok(Self::new(conn)) + } + } + impl ResourceGroupManagementServiceClient + where + T: tonic::client::GrpcService, + T::Error: Into, + T::ResponseBody: Body + std::marker::Send + 'static, + ::Error: Into + std::marker::Send, + { + pub fn new(inner: T) -> Self { + let inner = tonic::client::Grpc::new(inner); + Self { inner } + } + pub fn with_origin(inner: T, origin: Uri) -> Self { + let inner = tonic::client::Grpc::with_origin(inner, origin); + Self { inner } + } + pub fn with_interceptor( + inner: T, + interceptor: F, + ) -> ResourceGroupManagementServiceClient> + where + F: tonic::service::Interceptor, + T::ResponseBody: Default, + T: tonic::codegen::Service< + http::Request, + Response = http::Response< + >::ResponseBody, + >, + >, + , + >>::Error: Into + std::marker::Send + std::marker::Sync, + { + ResourceGroupManagementServiceClient::new( + InterceptedService::new(inner, interceptor), + ) + } + /// Compress requests with the given encoding. + /// + /// This requires the server to support it otherwise it might respond with an + /// error. + #[must_use] + pub fn send_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.inner = self.inner.send_compressed(encoding); + self + } + /// Enable decompressing responses. + #[must_use] + pub fn accept_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.inner = self.inner.accept_compressed(encoding); + self + } + /// Limits the maximum size of a decoded message. + /// + /// Default: `4MB` + #[must_use] + pub fn max_decoding_message_size(mut self, limit: usize) -> Self { + self.inner = self.inner.max_decoding_message_size(limit); + self + } + /// Limits the maximum size of an encoded message. + /// + /// Default: `usize::MAX` + #[must_use] + pub fn max_encoding_message_size(mut self, limit: usize) -> Self { + self.inner = self.inner.max_encoding_message_size(limit); + self + } + pub async fn add_resource_group( + &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::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/storage.ResourceGroupManagementService/AddResourceGroup", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert( + GrpcMethod::new( + "storage.ResourceGroupManagementService", + "AddResourceGroup", + ), + ); + self.inner.unary(req, path, codec).await + } + pub async fn verify_resource_group( + &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::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/storage.ResourceGroupManagementService/VerifyResourceGroup", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert( + GrpcMethod::new( + "storage.ResourceGroupManagementService", + "VerifyResourceGroup", + ), + ); + self.inner.unary(req, path, codec).await + } + pub async fn delete_resource_group( + &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::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/storage.ResourceGroupManagementService/DeleteResourceGroup", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert( + GrpcMethod::new( + "storage.ResourceGroupManagementService", + "DeleteResourceGroup", + ), + ); + self.inner.unary(req, path, codec).await + } + } +} +/// Generated client implementations. +pub mod execution_manager_liveness_service_client { + #![allow( + unused_variables, + dead_code, + missing_docs, + clippy::wildcard_imports, + clippy::let_unit_value, + )] + use tonic::codegen::*; + use tonic::codegen::http::Uri; + #[derive(Debug, Clone)] + pub struct ExecutionManagerLivenessServiceClient { + inner: tonic::client::Grpc, + } + impl ExecutionManagerLivenessServiceClient { + /// Attempt to create a new client by connecting to a given endpoint. + pub async fn connect(dst: D) -> Result + where + D: TryInto, + D::Error: Into, + { + let conn = tonic::transport::Endpoint::new(dst)?.connect().await?; + Ok(Self::new(conn)) + } + } + impl ExecutionManagerLivenessServiceClient + where + T: tonic::client::GrpcService, + T::Error: Into, + T::ResponseBody: Body + std::marker::Send + 'static, + ::Error: Into + std::marker::Send, + { + pub fn new(inner: T) -> Self { + let inner = tonic::client::Grpc::new(inner); + Self { inner } + } + pub fn with_origin(inner: T, origin: Uri) -> Self { + let inner = tonic::client::Grpc::with_origin(inner, origin); + Self { inner } + } + pub fn with_interceptor( + inner: T, + interceptor: F, + ) -> ExecutionManagerLivenessServiceClient> + where + F: tonic::service::Interceptor, + T::ResponseBody: Default, + T: tonic::codegen::Service< + http::Request, + Response = http::Response< + >::ResponseBody, + >, + >, + , + >>::Error: Into + std::marker::Send + std::marker::Sync, + { + ExecutionManagerLivenessServiceClient::new( + InterceptedService::new(inner, interceptor), + ) + } + /// Compress requests with the given encoding. + /// + /// This requires the server to support it otherwise it might respond with an + /// error. + #[must_use] + pub fn send_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.inner = self.inner.send_compressed(encoding); + self + } + /// Enable decompressing responses. + #[must_use] + pub fn accept_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.inner = self.inner.accept_compressed(encoding); + self + } + /// Limits the maximum size of a decoded message. + /// + /// Default: `4MB` + #[must_use] + pub fn max_decoding_message_size(mut self, limit: usize) -> Self { + self.inner = self.inner.max_decoding_message_size(limit); + self + } + /// Limits the maximum size of an encoded message. + /// + /// Default: `usize::MAX` + #[must_use] + pub fn max_encoding_message_size(mut self, limit: usize) -> Self { + self.inner = self.inner.max_encoding_message_size(limit); + self + } + pub async fn register_execution_manager( + &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::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/storage.ExecutionManagerLivenessService/RegisterExecutionManager", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert( + GrpcMethod::new( + "storage.ExecutionManagerLivenessService", + "RegisterExecutionManager", + ), + ); + self.inner.unary(req, path, codec).await + } + pub async fn update_execution_manager_heartbeat( + &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::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/storage.ExecutionManagerLivenessService/UpdateExecutionManagerHeartbeat", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert( + GrpcMethod::new( + "storage.ExecutionManagerLivenessService", + "UpdateExecutionManagerHeartbeat", + ), + ); + self.inner.unary(req, path, codec).await + } + pub async fn is_execution_manager_alive( + &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::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/storage.ExecutionManagerLivenessService/IsExecutionManagerAlive", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert( + GrpcMethod::new( + "storage.ExecutionManagerLivenessService", + "IsExecutionManagerAlive", + ), + ); + self.inner.unary(req, path, codec).await + } + pub async fn get_dead_execution_managers( + &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::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/storage.ExecutionManagerLivenessService/GetDeadExecutionManagers", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert( + GrpcMethod::new( + "storage.ExecutionManagerLivenessService", + "GetDeadExecutionManagers", + ), + ); + self.inner.unary(req, path, codec).await + } + } +} +/// Generated client implementations. +pub mod session_management_service_client { + #![allow( + unused_variables, + dead_code, + missing_docs, + clippy::wildcard_imports, + clippy::let_unit_value, + )] + use tonic::codegen::*; + use tonic::codegen::http::Uri; + #[derive(Debug, Clone)] + pub struct SessionManagementServiceClient { + inner: tonic::client::Grpc, + } + impl SessionManagementServiceClient { + /// Attempt to create a new client by connecting to a given endpoint. + pub async fn connect(dst: D) -> Result + where + D: TryInto, + D::Error: Into, + { + let conn = tonic::transport::Endpoint::new(dst)?.connect().await?; + Ok(Self::new(conn)) + } + } + impl SessionManagementServiceClient + where + T: tonic::client::GrpcService, + T::Error: Into, + T::ResponseBody: Body + std::marker::Send + 'static, + ::Error: Into + std::marker::Send, + { + pub fn new(inner: T) -> Self { + let inner = tonic::client::Grpc::new(inner); + Self { inner } + } + pub fn with_origin(inner: T, origin: Uri) -> Self { + let inner = tonic::client::Grpc::with_origin(inner, origin); + Self { inner } + } + pub fn with_interceptor( + inner: T, + interceptor: F, + ) -> SessionManagementServiceClient> + where + F: tonic::service::Interceptor, + T::ResponseBody: Default, + T: tonic::codegen::Service< + http::Request, + Response = http::Response< + >::ResponseBody, + >, + >, + , + >>::Error: Into + std::marker::Send + std::marker::Sync, + { + SessionManagementServiceClient::new( + InterceptedService::new(inner, interceptor), + ) + } + /// Compress requests with the given encoding. + /// + /// This requires the server to support it otherwise it might respond with an + /// error. + #[must_use] + pub fn send_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.inner = self.inner.send_compressed(encoding); + self + } + /// Enable decompressing responses. + #[must_use] + pub fn accept_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.inner = self.inner.accept_compressed(encoding); + self + } + /// Limits the maximum size of a decoded message. + /// + /// Default: `4MB` + #[must_use] + pub fn max_decoding_message_size(mut self, limit: usize) -> Self { + self.inner = self.inner.max_decoding_message_size(limit); + self + } + /// Limits the maximum size of an encoded message. + /// + /// Default: `usize::MAX` + #[must_use] + pub fn max_encoding_message_size(mut self, limit: usize) -> Self { + self.inner = self.inner.max_encoding_message_size(limit); + self + } + pub async fn get_session( + &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::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/storage.SessionManagementService/GetSession", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert( + GrpcMethod::new("storage.SessionManagementService", "GetSession"), + ); + self.inner.unary(req, path, codec).await + } + } +} +/// Generated server implementations. +pub mod job_management_service_server { + #![allow( + unused_variables, + dead_code, + missing_docs, + clippy::wildcard_imports, + clippy::let_unit_value, + )] + use tonic::codegen::*; + /// Generated trait containing gRPC methods that should be implemented for use with JobManagementServiceServer. + #[async_trait] + pub trait JobManagementService: std::marker::Send + std::marker::Sync + 'static { + async fn submit_job( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; + async fn start_job( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; + async fn cancel_job( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; + async fn get_job_state( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; + async fn get_job_outputs( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; + async fn get_job_error( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; + async fn delete_expired_terminated_jobs( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; + async fn resend_ready_tasks( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; + } + #[derive(Debug)] + pub struct JobManagementServiceServer { + inner: Arc, + accept_compression_encodings: EnabledCompressionEncodings, + send_compression_encodings: EnabledCompressionEncodings, + max_decoding_message_size: Option, + max_encoding_message_size: Option, + } + impl JobManagementServiceServer { + pub fn new(inner: T) -> Self { + Self::from_arc(Arc::new(inner)) + } + pub fn from_arc(inner: Arc) -> Self { + Self { + inner, + accept_compression_encodings: Default::default(), + send_compression_encodings: Default::default(), + max_decoding_message_size: None, + max_encoding_message_size: None, + } + } + pub fn with_interceptor( + inner: T, + interceptor: F, + ) -> InterceptedService + where + F: tonic::service::Interceptor, + { + InterceptedService::new(Self::new(inner), interceptor) + } + /// Enable decompressing requests with the given encoding. + #[must_use] + pub fn accept_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.accept_compression_encodings.enable(encoding); + self + } + /// Compress responses with the given encoding, if the client supports it. + #[must_use] + pub fn send_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.send_compression_encodings.enable(encoding); + self + } + /// Limits the maximum size of a decoded message. + /// + /// Default: `4MB` + #[must_use] + pub fn max_decoding_message_size(mut self, limit: usize) -> Self { + self.max_decoding_message_size = Some(limit); + self + } + /// Limits the maximum size of an encoded message. + /// + /// Default: `usize::MAX` + #[must_use] + pub fn max_encoding_message_size(mut self, limit: usize) -> Self { + self.max_encoding_message_size = Some(limit); + self + } + } + impl tonic::codegen::Service> + for JobManagementServiceServer + where + T: JobManagementService, + B: Body + std::marker::Send + 'static, + B::Error: Into + std::marker::Send + 'static, + { + type Response = http::Response; + type Error = std::convert::Infallible; + type Future = BoxFuture; + fn poll_ready( + &mut self, + _cx: &mut Context<'_>, + ) -> Poll> { + Poll::Ready(Ok(())) + } + fn call(&mut self, req: http::Request) -> Self::Future { + match req.uri().path() { + "/storage.JobManagementService/SubmitJob" => { + #[allow(non_camel_case_types)] + struct SubmitJobSvc(pub Arc); + impl< + T: JobManagementService, + > tonic::server::UnaryService + for SubmitJobSvc { + type Response = super::SubmitJobResponse; + 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 { + ::submit_job(&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 = SubmitJobSvc(inner); + let codec = tonic::codec::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) + } + "/storage.JobManagementService/StartJob" => { + #[allow(non_camel_case_types)] + struct StartJobSvc(pub Arc); + impl< + T: JobManagementService, + > tonic::server::UnaryService + for StartJobSvc { + type Response = super::JobStateResponse; + 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 { + ::start_job(&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 = StartJobSvc(inner); + let codec = tonic::codec::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) + } + "/storage.JobManagementService/CancelJob" => { + #[allow(non_camel_case_types)] + struct CancelJobSvc(pub Arc); + impl< + T: JobManagementService, + > tonic::server::UnaryService + for CancelJobSvc { + type Response = super::JobStateResponse; + 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 { + ::cancel_job(&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 = CancelJobSvc(inner); + let codec = tonic::codec::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) + } + "/storage.JobManagementService/GetJobState" => { + #[allow(non_camel_case_types)] + struct GetJobStateSvc(pub Arc); + impl< + T: JobManagementService, + > tonic::server::UnaryService + for GetJobStateSvc { + type Response = super::JobStateResponse; + 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 { + ::get_job_state(&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 = GetJobStateSvc(inner); + let codec = tonic::codec::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) + } + "/storage.JobManagementService/GetJobOutputs" => { + #[allow(non_camel_case_types)] + struct GetJobOutputsSvc(pub Arc); + impl< + T: JobManagementService, + > tonic::server::UnaryService + for GetJobOutputsSvc { + type Response = super::JobOutputsResponse; + 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 { + ::get_job_outputs( + &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 = GetJobOutputsSvc(inner); + let codec = tonic::codec::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) + } + "/storage.JobManagementService/GetJobError" => { + #[allow(non_camel_case_types)] + struct GetJobErrorSvc(pub Arc); + impl< + T: JobManagementService, + > tonic::server::UnaryService + for GetJobErrorSvc { + type Response = super::JobErrorResponse; + 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 { + ::get_job_error(&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 = GetJobErrorSvc(inner); + let codec = tonic::codec::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) + } + "/storage.JobManagementService/DeleteExpiredTerminatedJobs" => { + #[allow(non_camel_case_types)] + struct DeleteExpiredTerminatedJobsSvc( + pub Arc, + ); + impl< + T: JobManagementService, + > tonic::server::UnaryService< + super::DeleteExpiredTerminatedJobsRequest, + > for DeleteExpiredTerminatedJobsSvc { + type Response = super::DeleteExpiredTerminatedJobsResponse; + type Future = BoxFuture< + tonic::Response, + tonic::Status, + >; + fn call( + &mut self, + request: tonic::Request< + super::DeleteExpiredTerminatedJobsRequest, + >, + ) -> Self::Future { + let inner = Arc::clone(&self.0); + let fut = async move { + ::delete_expired_terminated_jobs( + &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 = DeleteExpiredTerminatedJobsSvc(inner); + let codec = tonic::codec::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) + } + "/storage.JobManagementService/ResendReadyTasks" => { + #[allow(non_camel_case_types)] + struct ResendReadyTasksSvc(pub Arc); + impl< + T: JobManagementService, + > tonic::server::UnaryService + for ResendReadyTasksSvc { + type Response = super::StorageOperationResponse; + 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::codec::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(empty_body()); + let headers = response.headers_mut(); + headers + .insert( + tonic::Status::GRPC_STATUS, + (tonic::Code::Unimplemented as i32).into(), + ); + headers + .insert( + http::header::CONTENT_TYPE, + tonic::metadata::GRPC_CONTENT_TYPE, + ); + Ok(response) + }) + } + } + } + } + impl Clone for JobManagementServiceServer { + fn clone(&self) -> Self { + let inner = self.inner.clone(); + Self { + inner, + accept_compression_encodings: self.accept_compression_encodings, + send_compression_encodings: self.send_compression_encodings, + max_decoding_message_size: self.max_decoding_message_size, + max_encoding_message_size: self.max_encoding_message_size, + } + } + } + /// Generated gRPC service name + pub const SERVICE_NAME: &str = "storage.JobManagementService"; + impl tonic::server::NamedService for JobManagementServiceServer { + const NAME: &'static str = SERVICE_NAME; + } +} +/// Generated server implementations. +pub mod task_instance_management_service_server { + #![allow( + unused_variables, + dead_code, + missing_docs, + clippy::wildcard_imports, + clippy::let_unit_value, + )] + use tonic::codegen::*; + /// Generated trait containing gRPC methods that should be implemented for use with TaskInstanceManagementServiceServer. + #[async_trait] + pub trait TaskInstanceManagementService: std::marker::Send + std::marker::Sync + 'static { + async fn register_task_instance( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; + async fn report_task_success( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; + async fn report_task_failure( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; + } + #[derive(Debug)] + pub struct TaskInstanceManagementServiceServer { + inner: Arc, + accept_compression_encodings: EnabledCompressionEncodings, + send_compression_encodings: EnabledCompressionEncodings, + max_decoding_message_size: Option, + max_encoding_message_size: Option, + } + impl TaskInstanceManagementServiceServer { + pub fn new(inner: T) -> Self { + Self::from_arc(Arc::new(inner)) + } + pub fn from_arc(inner: Arc) -> Self { + Self { + inner, + accept_compression_encodings: Default::default(), + send_compression_encodings: Default::default(), + max_decoding_message_size: None, + max_encoding_message_size: None, + } + } + pub fn with_interceptor( + inner: T, + interceptor: F, + ) -> InterceptedService + where + F: tonic::service::Interceptor, + { + InterceptedService::new(Self::new(inner), interceptor) + } + /// Enable decompressing requests with the given encoding. + #[must_use] + pub fn accept_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.accept_compression_encodings.enable(encoding); + self + } + /// Compress responses with the given encoding, if the client supports it. + #[must_use] + pub fn send_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.send_compression_encodings.enable(encoding); + self + } + /// Limits the maximum size of a decoded message. + /// + /// Default: `4MB` + #[must_use] + pub fn max_decoding_message_size(mut self, limit: usize) -> Self { + self.max_decoding_message_size = Some(limit); + self + } + /// Limits the maximum size of an encoded message. + /// + /// Default: `usize::MAX` + #[must_use] + pub fn max_encoding_message_size(mut self, limit: usize) -> Self { + self.max_encoding_message_size = Some(limit); + self + } + } + impl tonic::codegen::Service> + for TaskInstanceManagementServiceServer + where + T: TaskInstanceManagementService, + B: Body + std::marker::Send + 'static, + B::Error: Into + std::marker::Send + 'static, + { + type Response = http::Response; + type Error = std::convert::Infallible; + type Future = BoxFuture; + fn poll_ready( + &mut self, + _cx: &mut Context<'_>, + ) -> Poll> { + Poll::Ready(Ok(())) + } + fn call(&mut self, req: http::Request) -> Self::Future { + match req.uri().path() { + "/storage.TaskInstanceManagementService/RegisterTaskInstance" => { + #[allow(non_camel_case_types)] + struct RegisterTaskInstanceSvc( + pub Arc, + ); + impl< + T: TaskInstanceManagementService, + > tonic::server::UnaryService + for RegisterTaskInstanceSvc { + type Response = super::RegisterTaskInstanceResponse; + 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 { + ::register_task_instance( + &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 = RegisterTaskInstanceSvc(inner); + let codec = tonic::codec::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) + } + "/storage.TaskInstanceManagementService/ReportTaskSuccess" => { + #[allow(non_camel_case_types)] + struct ReportTaskSuccessSvc( + pub Arc, + ); + impl< + T: TaskInstanceManagementService, + > tonic::server::UnaryService + for ReportTaskSuccessSvc { + type Response = super::StorageOperationResponse; + 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 { + ::report_task_success( + &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 = ReportTaskSuccessSvc(inner); + let codec = tonic::codec::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) + } + "/storage.TaskInstanceManagementService/ReportTaskFailure" => { + #[allow(non_camel_case_types)] + struct ReportTaskFailureSvc( + pub Arc, + ); + impl< + T: TaskInstanceManagementService, + > tonic::server::UnaryService + for ReportTaskFailureSvc { + type Response = super::StorageOperationResponse; + 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 { + ::report_task_failure( + &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 = ReportTaskFailureSvc(inner); + let codec = tonic::codec::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(empty_body()); + let headers = response.headers_mut(); + headers + .insert( + tonic::Status::GRPC_STATUS, + (tonic::Code::Unimplemented as i32).into(), + ); + headers + .insert( + http::header::CONTENT_TYPE, + tonic::metadata::GRPC_CONTENT_TYPE, + ); + Ok(response) + }) + } + } + } + } + impl Clone for TaskInstanceManagementServiceServer { + fn clone(&self) -> Self { + let inner = self.inner.clone(); + Self { + inner, + accept_compression_encodings: self.accept_compression_encodings, + send_compression_encodings: self.send_compression_encodings, + max_decoding_message_size: self.max_decoding_message_size, + max_encoding_message_size: self.max_encoding_message_size, + } + } + } + /// Generated gRPC service name + pub const SERVICE_NAME: &str = "storage.TaskInstanceManagementService"; + impl tonic::server::NamedService for TaskInstanceManagementServiceServer { + const NAME: &'static str = SERVICE_NAME; + } +} +/// Generated server implementations. +pub mod resource_group_management_service_server { + #![allow( + unused_variables, + dead_code, + missing_docs, + clippy::wildcard_imports, + clippy::let_unit_value, + )] + use tonic::codegen::*; + /// Generated trait containing gRPC methods that should be implemented for use with ResourceGroupManagementServiceServer. + #[async_trait] + pub trait ResourceGroupManagementService: std::marker::Send + std::marker::Sync + 'static { + async fn add_resource_group( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; + async fn verify_resource_group( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; + async fn delete_resource_group( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; + } + #[derive(Debug)] + pub struct ResourceGroupManagementServiceServer { + inner: Arc, + accept_compression_encodings: EnabledCompressionEncodings, + send_compression_encodings: EnabledCompressionEncodings, + max_decoding_message_size: Option, + max_encoding_message_size: Option, + } + impl ResourceGroupManagementServiceServer { + pub fn new(inner: T) -> Self { + Self::from_arc(Arc::new(inner)) + } + pub fn from_arc(inner: Arc) -> Self { + Self { + inner, + accept_compression_encodings: Default::default(), + send_compression_encodings: Default::default(), + max_decoding_message_size: None, + max_encoding_message_size: None, + } + } + pub fn with_interceptor( + inner: T, + interceptor: F, + ) -> InterceptedService + where + F: tonic::service::Interceptor, + { + InterceptedService::new(Self::new(inner), interceptor) + } + /// Enable decompressing requests with the given encoding. + #[must_use] + pub fn accept_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.accept_compression_encodings.enable(encoding); + self + } + /// Compress responses with the given encoding, if the client supports it. + #[must_use] + pub fn send_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.send_compression_encodings.enable(encoding); + self + } + /// Limits the maximum size of a decoded message. + /// + /// Default: `4MB` + #[must_use] + pub fn max_decoding_message_size(mut self, limit: usize) -> Self { + self.max_decoding_message_size = Some(limit); + self + } + /// Limits the maximum size of an encoded message. + /// + /// Default: `usize::MAX` + #[must_use] + pub fn max_encoding_message_size(mut self, limit: usize) -> Self { + self.max_encoding_message_size = Some(limit); + self + } + } + impl tonic::codegen::Service> + for ResourceGroupManagementServiceServer + where + T: ResourceGroupManagementService, + B: Body + std::marker::Send + 'static, + B::Error: Into + std::marker::Send + 'static, + { + type Response = http::Response; + type Error = std::convert::Infallible; + type Future = BoxFuture; + fn poll_ready( + &mut self, + _cx: &mut Context<'_>, + ) -> Poll> { + Poll::Ready(Ok(())) + } + fn call(&mut self, req: http::Request) -> Self::Future { + match req.uri().path() { + "/storage.ResourceGroupManagementService/AddResourceGroup" => { + #[allow(non_camel_case_types)] + struct AddResourceGroupSvc( + pub Arc, + ); + impl< + T: ResourceGroupManagementService, + > tonic::server::UnaryService + for AddResourceGroupSvc { + type Response = super::ResourceGroupIdResponse; + 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 { + ::add_resource_group( + &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 = AddResourceGroupSvc(inner); + let codec = tonic::codec::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) + } + "/storage.ResourceGroupManagementService/VerifyResourceGroup" => { + #[allow(non_camel_case_types)] + struct VerifyResourceGroupSvc( + pub Arc, + ); + impl< + T: ResourceGroupManagementService, + > tonic::server::UnaryService + for VerifyResourceGroupSvc { + type Response = super::StorageOperationResponse; + 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 { + ::verify_resource_group( + &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 = VerifyResourceGroupSvc(inner); + let codec = tonic::codec::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) + } + "/storage.ResourceGroupManagementService/DeleteResourceGroup" => { + #[allow(non_camel_case_types)] + struct DeleteResourceGroupSvc( + pub Arc, + ); + impl< + T: ResourceGroupManagementService, + > tonic::server::UnaryService + for DeleteResourceGroupSvc { + type Response = super::StorageOperationResponse; + 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 { + ::delete_resource_group( + &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 = DeleteResourceGroupSvc(inner); + let codec = tonic::codec::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(empty_body()); + let headers = response.headers_mut(); + headers + .insert( + tonic::Status::GRPC_STATUS, + (tonic::Code::Unimplemented as i32).into(), + ); + headers + .insert( + http::header::CONTENT_TYPE, + tonic::metadata::GRPC_CONTENT_TYPE, + ); + Ok(response) + }) + } + } + } + } + impl Clone for ResourceGroupManagementServiceServer { + fn clone(&self) -> Self { + let inner = self.inner.clone(); + Self { + inner, + accept_compression_encodings: self.accept_compression_encodings, + send_compression_encodings: self.send_compression_encodings, + max_decoding_message_size: self.max_decoding_message_size, + max_encoding_message_size: self.max_encoding_message_size, + } + } + } + /// Generated gRPC service name + pub const SERVICE_NAME: &str = "storage.ResourceGroupManagementService"; + impl tonic::server::NamedService for ResourceGroupManagementServiceServer { + const NAME: &'static str = SERVICE_NAME; + } +} +/// Generated server implementations. +pub mod execution_manager_liveness_service_server { + #![allow( + unused_variables, + dead_code, + missing_docs, + clippy::wildcard_imports, + clippy::let_unit_value, + )] + use tonic::codegen::*; + /// Generated trait containing gRPC methods that should be implemented for use with ExecutionManagerLivenessServiceServer. + #[async_trait] + pub trait ExecutionManagerLivenessService: std::marker::Send + std::marker::Sync + 'static { + async fn register_execution_manager( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; + async fn update_execution_manager_heartbeat( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; + async fn is_execution_manager_alive( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; + async fn get_dead_execution_managers( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; + } + #[derive(Debug)] + pub struct ExecutionManagerLivenessServiceServer { + inner: Arc, + accept_compression_encodings: EnabledCompressionEncodings, + send_compression_encodings: EnabledCompressionEncodings, + max_decoding_message_size: Option, + max_encoding_message_size: Option, + } + impl ExecutionManagerLivenessServiceServer { + pub fn new(inner: T) -> Self { + Self::from_arc(Arc::new(inner)) + } + pub fn from_arc(inner: Arc) -> Self { + Self { inner, accept_compression_encodings: Default::default(), send_compression_encodings: Default::default(), @@ -422,9 +2687,9 @@ pub mod task_instance_management_service_server { } } impl tonic::codegen::Service> - for TaskInstanceManagementServiceServer + for ExecutionManagerLivenessServiceServer where - T: TaskInstanceManagementService, + T: ExecutionManagerLivenessService, B: Body + std::marker::Send + 'static, B::Error: Into + std::marker::Send + 'static, { @@ -439,27 +2704,31 @@ pub mod task_instance_management_service_server { } fn call(&mut self, req: http::Request) -> Self::Future { match req.uri().path() { - "/storage.TaskInstanceManagementService/RegisterTaskInstance" => { + "/storage.ExecutionManagerLivenessService/RegisterExecutionManager" => { #[allow(non_camel_case_types)] - struct RegisterTaskInstanceSvc( + struct RegisterExecutionManagerSvc< + T: ExecutionManagerLivenessService, + >( pub Arc, ); impl< - T: TaskInstanceManagementService, - > tonic::server::UnaryService - for RegisterTaskInstanceSvc { - type Response = super::RegisterTaskInstanceResponse; + T: ExecutionManagerLivenessService, + > tonic::server::UnaryService + for RegisterExecutionManagerSvc { + type Response = super::ExecutionManagerIdResponse; type Future = BoxFuture< tonic::Response, tonic::Status, >; fn call( &mut self, - request: tonic::Request, + request: tonic::Request< + super::RegisterExecutionManagerRequest, + >, ) -> Self::Future { let inner = Arc::clone(&self.0); let fut = async move { - ::register_task_instance( + ::register_execution_manager( &inner, request, ) @@ -474,7 +2743,7 @@ pub mod task_instance_management_service_server { let max_encoding_message_size = self.max_encoding_message_size; let inner = self.inner.clone(); let fut = async move { - let method = RegisterTaskInstanceSvc(inner); + let method = RegisterExecutionManagerSvc(inner); let codec = tonic::codec::ProstCodec::default(); let mut grpc = tonic::server::Grpc::new(codec) .apply_compression_config( @@ -490,15 +2759,17 @@ pub mod task_instance_management_service_server { }; Box::pin(fut) } - "/storage.TaskInstanceManagementService/ReportTaskSuccess" => { + "/storage.ExecutionManagerLivenessService/UpdateExecutionManagerHeartbeat" => { #[allow(non_camel_case_types)] - struct ReportTaskSuccessSvc( + struct UpdateExecutionManagerHeartbeatSvc< + T: ExecutionManagerLivenessService, + >( pub Arc, ); impl< - T: TaskInstanceManagementService, - > tonic::server::UnaryService - for ReportTaskSuccessSvc { + T: ExecutionManagerLivenessService, + > tonic::server::UnaryService + for UpdateExecutionManagerHeartbeatSvc { type Response = super::StorageOperationResponse; type Future = BoxFuture< tonic::Response, @@ -506,11 +2777,11 @@ pub mod task_instance_management_service_server { >; fn call( &mut self, - request: tonic::Request, + request: tonic::Request, ) -> Self::Future { let inner = Arc::clone(&self.0); let fut = async move { - ::report_task_success( + ::update_execution_manager_heartbeat( &inner, request, ) @@ -525,7 +2796,7 @@ pub mod task_instance_management_service_server { let max_encoding_message_size = self.max_encoding_message_size; let inner = self.inner.clone(); let fut = async move { - let method = ReportTaskSuccessSvc(inner); + let method = UpdateExecutionManagerHeartbeatSvc(inner); let codec = tonic::codec::ProstCodec::default(); let mut grpc = tonic::server::Grpc::new(codec) .apply_compression_config( @@ -541,27 +2812,29 @@ pub mod task_instance_management_service_server { }; Box::pin(fut) } - "/storage.TaskInstanceManagementService/ReportTaskFailure" => { + "/storage.ExecutionManagerLivenessService/IsExecutionManagerAlive" => { #[allow(non_camel_case_types)] - struct ReportTaskFailureSvc( + struct IsExecutionManagerAliveSvc< + T: ExecutionManagerLivenessService, + >( pub Arc, ); impl< - T: TaskInstanceManagementService, - > tonic::server::UnaryService - for ReportTaskFailureSvc { - type Response = super::StorageOperationResponse; + T: ExecutionManagerLivenessService, + > tonic::server::UnaryService + for IsExecutionManagerAliveSvc { + type Response = super::IsExecutionManagerAliveResponse; type Future = BoxFuture< tonic::Response, tonic::Status, >; fn call( &mut self, - request: tonic::Request, + request: tonic::Request, ) -> Self::Future { let inner = Arc::clone(&self.0); let fut = async move { - ::report_task_failure( + ::is_execution_manager_alive( &inner, request, ) @@ -576,7 +2849,62 @@ pub mod task_instance_management_service_server { let max_encoding_message_size = self.max_encoding_message_size; let inner = self.inner.clone(); let fut = async move { - let method = ReportTaskFailureSvc(inner); + let method = IsExecutionManagerAliveSvc(inner); + let codec = tonic::codec::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) + } + "/storage.ExecutionManagerLivenessService/GetDeadExecutionManagers" => { + #[allow(non_camel_case_types)] + struct GetDeadExecutionManagersSvc< + T: ExecutionManagerLivenessService, + >( + pub Arc, + ); + impl< + T: ExecutionManagerLivenessService, + > tonic::server::UnaryService + for GetDeadExecutionManagersSvc { + type Response = super::GetDeadExecutionManagersResponse; + type Future = BoxFuture< + tonic::Response, + tonic::Status, + >; + fn call( + &mut self, + request: tonic::Request< + super::GetDeadExecutionManagersRequest, + >, + ) -> Self::Future { + let inner = Arc::clone(&self.0); + let fut = async move { + ::get_dead_execution_managers( + &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 = GetDeadExecutionManagersSvc(inner); let codec = tonic::codec::ProstCodec::default(); let mut grpc = tonic::server::Grpc::new(codec) .apply_compression_config( @@ -612,7 +2940,7 @@ pub mod task_instance_management_service_server { } } } - impl Clone for TaskInstanceManagementServiceServer { + impl Clone for ExecutionManagerLivenessServiceServer { fn clone(&self) -> Self { let inner = self.inner.clone(); Self { @@ -625,8 +2953,192 @@ pub mod task_instance_management_service_server { } } /// Generated gRPC service name - pub const SERVICE_NAME: &str = "storage.TaskInstanceManagementService"; - impl tonic::server::NamedService for TaskInstanceManagementServiceServer { + pub const SERVICE_NAME: &str = "storage.ExecutionManagerLivenessService"; + impl tonic::server::NamedService for ExecutionManagerLivenessServiceServer { + const NAME: &'static str = SERVICE_NAME; + } +} +/// Generated server implementations. +pub mod session_management_service_server { + #![allow( + unused_variables, + dead_code, + missing_docs, + clippy::wildcard_imports, + clippy::let_unit_value, + )] + use tonic::codegen::*; + /// Generated trait containing gRPC methods that should be implemented for use with SessionManagementServiceServer. + #[async_trait] + pub trait SessionManagementService: std::marker::Send + std::marker::Sync + 'static { + async fn get_session( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; + } + #[derive(Debug)] + pub struct SessionManagementServiceServer { + inner: Arc, + accept_compression_encodings: EnabledCompressionEncodings, + send_compression_encodings: EnabledCompressionEncodings, + max_decoding_message_size: Option, + max_encoding_message_size: Option, + } + impl SessionManagementServiceServer { + pub fn new(inner: T) -> Self { + Self::from_arc(Arc::new(inner)) + } + pub fn from_arc(inner: Arc) -> Self { + Self { + inner, + accept_compression_encodings: Default::default(), + send_compression_encodings: Default::default(), + max_decoding_message_size: None, + max_encoding_message_size: None, + } + } + pub fn with_interceptor( + inner: T, + interceptor: F, + ) -> InterceptedService + where + F: tonic::service::Interceptor, + { + InterceptedService::new(Self::new(inner), interceptor) + } + /// Enable decompressing requests with the given encoding. + #[must_use] + pub fn accept_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.accept_compression_encodings.enable(encoding); + self + } + /// Compress responses with the given encoding, if the client supports it. + #[must_use] + pub fn send_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.send_compression_encodings.enable(encoding); + self + } + /// Limits the maximum size of a decoded message. + /// + /// Default: `4MB` + #[must_use] + pub fn max_decoding_message_size(mut self, limit: usize) -> Self { + self.max_decoding_message_size = Some(limit); + self + } + /// Limits the maximum size of an encoded message. + /// + /// Default: `usize::MAX` + #[must_use] + pub fn max_encoding_message_size(mut self, limit: usize) -> Self { + self.max_encoding_message_size = Some(limit); + self + } + } + impl tonic::codegen::Service> + for SessionManagementServiceServer + where + T: SessionManagementService, + B: Body + std::marker::Send + 'static, + B::Error: Into + std::marker::Send + 'static, + { + type Response = http::Response; + type Error = std::convert::Infallible; + type Future = BoxFuture; + fn poll_ready( + &mut self, + _cx: &mut Context<'_>, + ) -> Poll> { + Poll::Ready(Ok(())) + } + fn call(&mut self, req: http::Request) -> Self::Future { + match req.uri().path() { + "/storage.SessionManagementService/GetSession" => { + #[allow(non_camel_case_types)] + struct GetSessionSvc(pub Arc); + impl< + T: SessionManagementService, + > tonic::server::UnaryService for GetSessionSvc { + type Response = super::GetSessionResponse; + 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 { + ::get_session( + &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 = GetSessionSvc(inner); + let codec = tonic::codec::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(empty_body()); + let headers = response.headers_mut(); + headers + .insert( + tonic::Status::GRPC_STATUS, + (tonic::Code::Unimplemented as i32).into(), + ); + headers + .insert( + http::header::CONTENT_TYPE, + tonic::metadata::GRPC_CONTENT_TYPE, + ); + Ok(response) + }) + } + } + } + } + impl Clone for SessionManagementServiceServer { + fn clone(&self) -> Self { + let inner = self.inner.clone(); + Self { + inner, + accept_compression_encodings: self.accept_compression_encodings, + send_compression_encodings: self.send_compression_encodings, + max_decoding_message_size: self.max_decoding_message_size, + max_encoding_message_size: self.max_encoding_message_size, + } + } + } + /// Generated gRPC service name + pub const SERVICE_NAME: &str = "storage.SessionManagementService"; + impl tonic::server::NamedService for SessionManagementServiceServer { const NAME: &'static str = SERVICE_NAME; } } From 80768cc97be4e8411378e84971e9bbdc34822ff0 Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Wed, 10 Jun 2026 12:02:40 -0400 Subject: [PATCH 03/16] Add conversion helper --- components/spider-proto-rust/src/id.rs | 33 +++++++++++++ components/spider-proto-rust/src/job.rs | 66 +++++++++++++++++++++++++ components/spider-proto-rust/src/lib.rs | 1 + 3 files changed, 100 insertions(+) create mode 100644 components/spider-proto-rust/src/job.rs diff --git a/components/spider-proto-rust/src/id.rs b/components/spider-proto-rust/src/id.rs index ef21bcd8..1da62ca5 100644 --- a/components/spider-proto-rust/src/id.rs +++ b/components/spider-proto-rust/src/id.rs @@ -17,6 +17,21 @@ impl From for storage::TaskId { } } +impl TryFrom for TaskId { + type Error = String; + + fn try_from(task_id: storage::TaskId) -> Result { + match task_id.kind { + Some(task_id::Kind::Index(task_index)) => usize::try_from(task_index) + .map(TaskId::Index) + .map_err(|error| format!("task index does not fit in usize: {error}")), + Some(task_id::Kind::Commit(_)) => Ok(Self::Commit), + Some(task_id::Kind::Cleanup(_)) => Ok(Self::Cleanup), + None => Err("task id missing kind".to_owned()), + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -41,4 +56,22 @@ mod tests { assert!(matches!(task_id.kind, Some(task_id::Kind::Cleanup(_)))); } + + #[test] + fn protocol_task_id_to_core_converts_index_task() { + let task_id = TaskId::try_from(storage::TaskId { + kind: Some(task_id::Kind::Index(7)), + }) + .expect("protocol task id conversion should succeed"); + + assert_eq!(task_id, TaskId::Index(7)); + } + + #[test] + fn protocol_task_id_to_core_rejects_missing_kind() { + let error = TaskId::try_from(storage::TaskId { kind: None }) + .expect_err("missing task id kind should fail"); + + assert!(error.contains("missing kind")); + } } diff --git a/components/spider-proto-rust/src/job.rs b/components/spider-proto-rust/src/job.rs new file mode 100644 index 00000000..f8d62494 --- /dev/null +++ b/components/spider-proto-rust/src/job.rs @@ -0,0 +1,66 @@ +//! Helpers for converting Spider job values to protobuf fields. + +use spider_core::job::JobState; + +use crate::storage; + +impl From for storage::JobState { + fn from(state: JobState) -> Self { + match state { + JobState::Ready => Self::Ready, + JobState::Running => Self::Running, + JobState::CommitReady => Self::CommitReady, + JobState::CleanupReady => Self::CleanupReady, + JobState::Succeeded => Self::Succeeded, + JobState::Failed => Self::Failed, + JobState::Cancelled => Self::Cancelled, + } + } +} + +impl TryFrom for JobState { + type Error = String; + + fn try_from(state: storage::JobState) -> Result { + match state { + storage::JobState::Unspecified => Err("job state is unspecified".to_owned()), + storage::JobState::Ready => Ok(Self::Ready), + storage::JobState::Running => Ok(Self::Running), + storage::JobState::CommitReady => Ok(Self::CommitReady), + storage::JobState::CleanupReady => Ok(Self::CleanupReady), + storage::JobState::Succeeded => Ok(Self::Succeeded), + storage::JobState::Failed => Ok(Self::Failed), + storage::JobState::Cancelled => Ok(Self::Cancelled), + } + } +} + +#[cfg(test)] +mod tests { + use spider_core::job::JobState; + + use crate::storage::JobState as ProtocolJobState; + + #[test] + fn job_state_to_protocol_converts_succeeded() { + let protocol_state = ProtocolJobState::from(JobState::Succeeded); + + assert_eq!(protocol_state, ProtocolJobState::Succeeded); + } + + #[test] + fn protocol_job_state_to_core_converts_cleanup_ready() { + let state = JobState::try_from(ProtocolJobState::CleanupReady) + .expect("protocol job state conversion should succeed"); + + assert_eq!(state, JobState::CleanupReady); + } + + #[test] + fn protocol_job_state_to_core_rejects_unspecified() { + let error = JobState::try_from(ProtocolJobState::Unspecified) + .expect_err("unspecified job state should fail"); + + assert!(error.contains("unspecified")); + } +} diff --git a/components/spider-proto-rust/src/lib.rs b/components/spider-proto-rust/src/lib.rs index d78e8f0d..06f1c0ba 100644 --- a/components/spider-proto-rust/src/lib.rs +++ b/components/spider-proto-rust/src/lib.rs @@ -1,6 +1,7 @@ //! Rust gRPC protocol definitions generated from Spider protobuf files. pub mod id; +pub mod job; #[allow(clippy::all, clippy::nursery, clippy::pedantic)] pub mod storage { From 46d9bac6b6db68d03d3f39f4070313f31058a5ee Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Wed, 10 Jun 2026 12:13:24 -0400 Subject: [PATCH 04/16] Embed session id in em register response --- .../src/generated/storage.rs | 51 ++++++++++++++----- components/spider-proto/storage/storage.proto | 20 ++++++-- 2 files changed, 52 insertions(+), 19 deletions(-) diff --git a/components/spider-proto-rust/src/generated/storage.rs b/components/spider-proto-rust/src/generated/storage.rs index 8b6b2c00..d228c04c 100644 --- a/components/spider-proto-rust/src/generated/storage.rs +++ b/components/spider-proto-rust/src/generated/storage.rs @@ -212,27 +212,50 @@ pub struct VerifyResourceGroupRequest { pub struct RegisterExecutionManagerRequest { #[prost(string, tag = "1")] pub ip_address: ::prost::alloc::string::String, - #[prost(uint64, tag = "2")] - pub session_id: u64, } #[derive(Clone, Copy, PartialEq, ::prost::Message)] pub struct ExecutionManagerIdRequest { #[prost(uint64, tag = "1")] pub execution_manager_id: u64, +} +#[derive(Clone, Copy, PartialEq, ::prost::Message)] +pub struct ExecutionManagerRegistration { + #[prost(uint64, tag = "1")] + pub execution_manager_id: u64, #[prost(uint64, tag = "2")] pub session_id: u64, } #[derive(Clone, PartialEq, ::prost::Message)] -pub struct ExecutionManagerIdResponse { - #[prost(oneof = "execution_manager_id_response::Result", tags = "1, 2")] - pub result: ::core::option::Option, +pub struct RegisterExecutionManagerResponse { + #[prost(oneof = "register_execution_manager_response::Result", tags = "1, 2")] + pub result: ::core::option::Option, } -/// Nested message and enum types in `ExecutionManagerIdResponse`. -pub mod execution_manager_id_response { +/// 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::StorageError), + } +} +#[derive(Clone, 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")] - ExecutionManagerId(u64), + SessionId(u64), #[prost(message, tag = "2")] Error(super::StorageError), } @@ -1177,7 +1200,7 @@ pub mod execution_manager_liveness_service_client { &mut self, request: impl tonic::IntoRequest, ) -> std::result::Result< - tonic::Response, + tonic::Response, tonic::Status, > { self.inner @@ -1206,7 +1229,7 @@ pub mod execution_manager_liveness_service_client { &mut self, request: impl tonic::IntoRequest, ) -> std::result::Result< - tonic::Response, + tonic::Response, tonic::Status, > { self.inner @@ -2602,14 +2625,14 @@ pub mod execution_manager_liveness_service_server { &self, request: tonic::Request, ) -> std::result::Result< - tonic::Response, + tonic::Response, tonic::Status, >; async fn update_execution_manager_heartbeat( &self, request: tonic::Request, ) -> std::result::Result< - tonic::Response, + tonic::Response, tonic::Status, >; async fn is_execution_manager_alive( @@ -2715,7 +2738,7 @@ pub mod execution_manager_liveness_service_server { T: ExecutionManagerLivenessService, > tonic::server::UnaryService for RegisterExecutionManagerSvc { - type Response = super::ExecutionManagerIdResponse; + type Response = super::RegisterExecutionManagerResponse; type Future = BoxFuture< tonic::Response, tonic::Status, @@ -2770,7 +2793,7 @@ pub mod execution_manager_liveness_service_server { T: ExecutionManagerLivenessService, > tonic::server::UnaryService for UpdateExecutionManagerHeartbeatSvc { - type Response = super::StorageOperationResponse; + type Response = super::UpdateExecutionManagerHeartbeatResponse; type Future = BoxFuture< tonic::Response, tonic::Status, diff --git a/components/spider-proto/storage/storage.proto b/components/spider-proto/storage/storage.proto index d7802d82..95c06d3a 100644 --- a/components/spider-proto/storage/storage.proto +++ b/components/spider-proto/storage/storage.proto @@ -28,9 +28,9 @@ service ResourceGroupManagementService { service ExecutionManagerLivenessService { rpc RegisterExecutionManager(RegisterExecutionManagerRequest) - returns (ExecutionManagerIdResponse); + returns (RegisterExecutionManagerResponse); rpc UpdateExecutionManagerHeartbeat(ExecutionManagerIdRequest) - returns (StorageOperationResponse); + returns (UpdateExecutionManagerHeartbeatResponse); rpc IsExecutionManagerAlive(ExecutionManagerIdRequest) returns (IsExecutionManagerAliveResponse); rpc GetDeadExecutionManagers(GetDeadExecutionManagersRequest) @@ -162,17 +162,27 @@ message VerifyResourceGroupRequest { message RegisterExecutionManagerRequest { string ip_address = 1; - uint64 session_id = 2; } message ExecutionManagerIdRequest { uint64 execution_manager_id = 1; +} + +message ExecutionManagerRegistration { + uint64 execution_manager_id = 1; uint64 session_id = 2; } -message ExecutionManagerIdResponse { +message RegisterExecutionManagerResponse { + oneof result { + ExecutionManagerRegistration registration = 1; + StorageError error = 2; + } +} + +message UpdateExecutionManagerHeartbeatResponse { oneof result { - uint64 execution_manager_id = 1; + uint64 session_id = 1; StorageError error = 2; } } From 8da124dd0005011d73ed03d694972e6ee7b5d7e8 Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Wed, 10 Jun 2026 12:14:52 -0400 Subject: [PATCH 05/16] Add grpc liveness client --- .../spider-execution-manager/src/client.rs | 2 +- .../src/client/grpc/liveness.rs | 210 ++++++++++++++++++ .../src/client/grpc/mod.rs | 2 + 3 files changed, 213 insertions(+), 1 deletion(-) create mode 100644 components/spider-execution-manager/src/client/grpc/liveness.rs diff --git a/components/spider-execution-manager/src/client.rs b/components/spider-execution-manager/src/client.rs index 63b132ce..c2ac95e4 100644 --- a/components/spider-execution-manager/src/client.rs +++ b/components/spider-execution-manager/src/client.rs @@ -11,7 +11,7 @@ pub mod liveness; pub mod scheduler; pub mod storage; -pub use grpc::GrpcStorageClient; +pub use grpc::{GrpcLivenessClient, GrpcStorageClient}; pub use liveness::{LivenessClient, LivenessResponseError, RegistrationResponse}; pub use scheduler::{SchedulerClient, SchedulerError, SchedulerResponse}; pub use storage::{StorageClient, StorageResponseError}; diff --git a/components/spider-execution-manager/src/client/grpc/liveness.rs b/components/spider-execution-manager/src/client/grpc/liveness.rs new file mode 100644 index 00000000..3cfe32ea --- /dev/null +++ b/components/spider-execution-manager/src/client/grpc/liveness.rs @@ -0,0 +1,210 @@ +//! gRPC-backed [`LivenessClient`] implementation. + +use std::net::IpAddr; + +use async_trait::async_trait; +use spider_core::types::id::{ExecutionManagerId, SessionId}; +use spider_proto_rust::storage::{ + self, + execution_manager_liveness_service_client::ExecutionManagerLivenessServiceClient, + register_execution_manager_response, + storage_error, + update_execution_manager_heartbeat_response, +}; +use tonic::transport::{Channel, Endpoint}; + +use crate::client::liveness::{LivenessClient, LivenessResponseError, RegistrationResponse}; + +/// gRPC-backed [`LivenessClient`] implementation. +#[derive(Debug, Clone)] +pub struct GrpcLivenessClient { + client: ExecutionManagerLivenessServiceClient, +} + +impl GrpcLivenessClient { + /// Connects to the storage gRPC endpoint. + /// + /// # Returns + /// + /// A new [`GrpcLivenessClient`] connected to `endpoint` on success. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`LivenessResponseError::Transport`] if tonic cannot create or connect to the endpoint. + pub async fn connect(endpoint: Endpoint) -> Result { + ExecutionManagerLivenessServiceClient::connect(endpoint) + .await + .map(|inner| Self { client: inner }) + .map_err(to_transport_error) + } +} + +#[async_trait] +impl LivenessClient for GrpcLivenessClient { + async fn register(&self, ip: IpAddr) -> Result { + let request = storage::RegisterExecutionManagerRequest { + ip_address: ip.to_string(), + }; + let response = self + .client + .clone() + .register_execution_manager(request) + .await + .map_err(to_transport_error)? + .into_inner(); + + register_response_to_result(response) + } + + async fn heartbeat( + &self, + em_id: ExecutionManagerId, + ) -> Result { + let request = storage::ExecutionManagerIdRequest { + execution_manager_id: em_id.get(), + }; + let response = self + .client + .clone() + .update_execution_manager_heartbeat(request) + .await + .map_err(to_transport_error)? + .into_inner(); + + heartbeat_response_to_result(response) + } +} + +/// # Returns +/// +/// [`storage::RegisterExecutionManagerResponse`] converted into +/// [`Result`]. +fn register_response_to_result( + response: storage::RegisterExecutionManagerResponse, +) -> Result { + match response.result { + Some(register_execution_manager_response::Result::Registration(registration)) => { + Ok(RegistrationResponse { + em_id: ExecutionManagerId::from(registration.execution_manager_id), + session_id: registration.session_id, + }) + } + Some(register_execution_manager_response::Result::Error(error)) => { + Err(storage_error_to_liveness_error(error)) + } + None => Err(LivenessResponseError::Transport( + "register execution manager response missing result".to_owned(), + )), + } +} + +/// # Returns +/// +/// [`storage::UpdateExecutionManagerHeartbeatResponse`] converted into +/// [`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(storage_error_to_liveness_error(error)) + } + None => Err(LivenessResponseError::Transport( + "update execution manager heartbeat response missing result".to_owned(), + )), + } +} + +/// Converts a protobuf storage error into a liveness client error. +/// +/// # Returns +/// +/// The corresponding [`LivenessResponseError`]. +fn storage_error_to_liveness_error(error: storage::StorageError) -> LivenessResponseError { + match storage_error::ErrCode::try_from(error.err_code) { + Ok(storage_error::ErrCode::CacheStale) => LivenessResponseError::MarkedDead, + Ok(storage_error::ErrCode::InvalidInput) => LivenessResponseError::IllegalId(error.message), + Ok( + storage_error::ErrCode::Transport + | storage_error::ErrCode::Server + | storage_error::ErrCode::StaleSession + | storage_error::ErrCode::Unspecified, + ) => LivenessResponseError::Transport(error.message), + Err(error) => { + LivenessResponseError::Transport(format!("unknown storage error kind: {error}")) + } + } +} + +/// Converts a displayable transport-layer error into [`LivenessResponseError::Transport`]. +/// +/// # Returns +/// +/// A [`LivenessResponseError::Transport`] containing `error`'s display string. +fn to_transport_error(error: impl std::fmt::Display) -> LivenessResponseError { + LivenessResponseError::Transport(error.to_string()) +} + +#[cfg(test)] +mod tests { + use spider_core::types::id::ExecutionManagerId; + + use super::*; + use crate::client::{LivenessResponseError, RegistrationResponse}; + + #[test] + fn register_response_to_result_returns_registration() { + let response = storage::RegisterExecutionManagerResponse { + result: Some(register_execution_manager_response::Result::Registration( + storage::ExecutionManagerRegistration { + execution_manager_id: 5, + session_id: 7, + }, + )), + }; + + let registration = register_response_to_result(response) + .expect("registration response conversion should succeed"); + + assert_eq!( + registration, + RegistrationResponse { + em_id: ExecutionManagerId::from(5), + session_id: 7, + } + ); + } + + #[test] + fn heartbeat_response_to_result_returns_session_id() { + let response = storage::UpdateExecutionManagerHeartbeatResponse { + result: Some(update_execution_manager_heartbeat_response::Result::SessionId(9)), + }; + + let session_id = heartbeat_response_to_result(response) + .expect("heartbeat response conversion should succeed"); + + assert_eq!(session_id, 9); + } + + #[test] + fn liveness_storage_error_maps_invalid_input_to_illegal_id() { + let error = storage::StorageError { + err_code: storage_error::ErrCode::InvalidInput.into(), + message: "bad em id".to_owned(), + storage_session: 0, + }; + + match storage_error_to_liveness_error(error) { + LivenessResponseError::IllegalId(message) => { + assert_eq!(message, "bad em id"); + } + error => panic!("unexpected liveness response error: {error:?}"), + } + } +} diff --git a/components/spider-execution-manager/src/client/grpc/mod.rs b/components/spider-execution-manager/src/client/grpc/mod.rs index 9f15ee9a..93a0a50e 100644 --- a/components/spider-execution-manager/src/client/grpc/mod.rs +++ b/components/spider-execution-manager/src/client/grpc/mod.rs @@ -1,5 +1,7 @@ //! gRPC-backed implementations of the execution manager's client traits. +pub mod liveness; pub mod storage; +pub use liveness::GrpcLivenessClient; pub use storage::GrpcStorageClient; From 0d06cf81da0888e27ee4ffc30bff7fdbc9e2499b Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Thu, 11 Jun 2026 13:11:09 -0400 Subject: [PATCH 06/16] Add scheduler support in storage proto --- .../src/generated/storage.rs | 524 ++++++++++++++++++ components/spider-proto/storage/storage.proto | 31 ++ 2 files changed, 555 insertions(+) diff --git a/components/spider-proto-rust/src/generated/storage.rs b/components/spider-proto-rust/src/generated/storage.rs index d228c04c..e52bfa9e 100644 --- a/components/spider-proto-rust/src/generated/storage.rs +++ b/components/spider-proto-rust/src/generated/storage.rs @@ -117,6 +117,44 @@ pub struct ResendReadyTasksRequest { pub session_id: u64, } #[derive(Clone, Copy, PartialEq, ::prost::Message)] +pub struct PollReadyTasksRequest { + #[prost(uint64, tag = "1")] + pub max_items: u64, + #[prost(uint64, tag = "2")] + pub wait_ns: u64, +} +#[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::StorageError), + } +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ReadyTasks { + #[prost(uint64, tag = "1")] + pub session_id: u64, + #[prost(message, repeated, tag = "2")] + pub tasks: ::prost::alloc::vec::Vec, +} +#[derive(Clone, Copy, PartialEq, ::prost::Message)] +pub struct ReadyTask { + #[prost(uint64, tag = "1")] + pub resource_group_id: u64, + #[prost(uint64, tag = "2")] + pub job_id: u64, + #[prost(message, optional, tag = "3")] + pub task_id: ::core::option::Option, +} +#[derive(Clone, Copy, PartialEq, ::prost::Message)] pub struct RegisterTaskInstanceRequest { #[prost(uint64, tag = "1")] pub job_id: u64, @@ -371,6 +409,8 @@ pub mod storage_error { Transport = 3, Server = 4, InvalidInput = 5, + InboundClosed = 6, + JobNotFound = 7, } impl ErrCode { /// String value of the enum field names used in the ProtoBuf definition. @@ -385,6 +425,8 @@ pub mod storage_error { Self::Transport => "TRANSPORT", Self::Server => "SERVER", Self::InvalidInput => "INVALID_INPUT", + Self::InboundClosed => "INBOUND_CLOSED", + Self::JobNotFound => "JOB_NOT_FOUND", } } /// Creates an enum from field names used in the ProtoBuf definition. @@ -396,6 +438,8 @@ pub mod storage_error { "TRANSPORT" => Some(Self::Transport), "SERVER" => Some(Self::Server), "INVALID_INPUT" => Some(Self::InvalidInput), + "INBOUND_CLOSED" => Some(Self::InboundClosed), + "JOB_NOT_FOUND" => Some(Self::JobNotFound), _ => None, } } @@ -922,6 +966,185 @@ pub mod task_instance_management_service_client { } } /// Generated client implementations. +pub mod scheduler_storage_service_client { + #![allow( + unused_variables, + dead_code, + missing_docs, + clippy::wildcard_imports, + clippy::let_unit_value, + )] + use tonic::codegen::*; + use tonic::codegen::http::Uri; + #[derive(Debug, Clone)] + pub struct SchedulerStorageServiceClient { + inner: tonic::client::Grpc, + } + impl SchedulerStorageServiceClient { + /// Attempt to create a new client by connecting to a given endpoint. + pub async fn connect(dst: D) -> Result + where + D: TryInto, + D::Error: Into, + { + let conn = tonic::transport::Endpoint::new(dst)?.connect().await?; + Ok(Self::new(conn)) + } + } + impl SchedulerStorageServiceClient + where + T: tonic::client::GrpcService, + T::Error: Into, + T::ResponseBody: Body + std::marker::Send + 'static, + ::Error: Into + std::marker::Send, + { + pub fn new(inner: T) -> Self { + let inner = tonic::client::Grpc::new(inner); + Self { inner } + } + pub fn with_origin(inner: T, origin: Uri) -> Self { + let inner = tonic::client::Grpc::with_origin(inner, origin); + Self { inner } + } + pub fn with_interceptor( + inner: T, + interceptor: F, + ) -> SchedulerStorageServiceClient> + where + F: tonic::service::Interceptor, + T::ResponseBody: Default, + T: tonic::codegen::Service< + http::Request, + Response = http::Response< + >::ResponseBody, + >, + >, + , + >>::Error: Into + std::marker::Send + std::marker::Sync, + { + SchedulerStorageServiceClient::new( + InterceptedService::new(inner, interceptor), + ) + } + /// Compress requests with the given encoding. + /// + /// This requires the server to support it otherwise it might respond with an + /// error. + #[must_use] + pub fn send_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.inner = self.inner.send_compressed(encoding); + self + } + /// Enable decompressing responses. + #[must_use] + pub fn accept_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.inner = self.inner.accept_compressed(encoding); + self + } + /// Limits the maximum size of a decoded message. + /// + /// Default: `4MB` + #[must_use] + pub fn max_decoding_message_size(mut self, limit: usize) -> Self { + self.inner = self.inner.max_decoding_message_size(limit); + self + } + /// Limits the maximum size of an encoded message. + /// + /// Default: `usize::MAX` + #[must_use] + pub fn max_encoding_message_size(mut self, limit: usize) -> Self { + self.inner = self.inner.max_encoding_message_size(limit); + self + } + pub async fn poll_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::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/storage.SchedulerStorageService/PollReadyTasks", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert( + GrpcMethod::new("storage.SchedulerStorageService", "PollReadyTasks"), + ); + self.inner.unary(req, path, codec).await + } + pub async fn poll_ready_commit_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::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/storage.SchedulerStorageService/PollReadyCommitTasks", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert( + GrpcMethod::new( + "storage.SchedulerStorageService", + "PollReadyCommitTasks", + ), + ); + self.inner.unary(req, path, codec).await + } + pub async fn poll_ready_cleanup_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::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/storage.SchedulerStorageService/PollReadyCleanupTasks", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert( + GrpcMethod::new( + "storage.SchedulerStorageService", + "PollReadyCleanupTasks", + ), + ); + self.inner.unary(req, path, codec).await + } + } +} +/// Generated client implementations. pub mod resource_group_management_service_client { #![allow( unused_variables, @@ -2306,6 +2529,307 @@ pub mod task_instance_management_service_server { } } /// Generated server implementations. +pub mod scheduler_storage_service_server { + #![allow( + unused_variables, + dead_code, + missing_docs, + clippy::wildcard_imports, + clippy::let_unit_value, + )] + use tonic::codegen::*; + /// Generated trait containing gRPC methods that should be implemented for use with SchedulerStorageServiceServer. + #[async_trait] + pub trait SchedulerStorageService: std::marker::Send + std::marker::Sync + 'static { + async fn poll_ready_tasks( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; + async fn poll_ready_commit_tasks( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; + async fn poll_ready_cleanup_tasks( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; + } + #[derive(Debug)] + pub struct SchedulerStorageServiceServer { + inner: Arc, + accept_compression_encodings: EnabledCompressionEncodings, + send_compression_encodings: EnabledCompressionEncodings, + max_decoding_message_size: Option, + max_encoding_message_size: Option, + } + impl SchedulerStorageServiceServer { + pub fn new(inner: T) -> Self { + Self::from_arc(Arc::new(inner)) + } + pub fn from_arc(inner: Arc) -> Self { + Self { + inner, + accept_compression_encodings: Default::default(), + send_compression_encodings: Default::default(), + max_decoding_message_size: None, + max_encoding_message_size: None, + } + } + pub fn with_interceptor( + inner: T, + interceptor: F, + ) -> InterceptedService + where + F: tonic::service::Interceptor, + { + InterceptedService::new(Self::new(inner), interceptor) + } + /// Enable decompressing requests with the given encoding. + #[must_use] + pub fn accept_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.accept_compression_encodings.enable(encoding); + self + } + /// Compress responses with the given encoding, if the client supports it. + #[must_use] + pub fn send_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.send_compression_encodings.enable(encoding); + self + } + /// Limits the maximum size of a decoded message. + /// + /// Default: `4MB` + #[must_use] + pub fn max_decoding_message_size(mut self, limit: usize) -> Self { + self.max_decoding_message_size = Some(limit); + self + } + /// Limits the maximum size of an encoded message. + /// + /// Default: `usize::MAX` + #[must_use] + pub fn max_encoding_message_size(mut self, limit: usize) -> Self { + self.max_encoding_message_size = Some(limit); + self + } + } + impl tonic::codegen::Service> + for SchedulerStorageServiceServer + where + T: SchedulerStorageService, + B: Body + std::marker::Send + 'static, + B::Error: Into + std::marker::Send + 'static, + { + type Response = http::Response; + type Error = std::convert::Infallible; + type Future = BoxFuture; + fn poll_ready( + &mut self, + _cx: &mut Context<'_>, + ) -> Poll> { + Poll::Ready(Ok(())) + } + fn call(&mut self, req: http::Request) -> Self::Future { + match req.uri().path() { + "/storage.SchedulerStorageService/PollReadyTasks" => { + #[allow(non_camel_case_types)] + struct PollReadyTasksSvc(pub Arc); + impl< + T: SchedulerStorageService, + > tonic::server::UnaryService + for PollReadyTasksSvc { + type Response = super::PollReadyTasksResponse; + 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 { + ::poll_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 = PollReadyTasksSvc(inner); + let codec = tonic::codec::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) + } + "/storage.SchedulerStorageService/PollReadyCommitTasks" => { + #[allow(non_camel_case_types)] + struct PollReadyCommitTasksSvc( + pub Arc, + ); + impl< + T: SchedulerStorageService, + > tonic::server::UnaryService + for PollReadyCommitTasksSvc { + type Response = super::PollReadyTasksResponse; + 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 { + ::poll_ready_commit_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 = PollReadyCommitTasksSvc(inner); + let codec = tonic::codec::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) + } + "/storage.SchedulerStorageService/PollReadyCleanupTasks" => { + #[allow(non_camel_case_types)] + struct PollReadyCleanupTasksSvc( + pub Arc, + ); + impl< + T: SchedulerStorageService, + > tonic::server::UnaryService + for PollReadyCleanupTasksSvc { + type Response = super::PollReadyTasksResponse; + 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 { + ::poll_ready_cleanup_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 = PollReadyCleanupTasksSvc(inner); + let codec = tonic::codec::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(empty_body()); + let headers = response.headers_mut(); + headers + .insert( + tonic::Status::GRPC_STATUS, + (tonic::Code::Unimplemented as i32).into(), + ); + headers + .insert( + http::header::CONTENT_TYPE, + tonic::metadata::GRPC_CONTENT_TYPE, + ); + Ok(response) + }) + } + } + } + } + impl Clone for SchedulerStorageServiceServer { + fn clone(&self) -> Self { + let inner = self.inner.clone(); + Self { + inner, + accept_compression_encodings: self.accept_compression_encodings, + send_compression_encodings: self.send_compression_encodings, + max_decoding_message_size: self.max_decoding_message_size, + max_encoding_message_size: self.max_encoding_message_size, + } + } + } + /// Generated gRPC service name + pub const SERVICE_NAME: &str = "storage.SchedulerStorageService"; + impl tonic::server::NamedService for SchedulerStorageServiceServer { + const NAME: &'static str = SERVICE_NAME; + } +} +/// Generated server implementations. pub mod resource_group_management_service_server { #![allow( unused_variables, diff --git a/components/spider-proto/storage/storage.proto b/components/spider-proto/storage/storage.proto index 95c06d3a..256ca4fc 100644 --- a/components/spider-proto/storage/storage.proto +++ b/components/spider-proto/storage/storage.proto @@ -20,6 +20,12 @@ service TaskInstanceManagementService { rpc ReportTaskFailure(ReportTaskFailureRequest) returns (StorageOperationResponse); } +service SchedulerStorageService { + rpc PollReadyTasks(PollReadyTasksRequest) returns (PollReadyTasksResponse); + rpc PollReadyCommitTasks(PollReadyTasksRequest) returns (PollReadyTasksResponse); + rpc PollReadyCleanupTasks(PollReadyTasksRequest) returns (PollReadyTasksResponse); +} + service ResourceGroupManagementService { rpc AddResourceGroup(AddResourceGroupRequest) returns (ResourceGroupIdResponse); rpc VerifyResourceGroup(VerifyResourceGroupRequest) returns (StorageOperationResponse); @@ -106,6 +112,29 @@ message ResendReadyTasksRequest { uint64 session_id = 2; } +message PollReadyTasksRequest { + uint64 max_items = 1; + uint64 wait_ns = 2; +} + +message PollReadyTasksResponse { + oneof result { + ReadyTasks tasks = 1; + StorageError error = 2; + } +} + +message ReadyTasks { + uint64 session_id = 1; + repeated ReadyTask tasks = 2; +} + +message ReadyTask { + uint64 resource_group_id = 1; + uint64 job_id = 2; + TaskId task_id = 3; +} + message RegisterTaskInstanceRequest { uint64 job_id = 1; TaskId task_id = 2; @@ -250,6 +279,8 @@ message StorageError { TRANSPORT = 3; SERVER = 4; INVALID_INPUT = 5; + INBOUND_CLOSED = 6; + JOB_NOT_FOUND = 7; } ErrCode err_code = 1; From 187915dc8078b74abc50c9540544be88f9209c6c Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Thu, 11 Jun 2026 13:31:51 -0400 Subject: [PATCH 07/16] Split error type in proto --- .../src/generated/storage.rs | 330 +++++++++++++++--- components/spider-proto/storage/storage.proto | 112 ++++-- 2 files changed, 370 insertions(+), 72 deletions(-) diff --git a/components/spider-proto-rust/src/generated/storage.rs b/components/spider-proto-rust/src/generated/storage.rs index e52bfa9e..b71d2aca 100644 --- a/components/spider-proto-rust/src/generated/storage.rs +++ b/components/spider-proto-rust/src/generated/storage.rs @@ -22,7 +22,7 @@ pub mod submit_job_response { #[prost(uint64, tag = "1")] JobId(u64), #[prost(message, tag = "2")] - Error(super::StorageError), + Error(super::JobManagementError), } } #[derive(Clone, Copy, PartialEq, ::prost::Message)] @@ -44,7 +44,7 @@ pub mod job_state_response { #[prost(enumeration = "super::JobState", tag = "1")] State(i32), #[prost(message, tag = "2")] - Error(super::StorageError), + Error(super::JobManagementError), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -59,7 +59,7 @@ pub mod job_outputs_response { #[prost(message, tag = "1")] Outputs(super::JobOutputs), #[prost(message, tag = "2")] - Error(super::StorageError), + Error(super::JobManagementError), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -79,7 +79,7 @@ pub mod job_error_response { #[prost(string, tag = "1")] ErrorMessage(::prost::alloc::string::String), #[prost(message, tag = "2")] - Error(super::StorageError), + Error(super::JobManagementError), } } #[derive(Clone, Copy, PartialEq, ::prost::Message)] @@ -101,7 +101,7 @@ pub mod delete_expired_terminated_jobs_response { #[prost(message, tag = "1")] DeletedJobs(super::DeletedJobs), #[prost(message, tag = "2")] - Error(super::StorageError), + Error(super::JobManagementError), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -135,7 +135,7 @@ pub mod poll_ready_tasks_response { #[prost(message, tag = "1")] Tasks(super::ReadyTasks), #[prost(message, tag = "2")] - Error(super::StorageError), + Error(super::SchedulerStorageError), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -177,7 +177,7 @@ pub mod register_task_instance_response { #[prost(bytes, tag = "1")] ExecutionContext(::prost::alloc::vec::Vec), #[prost(message, tag = "2")] - Error(super::StorageError), + Error(super::TaskInstanceError), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -234,7 +234,7 @@ pub mod resource_group_id_response { #[prost(uint64, tag = "1")] ResourceGroupId(u64), #[prost(message, tag = "2")] - Error(super::StorageError), + Error(super::ResourceGroupError), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -275,7 +275,7 @@ pub mod register_execution_manager_response { #[prost(message, tag = "1")] Registration(super::ExecutionManagerRegistration), #[prost(message, tag = "2")] - Error(super::StorageError), + Error(super::ExecutionManagerLivenessError), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -295,7 +295,7 @@ pub mod update_execution_manager_heartbeat_response { #[prost(uint64, tag = "1")] SessionId(u64), #[prost(message, tag = "2")] - Error(super::StorageError), + Error(super::ExecutionManagerLivenessError), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -310,7 +310,7 @@ pub mod is_execution_manager_alive_response { #[prost(bool, tag = "1")] Alive(bool), #[prost(message, tag = "2")] - Error(super::StorageError), + Error(super::ExecutionManagerLivenessError), } } #[derive(Clone, Copy, PartialEq, ::prost::Message)] @@ -332,7 +332,7 @@ pub mod get_dead_execution_managers_response { #[prost(message, tag = "1")] DeadExecutionManagers(super::DeadExecutionManagers), #[prost(message, tag = "2")] - Error(super::StorageError), + Error(super::ExecutionManagerLivenessError), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -363,33 +363,120 @@ pub mod task_id { } } #[derive(Clone, PartialEq, ::prost::Message)] -pub struct StorageOperationResponse { - #[prost(oneof = "storage_operation_response::Result", tags = "1, 2")] - pub result: ::core::option::Option, +pub struct JobManagementOperationResponse { + #[prost(oneof = "job_management_operation_response::Result", tags = "1, 2")] + pub result: ::core::option::Option, } -/// Nested message and enum types in `StorageOperationResponse`. -pub mod storage_operation_response { +/// Nested message and enum types in `JobManagementOperationResponse`. +pub mod job_management_operation_response { #[derive(Clone, PartialEq, ::prost::Oneof)] pub enum Result { #[prost(message, tag = "1")] Ok(super::Void), #[prost(message, tag = "2")] - Error(super::StorageError), + Error(super::JobManagementError), + } +} +#[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::TaskInstanceError), + } +} +#[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::ResourceGroupError), } } #[derive(Clone, Copy, PartialEq, ::prost::Message)] pub struct Void {} #[derive(Clone, PartialEq, ::prost::Message)] -pub struct StorageError { - #[prost(enumeration = "storage_error::ErrCode", tag = "1")] +pub struct JobManagementError { + #[prost(enumeration = "job_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 `StorageError`. -pub mod storage_error { +/// Nested message and enum types in `JobManagementError`. +pub mod job_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, + JobNotFound = 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::Server => "SERVER", + Self::InvalidInput => "INVALID_INPUT", + Self::JobNotFound => "JOB_NOT_FOUND", + } + } + /// 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), + "JOB_NOT_FOUND" => Some(Self::JobNotFound), + _ => None, + } + } + } +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct TaskInstanceError { + #[prost(enumeration = "task_instance_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 `TaskInstanceError`. +pub mod task_instance_error { #[derive( Clone, Copy, @@ -406,11 +493,8 @@ pub mod storage_error { Unspecified = 0, StaleSession = 1, CacheStale = 2, - Transport = 3, - Server = 4, - InvalidInput = 5, - InboundClosed = 6, - JobNotFound = 7, + Server = 3, + InvalidInput = 4, } impl ErrCode { /// String value of the enum field names used in the ProtoBuf definition. @@ -422,11 +506,8 @@ pub mod storage_error { Self::Unspecified => "ERR_CODE_UNSPECIFIED", Self::StaleSession => "STALE_SESSION", Self::CacheStale => "CACHE_STALE", - Self::Transport => "TRANSPORT", Self::Server => "SERVER", Self::InvalidInput => "INVALID_INPUT", - Self::InboundClosed => "INBOUND_CLOSED", - Self::JobNotFound => "JOB_NOT_FOUND", } } /// Creates an enum from field names used in the ProtoBuf definition. @@ -435,11 +516,166 @@ pub mod storage_error { "ERR_CODE_UNSPECIFIED" => Some(Self::Unspecified), "STALE_SESSION" => Some(Self::StaleSession), "CACHE_STALE" => Some(Self::CacheStale), - "TRANSPORT" => Some(Self::Transport), "SERVER" => Some(Self::Server), "INVALID_INPUT" => Some(Self::InvalidInput), + _ => None, + } + } + } +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct SchedulerStorageError { + #[prost(enumeration = "scheduler_storage_error::ErrCode", tag = "1")] + pub err_code: i32, + #[prost(string, tag = "2")] + pub message: ::prost::alloc::string::String, +} +/// Nested message and enum types in `SchedulerStorageError`. +pub mod scheduler_storage_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), - "JOB_NOT_FOUND" => Some(Self::JobNotFound), + "SERVER" => Some(Self::Server), + "INVALID_INPUT" => Some(Self::InvalidInput), + _ => None, + } + } + } +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ResourceGroupError { + #[prost(enumeration = "resource_group_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 `ResourceGroupError`. +pub mod resource_group_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, } } @@ -759,7 +995,7 @@ pub mod job_management_service_client { &mut self, request: impl tonic::IntoRequest, ) -> std::result::Result< - tonic::Response, + tonic::Response, tonic::Status, > { self.inner @@ -909,7 +1145,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 @@ -938,7 +1174,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 @@ -1270,7 +1506,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 @@ -1299,7 +1535,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 @@ -1724,7 +1960,7 @@ pub mod job_management_service_server { &self, request: tonic::Request, ) -> std::result::Result< - tonic::Response, + tonic::Response, tonic::Status, >; } @@ -2145,7 +2381,7 @@ pub mod job_management_service_server { T: JobManagementService, > tonic::server::UnaryService for ResendReadyTasksSvc { - type Response = super::StorageOperationResponse; + type Response = super::JobManagementOperationResponse; type Future = BoxFuture< tonic::Response, tonic::Status, @@ -2249,14 +2485,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, >; } @@ -2397,7 +2633,7 @@ pub mod task_instance_management_service_server { T: TaskInstanceManagementService, > tonic::server::UnaryService for ReportTaskSuccessSvc { - type Response = super::StorageOperationResponse; + type Response = super::TaskInstanceOperationResponse; type Future = BoxFuture< tonic::Response, tonic::Status, @@ -2448,7 +2684,7 @@ pub mod task_instance_management_service_server { T: TaskInstanceManagementService, > tonic::server::UnaryService for ReportTaskFailureSvc { - type Response = super::StorageOperationResponse; + type Response = super::TaskInstanceOperationResponse; type Future = BoxFuture< tonic::Response, tonic::Status, @@ -2853,14 +3089,14 @@ pub mod resource_group_management_service_server { &self, request: tonic::Request, ) -> std::result::Result< - tonic::Response, + tonic::Response, tonic::Status, >; async fn delete_resource_group( &self, request: tonic::Request, ) -> std::result::Result< - tonic::Response, + tonic::Response, tonic::Status, >; } @@ -3001,7 +3237,7 @@ pub mod resource_group_management_service_server { T: ResourceGroupManagementService, > tonic::server::UnaryService for VerifyResourceGroupSvc { - type Response = super::StorageOperationResponse; + type Response = super::ResourceGroupOperationResponse; type Future = BoxFuture< tonic::Response, tonic::Status, @@ -3052,7 +3288,7 @@ pub mod resource_group_management_service_server { T: ResourceGroupManagementService, > tonic::server::UnaryService for DeleteResourceGroupSvc { - type Response = super::StorageOperationResponse; + type Response = super::ResourceGroupOperationResponse; type Future = BoxFuture< tonic::Response, tonic::Status, diff --git a/components/spider-proto/storage/storage.proto b/components/spider-proto/storage/storage.proto index 256ca4fc..7de2defd 100644 --- a/components/spider-proto/storage/storage.proto +++ b/components/spider-proto/storage/storage.proto @@ -11,13 +11,13 @@ service JobManagementService { rpc GetJobError(JobIdRequest) returns (JobErrorResponse); rpc DeleteExpiredTerminatedJobs(DeleteExpiredTerminatedJobsRequest) returns (DeleteExpiredTerminatedJobsResponse); - rpc ResendReadyTasks(ResendReadyTasksRequest) returns (StorageOperationResponse); + rpc ResendReadyTasks(ResendReadyTasksRequest) returns (JobManagementOperationResponse); } service TaskInstanceManagementService { rpc RegisterTaskInstance(RegisterTaskInstanceRequest) returns (RegisterTaskInstanceResponse); - rpc ReportTaskSuccess(ReportTaskSuccessRequest) returns (StorageOperationResponse); - rpc ReportTaskFailure(ReportTaskFailureRequest) returns (StorageOperationResponse); + rpc ReportTaskSuccess(ReportTaskSuccessRequest) returns (TaskInstanceOperationResponse); + rpc ReportTaskFailure(ReportTaskFailureRequest) returns (TaskInstanceOperationResponse); } service SchedulerStorageService { @@ -28,8 +28,8 @@ service SchedulerStorageService { service ResourceGroupManagementService { rpc AddResourceGroup(AddResourceGroupRequest) returns (ResourceGroupIdResponse); - rpc VerifyResourceGroup(VerifyResourceGroupRequest) returns (StorageOperationResponse); - rpc DeleteResourceGroup(ResourceGroupIdRequest) returns (StorageOperationResponse); + rpc VerifyResourceGroup(VerifyResourceGroupRequest) returns (ResourceGroupOperationResponse); + rpc DeleteResourceGroup(ResourceGroupIdRequest) returns (ResourceGroupOperationResponse); } service ExecutionManagerLivenessService { @@ -57,7 +57,7 @@ message SubmitJobRequest { message SubmitJobResponse { oneof result { uint64 job_id = 1; - StorageError error = 2; + JobManagementError error = 2; } } @@ -69,14 +69,14 @@ message JobIdRequest { message JobStateResponse { oneof result { JobState state = 1; - StorageError error = 2; + JobManagementError error = 2; } } message JobOutputsResponse { oneof result { JobOutputs outputs = 1; - StorageError error = 2; + JobManagementError error = 2; } } @@ -87,7 +87,7 @@ message JobOutputs { message JobErrorResponse { oneof result { string error_message = 1; - StorageError error = 2; + JobManagementError error = 2; } } @@ -99,7 +99,7 @@ message DeleteExpiredTerminatedJobsRequest { message DeleteExpiredTerminatedJobsResponse { oneof result { DeletedJobs deleted_jobs = 1; - StorageError error = 2; + JobManagementError error = 2; } } @@ -120,7 +120,7 @@ message PollReadyTasksRequest { message PollReadyTasksResponse { oneof result { ReadyTasks tasks = 1; - StorageError error = 2; + SchedulerStorageError error = 2; } } @@ -145,7 +145,7 @@ message RegisterTaskInstanceRequest { message RegisterTaskInstanceResponse { oneof result { bytes execution_context = 1; - StorageError error = 2; + TaskInstanceError error = 2; } } @@ -179,7 +179,7 @@ message ResourceGroupIdRequest { message ResourceGroupIdResponse { oneof result { uint64 resource_group_id = 1; - StorageError error = 2; + ResourceGroupError error = 2; } } @@ -205,21 +205,21 @@ message ExecutionManagerRegistration { message RegisterExecutionManagerResponse { oneof result { ExecutionManagerRegistration registration = 1; - StorageError error = 2; + ExecutionManagerLivenessError error = 2; } } message UpdateExecutionManagerHeartbeatResponse { oneof result { uint64 session_id = 1; - StorageError error = 2; + ExecutionManagerLivenessError error = 2; } } message IsExecutionManagerAliveResponse { oneof result { bool alive = 1; - StorageError error = 2; + ExecutionManagerLivenessError error = 2; } } @@ -231,7 +231,7 @@ message GetDeadExecutionManagersRequest { message GetDeadExecutionManagersResponse { oneof result { DeadExecutionManagers dead_execution_managers = 1; - StorageError error = 2; + ExecutionManagerLivenessError error = 2; } } @@ -262,28 +262,90 @@ enum JobState { CANCELLED = 7; } -message StorageOperationResponse { +message JobManagementOperationResponse { oneof result { Void ok = 1; - StorageError error = 2; + JobManagementError error = 2; + } +} + +message TaskInstanceOperationResponse { + oneof result { + Void ok = 1; + TaskInstanceError error = 2; + } +} + +message ResourceGroupOperationResponse { + oneof result { + Void ok = 1; + ResourceGroupError error = 2; } } message Void {} -message StorageError { +message JobManagementError { + enum ErrCode { + ERR_CODE_UNSPECIFIED = 0; + STALE_SESSION = 1; + SERVER = 2; + INVALID_INPUT = 3; + JOB_NOT_FOUND = 4; + } + + ErrCode err_code = 1; + string message = 2; + uint64 storage_session = 3; +} + +message TaskInstanceError { enum ErrCode { ERR_CODE_UNSPECIFIED = 0; STALE_SESSION = 1; CACHE_STALE = 2; - TRANSPORT = 3; - SERVER = 4; - INVALID_INPUT = 5; - INBOUND_CLOSED = 6; - JOB_NOT_FOUND = 7; + SERVER = 3; + INVALID_INPUT = 4; + } + + ErrCode err_code = 1; + string message = 2; + uint64 storage_session = 3; +} + +message SchedulerStorageError { + enum ErrCode { + ERR_CODE_UNSPECIFIED = 0; + INBOUND_CLOSED = 1; + SERVER = 2; + INVALID_INPUT = 3; + } + + ErrCode err_code = 1; + string message = 2; +} + +message ResourceGroupError { + 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; +} From c36c6e1823764ea18c6d918771fbf2db9e257200 Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Thu, 11 Jun 2026 13:32:08 -0400 Subject: [PATCH 08/16] Update error code change --- .../src/client/grpc/liveness.rs | 33 +++++++------- .../src/client/grpc/storage.rs | 44 +++++++++---------- 2 files changed, 40 insertions(+), 37 deletions(-) diff --git a/components/spider-execution-manager/src/client/grpc/liveness.rs b/components/spider-execution-manager/src/client/grpc/liveness.rs index 3cfe32ea..8292cef8 100644 --- a/components/spider-execution-manager/src/client/grpc/liveness.rs +++ b/components/spider-execution-manager/src/client/grpc/liveness.rs @@ -6,9 +6,9 @@ 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, - storage_error, update_execution_manager_heartbeat_response, }; use tonic::transport::{Channel, Endpoint}; @@ -125,19 +125,23 @@ fn heartbeat_response_to_result( /// # Returns /// /// The corresponding [`LivenessResponseError`]. -fn storage_error_to_liveness_error(error: storage::StorageError) -> LivenessResponseError { - match storage_error::ErrCode::try_from(error.err_code) { - Ok(storage_error::ErrCode::CacheStale) => LivenessResponseError::MarkedDead, - Ok(storage_error::ErrCode::InvalidInput) => LivenessResponseError::IllegalId(error.message), +fn storage_error_to_liveness_error( + error: storage::ExecutionManagerLivenessError, +) -> LivenessResponseError { + match execution_manager_liveness_error::ErrCode::try_from(error.err_code) { + Ok(execution_manager_liveness_error::ErrCode::MarkedDead) => { + LivenessResponseError::MarkedDead + } + Ok(execution_manager_liveness_error::ErrCode::InvalidInput) => { + LivenessResponseError::IllegalId(error.message) + } Ok( - storage_error::ErrCode::Transport - | storage_error::ErrCode::Server - | storage_error::ErrCode::StaleSession - | storage_error::ErrCode::Unspecified, + execution_manager_liveness_error::ErrCode::Server + | execution_manager_liveness_error::ErrCode::Unspecified, ) => LivenessResponseError::Transport(error.message), - Err(error) => { - LivenessResponseError::Transport(format!("unknown storage error kind: {error}")) - } + Err(error) => LivenessResponseError::Transport(format!( + "unknown execution manager liveness error kind: {error}" + )), } } @@ -194,10 +198,9 @@ mod tests { #[test] fn liveness_storage_error_maps_invalid_input_to_illegal_id() { - let error = storage::StorageError { - err_code: storage_error::ErrCode::InvalidInput.into(), + let error = storage::ExecutionManagerLivenessError { + err_code: execution_manager_liveness_error::ErrCode::InvalidInput.into(), message: "bad em id".to_owned(), - storage_session: 0, }; match storage_error_to_liveness_error(error) { diff --git a/components/spider-execution-manager/src/client/grpc/storage.rs b/components/spider-execution-manager/src/client/grpc/storage.rs index c037814a..0c9afca5 100644 --- a/components/spider-execution-manager/src/client/grpc/storage.rs +++ b/components/spider-execution-manager/src/client/grpc/storage.rs @@ -11,9 +11,9 @@ use spider_core::types::{ use spider_proto_rust::storage::{ self, register_task_instance_response, - storage_error, - storage_operation_response, + task_instance_error, task_instance_management_service_client::TaskInstanceManagementServiceClient, + task_instance_operation_response, }; use tonic::transport::{Channel, Endpoint}; @@ -136,32 +136,32 @@ impl StorageClient for GrpcStorageClient { } } -impl From for StorageResponseError { - fn from(error: storage::StorageError) -> Self { - match storage_error::ErrCode::try_from(error.err_code) { - Ok(storage_error::ErrCode::StaleSession) => Self::StaleSession { +impl From for StorageResponseError { + fn from(error: storage::TaskInstanceError) -> Self { + match task_instance_error::ErrCode::try_from(error.err_code) { + Ok(task_instance_error::ErrCode::StaleSession) => Self::StaleSession { storage_session: error.storage_session, }, - Ok(storage_error::ErrCode::CacheStale) => Self::CacheStale(error.message), - Ok(storage_error::ErrCode::Transport) => Self::Transport(error.message), - Ok(storage_error::ErrCode::Server | storage_error::ErrCode::Unspecified) => { - Self::Server(error.message) - } - Ok(storage_error::ErrCode::InvalidInput) => Self::InvalidInput(error.message), - Err(error) => Self::Transport(format!("unknown storage error kind: {error}")), + Ok(task_instance_error::ErrCode::CacheStale) => Self::CacheStale(error.message), + Ok( + task_instance_error::ErrCode::Server | task_instance_error::ErrCode::Unspecified, + ) => Self::Server(error.message), + Ok(task_instance_error::ErrCode::InvalidInput) => Self::InvalidInput(error.message), + Err(error) => Self::Transport(format!("unknown task instance error kind: {error}")), } } } /// # Returns /// -/// [`storage::StorageOperationResponse`] converted into [`Result<(), StorageResponseError>`]. +/// [`storage::TaskInstanceOperationResponse`] converted into +/// [`Result<(), StorageResponseError>`]. fn storage_operation_response_to_result( - response: storage::StorageOperationResponse, + response: storage::TaskInstanceOperationResponse, ) -> Result<(), StorageResponseError> { match response.result { - Some(storage_operation_response::Result::Ok(_)) => Ok(()), - Some(storage_operation_response::Result::Error(error)) => Err(error.into()), + 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(), )), @@ -183,8 +183,8 @@ mod tests { #[test] fn storage_error_maps_stale_session() { - let error = storage::StorageError { - err_code: storage_error::ErrCode::StaleSession.into(), + let error = storage::TaskInstanceError { + err_code: task_instance_error::ErrCode::StaleSession.into(), message: "stale".to_owned(), storage_session: 7, }; @@ -199,7 +199,7 @@ mod tests { #[test] fn storage_error_maps_unknown_kind_to_transport_error() { - let error = storage::StorageError { + let error = storage::TaskInstanceError { err_code: 99, message: "unknown".to_owned(), storage_session: 0, @@ -207,7 +207,7 @@ mod tests { match StorageResponseError::from(error) { StorageResponseError::Transport(message) => { - assert!(message.contains("unknown storage error kind")); + assert!(message.contains("unknown task instance error kind")); } error => panic!("unexpected storage response error: {error:?}"), } @@ -215,7 +215,7 @@ mod tests { #[test] fn missing_storage_operation_result_is_transport_error() { - match storage_operation_response_to_result(storage::StorageOperationResponse { + match storage_operation_response_to_result(storage::TaskInstanceOperationResponse { result: None, }) { Err(StorageResponseError::Transport(_)) => {} From dcffc75bdf7105ece768fb046d843b09a034a9dd Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Thu, 11 Jun 2026 13:38:58 -0400 Subject: [PATCH 09/16] Add grpc impl of scheduler's storage client --- components/spider-scheduler/Cargo.toml | 2 + components/spider-scheduler/src/error.rs | 19 + components/spider-scheduler/src/grpc/mod.rs | 5 + .../src/grpc/storage_client.rs | 392 ++++++++++++++++++ components/spider-scheduler/src/lib.rs | 2 + 5 files changed, 420 insertions(+) create mode 100644 components/spider-scheduler/src/grpc/mod.rs create mode 100644 components/spider-scheduler/src/grpc/storage_client.rs diff --git a/components/spider-scheduler/Cargo.toml b/components/spider-scheduler/Cargo.toml index 59e9b8f7..ac44cc48 100644 --- a/components/spider-scheduler/Cargo.toml +++ b/components/spider-scheduler/Cargo.toml @@ -10,5 +10,7 @@ path = "src/lib.rs" [dependencies] async-trait = "0.1.89" spider-core = { path = "../spider-core" } +spider-proto-rust = { path = "../spider-proto-rust" } thiserror = "2.0.18" +tonic = "0.12.3" tokio-util = "0.7.18" diff --git a/components/spider-scheduler/src/error.rs b/components/spider-scheduler/src/error.rs index 6a852c46..18519da2 100644 --- a/components/spider-scheduler/src/error.rs +++ b/components/spider-scheduler/src/error.rs @@ -12,6 +12,25 @@ pub enum StorageClientError { /// No job with the requested identifier exists. #[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), + + /// The storage server returned an otherwise-uncategorized error. + #[error("storage server error: {0}")] + Server(String), + + /// The storage transport failed or returned malformed data. + #[error("storage transport error: {0}")] + Transport(String), } /// Errors returned by the scheduler runtime and its components. diff --git a/components/spider-scheduler/src/grpc/mod.rs b/components/spider-scheduler/src/grpc/mod.rs new file mode 100644 index 00000000..84954f46 --- /dev/null +++ b/components/spider-scheduler/src/grpc/mod.rs @@ -0,0 +1,5 @@ +//! gRPC-backed scheduler client implementations. + +mod storage_client; + +pub use storage_client::GrpcSchedulerStorageClient; diff --git a/components/spider-scheduler/src/grpc/storage_client.rs b/components/spider-scheduler/src/grpc/storage_client.rs new file mode 100644 index 00000000..a3fdc5be --- /dev/null +++ b/components/spider-scheduler/src/grpc/storage_client.rs @@ -0,0 +1,392 @@ +//! gRPC-backed [`SchedulerStorageClient`] implementation. + +use std::time::Duration; + +use async_trait::async_trait; +use spider_core::{ + job::JobState, + session::SessionTracker, + types::id::{JobId, ResourceGroupId, SessionId, TaskId}, +}; +use spider_proto_rust::storage::{ + self, + job_management_error, + job_management_service_client::JobManagementServiceClient, + job_state_response, + poll_ready_tasks_response, + scheduler_storage_error, + scheduler_storage_service_client::SchedulerStorageServiceClient, + session_management_service_client::SessionManagementServiceClient, +}; +use tonic::transport::{Channel, Endpoint}; + +use crate::{ + error::StorageClientError, + storage_client::SchedulerStorageClient, + types::InboundEntry, +}; + +/// gRPC-backed [`SchedulerStorageClient`] implementation. +#[derive(Debug, Clone)] +pub struct GrpcSchedulerStorageClient { + scheduler_client: SchedulerStorageServiceClient, + job_client: JobManagementServiceClient, + session_tracker: SessionTracker, +} + +impl GrpcSchedulerStorageClient { + /// Connects to the storage gRPC endpoint. + /// + /// # Returns + /// + /// A new [`GrpcSchedulerStorageClient`] connected to `endpoint` on success. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`StorageClientError::Transport`] if tonic cannot create or connect to the endpoint. + /// * Forwards [`Self::get_initial_session`]'s return values on failure. + pub async fn connect(endpoint: Endpoint) -> Result { + let channel = endpoint.connect().await.map_err(to_transport_error)?; + let session_id = Self::get_initial_session(channel.clone()).await?; + + Ok(Self { + scheduler_client: SchedulerStorageServiceClient::new(channel.clone()), + job_client: JobManagementServiceClient::new(channel), + session_tracker: SessionTracker::new(session_id), + }) + } + + /// Fetches storage's current session ID. + /// + /// # Returns + /// + /// The current storage session ID on success. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`StorageClientError::Transport`] if the request fails. + async fn get_initial_session(channel: Channel) -> Result { + SessionManagementServiceClient::new(channel) + .get_session(storage::Void {}) + .await + .map(|response| response.into_inner().session_id) + .map_err(to_transport_error) + } +} + +#[async_trait] +impl SchedulerStorageClient for GrpcSchedulerStorageClient { + async fn poll_ready( + &self, + max_items: usize, + wait: Duration, + ) -> Result<(SessionId, Vec), StorageClientError> { + let request = poll_ready_tasks_request(max_items, wait)?; + let response = self + .scheduler_client + .clone() + .poll_ready_tasks(request) + .await + .map_err(to_transport_error)? + .into_inner(); + poll_ready_tasks_response_to_result(response) + } + + async fn poll_commit_ready( + &self, + max_items: usize, + wait: Duration, + ) -> Result<(SessionId, Vec), StorageClientError> { + let request = poll_ready_tasks_request(max_items, wait)?; + let response = self + .scheduler_client + .clone() + .poll_ready_commit_tasks(request) + .await + .map_err(to_transport_error)? + .into_inner(); + poll_ready_tasks_response_to_result(response) + } + + async fn poll_cleanup_ready( + &self, + max_items: usize, + wait: Duration, + ) -> Result<(SessionId, Vec), StorageClientError> { + let request = poll_ready_tasks_request(max_items, wait)?; + let response = self + .scheduler_client + .clone() + .poll_ready_cleanup_tasks(request) + .await + .map_err(to_transport_error)? + .into_inner(); + poll_ready_tasks_response_to_result(response) + } + + async fn job_state(&self, job_id: JobId) -> Result { + let request = storage::JobIdRequest { + job_id: job_id.get(), + session_id: self.session_tracker.current(), + }; + let response = self + .job_client + .clone() + .get_job_state(request) + .await + .map_err(to_transport_error)? + .into_inner(); + job_state_response_to_result(response, job_id) + } +} + +impl From for StorageClientError { + fn from(error: storage::SchedulerStorageError) -> Self { + scheduler_storage_error_to_client_error(error) + } +} + +/// # Returns +/// +/// A [`storage::PollReadyTasksRequest`] carrying `max_items` and `wait` on success. +/// +/// # Errors +/// +/// Returns an error if: +/// +/// * [`StorageClientError::Transport`] if either value cannot fit in the protobuf field type. +fn poll_ready_tasks_request( + max_items: usize, + wait: Duration, +) -> Result { + Ok(storage::PollReadyTasksRequest { + max_items: u64::try_from(max_items).map_err(to_transport_error)?, + wait_ns: u64::try_from(wait.as_nanos()).map_err(to_transport_error)?, + }) +} + +/// # Returns +/// +/// [`storage::PollReadyTasksResponse`] converted into scheduler entries on success. +/// +/// # Errors +/// +/// Returns an error if: +/// +/// * [`StorageClientError::Transport`] if the response is malformed. +/// * Forwards [`StorageClientError::from`]'s return values on failure. +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(), + )), + } +} + +/// # Returns +/// +/// [`storage::ReadyTasks`] converted into scheduler entries on success. +/// +/// # Errors +/// +/// Returns an error if: +/// +/// * [`StorageClientError::Transport`] if a ready task is missing or has an invalid task ID. +fn ready_tasks_to_result( + tasks: storage::ReadyTasks, +) -> Result<(SessionId, Vec), StorageClientError> { + let session_id = tasks.session_id; + let entries = tasks + .tasks + .into_iter() + .map(ready_task_to_inbound_entry) + .collect::, _>>()?; + Ok((session_id, entries)) +} + +/// # Returns +/// +/// `task` converted into an [`InboundEntry`] on success. +/// +/// # Errors +/// +/// Returns an error if: +/// +/// * [`StorageClientError::Transport`] if `task` is missing or has an invalid task ID. +fn ready_task_to_inbound_entry( + task: storage::ReadyTask, +) -> Result { + let task_id = task + .task_id + .ok_or_else(|| StorageClientError::Transport("ready task missing task ID".to_owned())) + .and_then(|task_id| TaskId::try_from(task_id).map_err(StorageClientError::Transport))?; + Ok(InboundEntry { + resource_group_id: ResourceGroupId::from(task.resource_group_id), + job_id: JobId::from(task.job_id), + task_id, + }) +} + +/// # Returns +/// +/// [`storage::JobStateResponse`] converted into [`JobState`] on success. +/// +/// # Errors +/// +/// Returns an error if: +/// +/// * [`StorageClientError::Transport`] if the response is malformed. +/// * Forwards [`StorageClientError::from`]'s return values on failure. +fn job_state_response_to_result( + response: storage::JobStateResponse, + job_id: JobId, +) -> Result { + match response.result { + Some(job_state_response::Result::State(state)) => storage::JobState::try_from(state) + .map_err(|error| StorageClientError::Transport(error.to_string())) + .and_then(|state| JobState::try_from(state).map_err(StorageClientError::Transport)), + Some(job_state_response::Result::Error(error)) => { + Err(job_management_error_to_client_error(error, job_id)) + } + None => Err(StorageClientError::Transport( + "job state response missing `result` message".to_owned(), + )), + } +} + +/// # Returns +/// +/// [`storage::SchedulerStorageError`] converted into [`StorageClientError`]. +fn scheduler_storage_error_to_client_error( + error: storage::SchedulerStorageError, +) -> StorageClientError { + match scheduler_storage_error::ErrCode::try_from(error.err_code) { + Ok(scheduler_storage_error::ErrCode::InboundClosed) => StorageClientError::InboundClosed, + Ok(scheduler_storage_error::ErrCode::InvalidInput) => { + StorageClientError::InvalidInput(error.message) + } + Ok( + scheduler_storage_error::ErrCode::Server + | scheduler_storage_error::ErrCode::Unspecified, + ) => StorageClientError::Server(error.message), + Err(error) => { + StorageClientError::Transport(format!("unknown scheduler storage error kind: {error}")) + } + } +} + +/// # Returns +/// +/// [`storage::JobManagementError`] converted into [`StorageClientError`]. +fn job_management_error_to_client_error( + error: storage::JobManagementError, + requested_job_id: JobId, +) -> StorageClientError { + match job_management_error::ErrCode::try_from(error.err_code) { + Ok(job_management_error::ErrCode::JobNotFound) => { + StorageClientError::JobNotFound(requested_job_id) + } + Ok(job_management_error::ErrCode::StaleSession) => StorageClientError::StaleSession { + storage_session: error.storage_session, + }, + Ok(job_management_error::ErrCode::InvalidInput) => { + StorageClientError::InvalidInput(error.message) + } + Ok(job_management_error::ErrCode::Server | job_management_error::ErrCode::Unspecified) => { + StorageClientError::Server(error.message) + } + Err(error) => { + StorageClientError::Transport(format!("unknown job management error kind: {error}")) + } + } +} + +/// Converts a displayable transport-layer error into [`StorageClientError::Transport`]. +/// +/// # Returns +/// +/// A [`StorageClientError::Transport`] containing `error`'s display string. +fn to_transport_error(error: impl std::fmt::Display) -> StorageClientError { + StorageClientError::Transport(error.to_string()) +} + +#[cfg(test)] +mod tests { + use spider_core::types::id::{JobId, ResourceGroupId, TaskId}; + use spider_proto_rust::storage::{self, poll_ready_tasks_response, scheduler_storage_error}; + + use super::*; + + #[test] + fn poll_ready_tasks_response_converts_entries() { + let response = storage::PollReadyTasksResponse { + result: Some(poll_ready_tasks_response::Result::Tasks( + storage::ReadyTasks { + session_id: 11, + tasks: vec![storage::ReadyTask { + resource_group_id: 3, + job_id: 5, + task_id: Some(storage::TaskId::from(TaskId::Index(7))), + }], + }, + )), + }; + + let (session_id, entries) = poll_ready_tasks_response_to_result(response) + .expect("poll response conversion should succeed"); + + assert_eq!(session_id, 11); + assert_eq!( + entries, + vec![InboundEntry { + resource_group_id: ResourceGroupId::from(3), + job_id: JobId::from(5), + task_id: TaskId::Index(7), + }] + ); + } + + #[test] + fn poll_ready_tasks_response_rejects_missing_task_id() { + let response = storage::PollReadyTasksResponse { + result: Some(poll_ready_tasks_response::Result::Tasks( + storage::ReadyTasks { + session_id: 11, + tasks: vec![storage::ReadyTask { + resource_group_id: 3, + job_id: 5, + task_id: None, + }], + }, + )), + }; + + match poll_ready_tasks_response_to_result(response) { + Err(StorageClientError::Transport(message)) => { + assert!(message.contains("missing task ID")); + } + result => panic!("unexpected poll response conversion result: {result:?}"), + } + } + + #[test] + fn scheduler_storage_error_maps_inbound_closed() { + let error = storage::SchedulerStorageError { + err_code: scheduler_storage_error::ErrCode::InboundClosed.into(), + message: "closed".to_owned(), + }; + + assert!(matches!( + StorageClientError::from(error), + StorageClientError::InboundClosed + )); + } +} diff --git a/components/spider-scheduler/src/lib.rs b/components/spider-scheduler/src/lib.rs index bddd0750..00037216 100644 --- a/components/spider-scheduler/src/lib.rs +++ b/components/spider-scheduler/src/lib.rs @@ -34,6 +34,7 @@ pub mod core; pub mod dispatch_queue; pub mod error; +pub mod grpc; pub mod storage_client; pub mod types; @@ -41,6 +42,7 @@ pub use crate::{ core::SchedulerCore, dispatch_queue::{DispatchQueueSink, DispatchQueueSource}, error::{SchedulerError, StorageClientError}, + grpc::GrpcSchedulerStorageClient, storage_client::SchedulerStorageClient, types::{InboundEntry, TaskAssignment}, }; From a3a3f0908a4d93152600e4533f9f6930aa793b5c Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Thu, 11 Jun 2026 13:39:13 -0400 Subject: [PATCH 10/16] Update lockfile --- Cargo.lock | 2 ++ 1 file changed, 2 insertions(+) diff --git a/Cargo.lock b/Cargo.lock index f766b5f7..9b5cbe34 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1869,8 +1869,10 @@ version = "0.1.0" dependencies = [ "async-trait", "spider-core", + "spider-proto-rust", "thiserror", "tokio-util", + "tonic", ] [[package]] From fbc1c4cd7751c817d1a753cf3a689ac8e58ec465 Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Thu, 11 Jun 2026 13:42:24 -0400 Subject: [PATCH 11/16] Fix toml lint --- components/spider-scheduler/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/components/spider-scheduler/Cargo.toml b/components/spider-scheduler/Cargo.toml index ac44cc48..aa3be90f 100644 --- a/components/spider-scheduler/Cargo.toml +++ b/components/spider-scheduler/Cargo.toml @@ -12,5 +12,5 @@ async-trait = "0.1.89" spider-core = { path = "../spider-core" } spider-proto-rust = { path = "../spider-proto-rust" } thiserror = "2.0.18" -tonic = "0.12.3" tokio-util = "0.7.18" +tonic = "0.12.3" From 5279ecc351a2a46beaa28fbd1821e1efadab818d Mon Sep 17 00:00:00 2001 From: LinZhihao-723 Date: Fri, 12 Jun 2026 15:11:03 -0400 Subject: [PATCH 12/16] Reviewed proto source and apply changes accordingly. --- .../src/client/grpc/storage.rs | 27 +- .../src/generated/storage.rs | 498 +++++------------- components/spider-proto/storage/storage.proto | 63 +-- .../src/grpc/storage_client.rs | 75 +-- 4 files changed, 217 insertions(+), 446 deletions(-) diff --git a/components/spider-execution-manager/src/client/grpc/storage.rs b/components/spider-execution-manager/src/client/grpc/storage.rs index 0c9afca5..1a7df78a 100644 --- a/components/spider-execution-manager/src/client/grpc/storage.rs +++ b/components/spider-execution-manager/src/client/grpc/storage.rs @@ -11,7 +11,7 @@ use spider_core::types::{ use spider_proto_rust::storage::{ self, register_task_instance_response, - task_instance_error, + task_instance_management_error, task_instance_management_service_client::TaskInstanceManagementServiceClient, task_instance_operation_response, }; @@ -136,17 +136,22 @@ impl StorageClient for GrpcStorageClient { } } -impl From for StorageResponseError { - fn from(error: storage::TaskInstanceError) -> Self { - match task_instance_error::ErrCode::try_from(error.err_code) { - Ok(task_instance_error::ErrCode::StaleSession) => Self::StaleSession { +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_error::ErrCode::CacheStale) => Self::CacheStale(error.message), + Ok(task_instance_management_error::ErrCode::CacheStale) => { + Self::CacheStale(error.message) + } Ok( - task_instance_error::ErrCode::Server | task_instance_error::ErrCode::Unspecified, + task_instance_management_error::ErrCode::Server + | task_instance_management_error::ErrCode::Unspecified, ) => Self::Server(error.message), - Ok(task_instance_error::ErrCode::InvalidInput) => Self::InvalidInput(error.message), + Ok(task_instance_management_error::ErrCode::InvalidInput) => { + Self::InvalidInput(error.message) + } Err(error) => Self::Transport(format!("unknown task instance error kind: {error}")), } } @@ -183,8 +188,8 @@ mod tests { #[test] fn storage_error_maps_stale_session() { - let error = storage::TaskInstanceError { - err_code: task_instance_error::ErrCode::StaleSession.into(), + let error = storage::TaskInstanceManagementError { + err_code: task_instance_management_error::ErrCode::StaleSession.into(), message: "stale".to_owned(), storage_session: 7, }; @@ -199,7 +204,7 @@ mod tests { #[test] fn storage_error_maps_unknown_kind_to_transport_error() { - let error = storage::TaskInstanceError { + let error = storage::TaskInstanceManagementError { err_code: 99, message: "unknown".to_owned(), storage_session: 0, diff --git a/components/spider-proto-rust/src/generated/storage.rs b/components/spider-proto-rust/src/generated/storage.rs index b71d2aca..4323ce13 100644 --- a/components/spider-proto-rust/src/generated/storage.rs +++ b/components/spider-proto-rust/src/generated/storage.rs @@ -5,8 +5,8 @@ pub struct SubmitJobRequest { pub resource_group_id: u64, #[prost(bytes = "vec", tag = "2")] pub serialized_task_graph: ::prost::alloc::vec::Vec, - #[prost(bytes = "vec", repeated, tag = "3")] - pub serialized_inputs: ::prost::alloc::vec::Vec<::prost::alloc::vec::Vec>, + #[prost(bytes = "vec", tag = "3")] + pub serialized_inputs: ::prost::alloc::vec::Vec, #[prost(uint64, tag = "4")] pub session_id: u64, } @@ -22,7 +22,7 @@ pub mod submit_job_response { #[prost(uint64, tag = "1")] JobId(u64), #[prost(message, tag = "2")] - Error(super::JobManagementError), + Error(super::JobOrchestrationError), } } #[derive(Clone, Copy, PartialEq, ::prost::Message)] @@ -44,7 +44,7 @@ pub mod job_state_response { #[prost(enumeration = "super::JobState", tag = "1")] State(i32), #[prost(message, tag = "2")] - Error(super::JobManagementError), + Error(super::JobOrchestrationError), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -59,7 +59,7 @@ pub mod job_outputs_response { #[prost(message, tag = "1")] Outputs(super::JobOutputs), #[prost(message, tag = "2")] - Error(super::JobManagementError), + Error(super::JobOrchestrationError), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -79,7 +79,7 @@ pub mod job_error_response { #[prost(string, tag = "1")] ErrorMessage(::prost::alloc::string::String), #[prost(message, tag = "2")] - Error(super::JobManagementError), + Error(super::JobOrchestrationError), } } #[derive(Clone, Copy, PartialEq, ::prost::Message)] @@ -101,7 +101,7 @@ pub mod delete_expired_terminated_jobs_response { #[prost(message, tag = "1")] DeletedJobs(super::DeletedJobs), #[prost(message, tag = "2")] - Error(super::JobManagementError), + Error(super::JobOrchestrationError), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -135,7 +135,7 @@ pub mod poll_ready_tasks_response { #[prost(message, tag = "1")] Tasks(super::ReadyTasks), #[prost(message, tag = "2")] - Error(super::SchedulerStorageError), + Error(super::InboundQueueResponseError), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -177,7 +177,7 @@ pub mod register_task_instance_response { #[prost(bytes, tag = "1")] ExecutionContext(::prost::alloc::vec::Vec), #[prost(message, tag = "2")] - Error(super::TaskInstanceError), + Error(super::TaskInstanceManagementError), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -234,7 +234,7 @@ pub mod resource_group_id_response { #[prost(uint64, tag = "1")] ResourceGroupId(u64), #[prost(message, tag = "2")] - Error(super::ResourceGroupError), + Error(super::ResourceGroupManagementError), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -298,48 +298,6 @@ pub mod update_execution_manager_heartbeat_response { Error(super::ExecutionManagerLivenessError), } } -#[derive(Clone, PartialEq, ::prost::Message)] -pub struct IsExecutionManagerAliveResponse { - #[prost(oneof = "is_execution_manager_alive_response::Result", tags = "1, 2")] - pub result: ::core::option::Option, -} -/// Nested message and enum types in `IsExecutionManagerAliveResponse`. -pub mod is_execution_manager_alive_response { - #[derive(Clone, PartialEq, ::prost::Oneof)] - pub enum Result { - #[prost(bool, tag = "1")] - Alive(bool), - #[prost(message, tag = "2")] - Error(super::ExecutionManagerLivenessError), - } -} -#[derive(Clone, Copy, PartialEq, ::prost::Message)] -pub struct GetDeadExecutionManagersRequest { - #[prost(uint64, tag = "1")] - pub stale_after_sec: u64, - #[prost(uint64, tag = "2")] - pub session_id: u64, -} -#[derive(Clone, PartialEq, ::prost::Message)] -pub struct GetDeadExecutionManagersResponse { - #[prost(oneof = "get_dead_execution_managers_response::Result", tags = "1, 2")] - pub result: ::core::option::Option, -} -/// Nested message and enum types in `GetDeadExecutionManagersResponse`. -pub mod get_dead_execution_managers_response { - #[derive(Clone, PartialEq, ::prost::Oneof)] - pub enum Result { - #[prost(message, tag = "1")] - DeadExecutionManagers(super::DeadExecutionManagers), - #[prost(message, tag = "2")] - Error(super::ExecutionManagerLivenessError), - } -} -#[derive(Clone, PartialEq, ::prost::Message)] -pub struct DeadExecutionManagers { - #[prost(uint64, repeated, tag = "1")] - pub execution_manager_ids: ::prost::alloc::vec::Vec, -} #[derive(Clone, Copy, PartialEq, ::prost::Message)] pub struct GetSessionResponse { #[prost(uint64, tag = "1")] @@ -374,7 +332,7 @@ pub mod job_management_operation_response { #[prost(message, tag = "1")] Ok(super::Void), #[prost(message, tag = "2")] - Error(super::JobManagementError), + Error(super::JobOrchestrationError), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -389,7 +347,7 @@ pub mod task_instance_operation_response { #[prost(message, tag = "1")] Ok(super::Void), #[prost(message, tag = "2")] - Error(super::TaskInstanceError), + Error(super::TaskInstanceManagementError), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -404,22 +362,22 @@ pub mod resource_group_operation_response { #[prost(message, tag = "1")] Ok(super::Void), #[prost(message, tag = "2")] - Error(super::ResourceGroupError), + Error(super::ResourceGroupManagementError), } } #[derive(Clone, Copy, PartialEq, ::prost::Message)] pub struct Void {} #[derive(Clone, PartialEq, ::prost::Message)] -pub struct JobManagementError { - #[prost(enumeration = "job_management_error::ErrCode", tag = "1")] +pub struct JobOrchestrationError { + #[prost(enumeration = "job_orchestration_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 `JobManagementError`. -pub mod job_management_error { +/// Nested message and enum types in `JobOrchestrationError`. +pub mod job_orchestration_error { #[derive( Clone, Copy, @@ -467,16 +425,16 @@ pub mod job_management_error { } } #[derive(Clone, PartialEq, ::prost::Message)] -pub struct TaskInstanceError { - #[prost(enumeration = "task_instance_error::ErrCode", tag = "1")] +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 `TaskInstanceError`. -pub mod task_instance_error { +/// Nested message and enum types in `TaskInstanceManagementError`. +pub mod task_instance_management_error { #[derive( Clone, Copy, @@ -524,14 +482,14 @@ pub mod task_instance_error { } } #[derive(Clone, PartialEq, ::prost::Message)] -pub struct SchedulerStorageError { - #[prost(enumeration = "scheduler_storage_error::ErrCode", tag = "1")] +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 `SchedulerStorageError`. -pub mod scheduler_storage_error { +/// Nested message and enum types in `InboundQueueResponseError`. +pub mod inbound_queue_response_error { #[derive( Clone, Copy, @@ -576,16 +534,16 @@ pub mod scheduler_storage_error { } } #[derive(Clone, PartialEq, ::prost::Message)] -pub struct ResourceGroupError { - #[prost(enumeration = "resource_group_error::ErrCode", tag = "1")] +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 `ResourceGroupError`. -pub mod resource_group_error { +/// Nested message and enum types in `ResourceGroupManagementError`. +pub mod resource_group_management_error { #[derive( Clone, Copy, @@ -726,7 +684,7 @@ impl JobState { } } /// Generated client implementations. -pub mod job_management_service_client { +pub mod job_orchestration_service_client { #![allow( unused_variables, dead_code, @@ -737,10 +695,10 @@ pub mod job_management_service_client { use tonic::codegen::*; use tonic::codegen::http::Uri; #[derive(Debug, Clone)] - pub struct JobManagementServiceClient { + pub struct JobOrchestrationServiceClient { inner: tonic::client::Grpc, } - impl JobManagementServiceClient { + impl JobOrchestrationServiceClient { /// Attempt to create a new client by connecting to a given endpoint. pub async fn connect(dst: D) -> Result where @@ -751,7 +709,7 @@ pub mod job_management_service_client { Ok(Self::new(conn)) } } - impl JobManagementServiceClient + impl JobOrchestrationServiceClient where T: tonic::client::GrpcService, T::Error: Into, @@ -769,7 +727,7 @@ pub mod job_management_service_client { pub fn with_interceptor( inner: T, interceptor: F, - ) -> JobManagementServiceClient> + ) -> JobOrchestrationServiceClient> where F: tonic::service::Interceptor, T::ResponseBody: Default, @@ -783,7 +741,9 @@ pub mod job_management_service_client { http::Request, >>::Error: Into + std::marker::Send + std::marker::Sync, { - JobManagementServiceClient::new(InterceptedService::new(inner, interceptor)) + JobOrchestrationServiceClient::new( + InterceptedService::new(inner, interceptor), + ) } /// Compress requests with the given encoding. /// @@ -833,11 +793,11 @@ pub mod job_management_service_client { })?; let codec = tonic::codec::ProstCodec::default(); let path = http::uri::PathAndQuery::from_static( - "/storage.JobManagementService/SubmitJob", + "/storage.JobOrchestrationService/SubmitJob", ); let mut req = request.into_request(); req.extensions_mut() - .insert(GrpcMethod::new("storage.JobManagementService", "SubmitJob")); + .insert(GrpcMethod::new("storage.JobOrchestrationService", "SubmitJob")); self.inner.unary(req, path, codec).await } pub async fn start_job( @@ -857,11 +817,11 @@ pub mod job_management_service_client { })?; let codec = tonic::codec::ProstCodec::default(); let path = http::uri::PathAndQuery::from_static( - "/storage.JobManagementService/StartJob", + "/storage.JobOrchestrationService/StartJob", ); let mut req = request.into_request(); req.extensions_mut() - .insert(GrpcMethod::new("storage.JobManagementService", "StartJob")); + .insert(GrpcMethod::new("storage.JobOrchestrationService", "StartJob")); self.inner.unary(req, path, codec).await } pub async fn cancel_job( @@ -881,11 +841,11 @@ pub mod job_management_service_client { })?; let codec = tonic::codec::ProstCodec::default(); let path = http::uri::PathAndQuery::from_static( - "/storage.JobManagementService/CancelJob", + "/storage.JobOrchestrationService/CancelJob", ); let mut req = request.into_request(); req.extensions_mut() - .insert(GrpcMethod::new("storage.JobManagementService", "CancelJob")); + .insert(GrpcMethod::new("storage.JobOrchestrationService", "CancelJob")); self.inner.unary(req, path, codec).await } pub async fn get_job_state( @@ -905,11 +865,13 @@ pub mod job_management_service_client { })?; let codec = tonic::codec::ProstCodec::default(); let path = http::uri::PathAndQuery::from_static( - "/storage.JobManagementService/GetJobState", + "/storage.JobOrchestrationService/GetJobState", ); let mut req = request.into_request(); req.extensions_mut() - .insert(GrpcMethod::new("storage.JobManagementService", "GetJobState")); + .insert( + GrpcMethod::new("storage.JobOrchestrationService", "GetJobState"), + ); self.inner.unary(req, path, codec).await } pub async fn get_job_outputs( @@ -929,12 +891,12 @@ pub mod job_management_service_client { })?; let codec = tonic::codec::ProstCodec::default(); let path = http::uri::PathAndQuery::from_static( - "/storage.JobManagementService/GetJobOutputs", + "/storage.JobOrchestrationService/GetJobOutputs", ); let mut req = request.into_request(); req.extensions_mut() .insert( - GrpcMethod::new("storage.JobManagementService", "GetJobOutputs"), + GrpcMethod::new("storage.JobOrchestrationService", "GetJobOutputs"), ); self.inner.unary(req, path, codec).await } @@ -955,11 +917,13 @@ pub mod job_management_service_client { })?; let codec = tonic::codec::ProstCodec::default(); let path = http::uri::PathAndQuery::from_static( - "/storage.JobManagementService/GetJobError", + "/storage.JobOrchestrationService/GetJobError", ); let mut req = request.into_request(); req.extensions_mut() - .insert(GrpcMethod::new("storage.JobManagementService", "GetJobError")); + .insert( + GrpcMethod::new("storage.JobOrchestrationService", "GetJobError"), + ); self.inner.unary(req, path, codec).await } pub async fn delete_expired_terminated_jobs( @@ -979,13 +943,13 @@ pub mod job_management_service_client { })?; let codec = tonic::codec::ProstCodec::default(); let path = http::uri::PathAndQuery::from_static( - "/storage.JobManagementService/DeleteExpiredTerminatedJobs", + "/storage.JobOrchestrationService/DeleteExpiredTerminatedJobs", ); let mut req = request.into_request(); req.extensions_mut() .insert( GrpcMethod::new( - "storage.JobManagementService", + "storage.JobOrchestrationService", "DeleteExpiredTerminatedJobs", ), ); @@ -1008,12 +972,15 @@ pub mod job_management_service_client { })?; let codec = tonic::codec::ProstCodec::default(); let path = http::uri::PathAndQuery::from_static( - "/storage.JobManagementService/ResendReadyTasks", + "/storage.JobOrchestrationService/ResendReadyTasks", ); let mut req = request.into_request(); req.extensions_mut() .insert( - GrpcMethod::new("storage.JobManagementService", "ResendReadyTasks"), + GrpcMethod::new( + "storage.JobOrchestrationService", + "ResendReadyTasks", + ), ); self.inner.unary(req, path, codec).await } @@ -1202,7 +1169,7 @@ pub mod task_instance_management_service_client { } } /// Generated client implementations. -pub mod scheduler_storage_service_client { +pub mod inbound_queue_service_client { #![allow( unused_variables, dead_code, @@ -1213,10 +1180,10 @@ pub mod scheduler_storage_service_client { use tonic::codegen::*; use tonic::codegen::http::Uri; #[derive(Debug, Clone)] - pub struct SchedulerStorageServiceClient { + pub struct InboundQueueServiceClient { inner: tonic::client::Grpc, } - impl SchedulerStorageServiceClient { + impl InboundQueueServiceClient { /// Attempt to create a new client by connecting to a given endpoint. pub async fn connect(dst: D) -> Result where @@ -1227,7 +1194,7 @@ pub mod scheduler_storage_service_client { Ok(Self::new(conn)) } } - impl SchedulerStorageServiceClient + impl InboundQueueServiceClient where T: tonic::client::GrpcService, T::Error: Into, @@ -1245,7 +1212,7 @@ pub mod scheduler_storage_service_client { pub fn with_interceptor( inner: T, interceptor: F, - ) -> SchedulerStorageServiceClient> + ) -> InboundQueueServiceClient> where F: tonic::service::Interceptor, T::ResponseBody: Default, @@ -1259,9 +1226,7 @@ pub mod scheduler_storage_service_client { http::Request, >>::Error: Into + std::marker::Send + std::marker::Sync, { - SchedulerStorageServiceClient::new( - InterceptedService::new(inner, interceptor), - ) + InboundQueueServiceClient::new(InterceptedService::new(inner, interceptor)) } /// Compress requests with the given encoding. /// @@ -1311,12 +1276,12 @@ pub mod scheduler_storage_service_client { })?; let codec = tonic::codec::ProstCodec::default(); let path = http::uri::PathAndQuery::from_static( - "/storage.SchedulerStorageService/PollReadyTasks", + "/storage.InboundQueueService/PollReadyTasks", ); let mut req = request.into_request(); req.extensions_mut() .insert( - GrpcMethod::new("storage.SchedulerStorageService", "PollReadyTasks"), + GrpcMethod::new("storage.InboundQueueService", "PollReadyTasks"), ); self.inner.unary(req, path, codec).await } @@ -1337,13 +1302,13 @@ pub mod scheduler_storage_service_client { })?; let codec = tonic::codec::ProstCodec::default(); let path = http::uri::PathAndQuery::from_static( - "/storage.SchedulerStorageService/PollReadyCommitTasks", + "/storage.InboundQueueService/PollReadyCommitTasks", ); let mut req = request.into_request(); req.extensions_mut() .insert( GrpcMethod::new( - "storage.SchedulerStorageService", + "storage.InboundQueueService", "PollReadyCommitTasks", ), ); @@ -1366,13 +1331,13 @@ pub mod scheduler_storage_service_client { })?; let codec = tonic::codec::ProstCodec::default(); let path = http::uri::PathAndQuery::from_static( - "/storage.SchedulerStorageService/PollReadyCleanupTasks", + "/storage.InboundQueueService/PollReadyCleanupTasks", ); let mut req = request.into_request(); req.extensions_mut() .insert( GrpcMethod::new( - "storage.SchedulerStorageService", + "storage.InboundQueueService", "PollReadyCleanupTasks", ), ); @@ -1713,64 +1678,6 @@ pub mod execution_manager_liveness_service_client { ); self.inner.unary(req, path, codec).await } - pub async fn is_execution_manager_alive( - &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::codec::ProstCodec::default(); - let path = http::uri::PathAndQuery::from_static( - "/storage.ExecutionManagerLivenessService/IsExecutionManagerAlive", - ); - let mut req = request.into_request(); - req.extensions_mut() - .insert( - GrpcMethod::new( - "storage.ExecutionManagerLivenessService", - "IsExecutionManagerAlive", - ), - ); - self.inner.unary(req, path, codec).await - } - pub async fn get_dead_execution_managers( - &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::codec::ProstCodec::default(); - let path = http::uri::PathAndQuery::from_static( - "/storage.ExecutionManagerLivenessService/GetDeadExecutionManagers", - ); - let mut req = request.into_request(); - req.extensions_mut() - .insert( - GrpcMethod::new( - "storage.ExecutionManagerLivenessService", - "GetDeadExecutionManagers", - ), - ); - self.inner.unary(req, path, codec).await - } } } /// Generated client implementations. @@ -1895,7 +1802,7 @@ pub mod session_management_service_client { } } /// Generated server implementations. -pub mod job_management_service_server { +pub mod job_orchestration_service_server { #![allow( unused_variables, dead_code, @@ -1904,9 +1811,9 @@ pub mod job_management_service_server { clippy::let_unit_value, )] use tonic::codegen::*; - /// Generated trait containing gRPC methods that should be implemented for use with JobManagementServiceServer. + /// Generated trait containing gRPC methods that should be implemented for use with JobOrchestrationServiceServer. #[async_trait] - pub trait JobManagementService: std::marker::Send + std::marker::Sync + 'static { + pub trait JobOrchestrationService: std::marker::Send + std::marker::Sync + 'static { async fn submit_job( &self, request: tonic::Request, @@ -1965,14 +1872,14 @@ pub mod job_management_service_server { >; } #[derive(Debug)] - pub struct JobManagementServiceServer { + pub struct JobOrchestrationServiceServer { inner: Arc, accept_compression_encodings: EnabledCompressionEncodings, send_compression_encodings: EnabledCompressionEncodings, max_decoding_message_size: Option, max_encoding_message_size: Option, } - impl JobManagementServiceServer { + impl JobOrchestrationServiceServer { pub fn new(inner: T) -> Self { Self::from_arc(Arc::new(inner)) } @@ -2024,9 +1931,9 @@ pub mod job_management_service_server { } } impl tonic::codegen::Service> - for JobManagementServiceServer + for JobOrchestrationServiceServer where - T: JobManagementService, + T: JobOrchestrationService, B: Body + std::marker::Send + 'static, B::Error: Into + std::marker::Send + 'static, { @@ -2041,11 +1948,11 @@ pub mod job_management_service_server { } fn call(&mut self, req: http::Request) -> Self::Future { match req.uri().path() { - "/storage.JobManagementService/SubmitJob" => { + "/storage.JobOrchestrationService/SubmitJob" => { #[allow(non_camel_case_types)] - struct SubmitJobSvc(pub Arc); + struct SubmitJobSvc(pub Arc); impl< - T: JobManagementService, + T: JobOrchestrationService, > tonic::server::UnaryService for SubmitJobSvc { type Response = super::SubmitJobResponse; @@ -2059,7 +1966,7 @@ pub mod job_management_service_server { ) -> Self::Future { let inner = Arc::clone(&self.0); let fut = async move { - ::submit_job(&inner, request) + ::submit_job(&inner, request) .await }; Box::pin(fut) @@ -2087,11 +1994,11 @@ pub mod job_management_service_server { }; Box::pin(fut) } - "/storage.JobManagementService/StartJob" => { + "/storage.JobOrchestrationService/StartJob" => { #[allow(non_camel_case_types)] - struct StartJobSvc(pub Arc); + struct StartJobSvc(pub Arc); impl< - T: JobManagementService, + T: JobOrchestrationService, > tonic::server::UnaryService for StartJobSvc { type Response = super::JobStateResponse; @@ -2105,7 +2012,7 @@ pub mod job_management_service_server { ) -> Self::Future { let inner = Arc::clone(&self.0); let fut = async move { - ::start_job(&inner, request) + ::start_job(&inner, request) .await }; Box::pin(fut) @@ -2133,11 +2040,11 @@ pub mod job_management_service_server { }; Box::pin(fut) } - "/storage.JobManagementService/CancelJob" => { + "/storage.JobOrchestrationService/CancelJob" => { #[allow(non_camel_case_types)] - struct CancelJobSvc(pub Arc); + struct CancelJobSvc(pub Arc); impl< - T: JobManagementService, + T: JobOrchestrationService, > tonic::server::UnaryService for CancelJobSvc { type Response = super::JobStateResponse; @@ -2151,7 +2058,7 @@ pub mod job_management_service_server { ) -> Self::Future { let inner = Arc::clone(&self.0); let fut = async move { - ::cancel_job(&inner, request) + ::cancel_job(&inner, request) .await }; Box::pin(fut) @@ -2179,11 +2086,11 @@ pub mod job_management_service_server { }; Box::pin(fut) } - "/storage.JobManagementService/GetJobState" => { + "/storage.JobOrchestrationService/GetJobState" => { #[allow(non_camel_case_types)] - struct GetJobStateSvc(pub Arc); + struct GetJobStateSvc(pub Arc); impl< - T: JobManagementService, + T: JobOrchestrationService, > tonic::server::UnaryService for GetJobStateSvc { type Response = super::JobStateResponse; @@ -2197,7 +2104,10 @@ pub mod job_management_service_server { ) -> Self::Future { let inner = Arc::clone(&self.0); let fut = async move { - ::get_job_state(&inner, request) + ::get_job_state( + &inner, + request, + ) .await }; Box::pin(fut) @@ -2225,11 +2135,11 @@ pub mod job_management_service_server { }; Box::pin(fut) } - "/storage.JobManagementService/GetJobOutputs" => { + "/storage.JobOrchestrationService/GetJobOutputs" => { #[allow(non_camel_case_types)] - struct GetJobOutputsSvc(pub Arc); + struct GetJobOutputsSvc(pub Arc); impl< - T: JobManagementService, + T: JobOrchestrationService, > tonic::server::UnaryService for GetJobOutputsSvc { type Response = super::JobOutputsResponse; @@ -2243,7 +2153,7 @@ pub mod job_management_service_server { ) -> Self::Future { let inner = Arc::clone(&self.0); let fut = async move { - ::get_job_outputs( + ::get_job_outputs( &inner, request, ) @@ -2274,11 +2184,11 @@ pub mod job_management_service_server { }; Box::pin(fut) } - "/storage.JobManagementService/GetJobError" => { + "/storage.JobOrchestrationService/GetJobError" => { #[allow(non_camel_case_types)] - struct GetJobErrorSvc(pub Arc); + struct GetJobErrorSvc(pub Arc); impl< - T: JobManagementService, + T: JobOrchestrationService, > tonic::server::UnaryService for GetJobErrorSvc { type Response = super::JobErrorResponse; @@ -2292,7 +2202,10 @@ pub mod job_management_service_server { ) -> Self::Future { let inner = Arc::clone(&self.0); let fut = async move { - ::get_job_error(&inner, request) + ::get_job_error( + &inner, + request, + ) .await }; Box::pin(fut) @@ -2320,13 +2233,13 @@ pub mod job_management_service_server { }; Box::pin(fut) } - "/storage.JobManagementService/DeleteExpiredTerminatedJobs" => { + "/storage.JobOrchestrationService/DeleteExpiredTerminatedJobs" => { #[allow(non_camel_case_types)] - struct DeleteExpiredTerminatedJobsSvc( + struct DeleteExpiredTerminatedJobsSvc( pub Arc, ); impl< - T: JobManagementService, + T: JobOrchestrationService, > tonic::server::UnaryService< super::DeleteExpiredTerminatedJobsRequest, > for DeleteExpiredTerminatedJobsSvc { @@ -2343,7 +2256,7 @@ pub mod job_management_service_server { ) -> Self::Future { let inner = Arc::clone(&self.0); let fut = async move { - ::delete_expired_terminated_jobs( + ::delete_expired_terminated_jobs( &inner, request, ) @@ -2374,11 +2287,11 @@ pub mod job_management_service_server { }; Box::pin(fut) } - "/storage.JobManagementService/ResendReadyTasks" => { + "/storage.JobOrchestrationService/ResendReadyTasks" => { #[allow(non_camel_case_types)] - struct ResendReadyTasksSvc(pub Arc); + struct ResendReadyTasksSvc(pub Arc); impl< - T: JobManagementService, + T: JobOrchestrationService, > tonic::server::UnaryService for ResendReadyTasksSvc { type Response = super::JobManagementOperationResponse; @@ -2392,7 +2305,7 @@ pub mod job_management_service_server { ) -> Self::Future { let inner = Arc::clone(&self.0); let fut = async move { - ::resend_ready_tasks( + ::resend_ready_tasks( &inner, request, ) @@ -2443,7 +2356,7 @@ pub mod job_management_service_server { } } } - impl Clone for JobManagementServiceServer { + impl Clone for JobOrchestrationServiceServer { fn clone(&self) -> Self { let inner = self.inner.clone(); Self { @@ -2456,8 +2369,8 @@ pub mod job_management_service_server { } } /// Generated gRPC service name - pub const SERVICE_NAME: &str = "storage.JobManagementService"; - impl tonic::server::NamedService for JobManagementServiceServer { + pub const SERVICE_NAME: &str = "storage.JobOrchestrationService"; + impl tonic::server::NamedService for JobOrchestrationServiceServer { const NAME: &'static str = SERVICE_NAME; } } @@ -2765,7 +2678,7 @@ pub mod task_instance_management_service_server { } } /// Generated server implementations. -pub mod scheduler_storage_service_server { +pub mod inbound_queue_service_server { #![allow( unused_variables, dead_code, @@ -2774,9 +2687,9 @@ pub mod scheduler_storage_service_server { clippy::let_unit_value, )] use tonic::codegen::*; - /// Generated trait containing gRPC methods that should be implemented for use with SchedulerStorageServiceServer. + /// Generated trait containing gRPC methods that should be implemented for use with InboundQueueServiceServer. #[async_trait] - pub trait SchedulerStorageService: std::marker::Send + std::marker::Sync + 'static { + pub trait InboundQueueService: std::marker::Send + std::marker::Sync + 'static { async fn poll_ready_tasks( &self, request: tonic::Request, @@ -2800,14 +2713,14 @@ pub mod scheduler_storage_service_server { >; } #[derive(Debug)] - pub struct SchedulerStorageServiceServer { + pub struct InboundQueueServiceServer { inner: Arc, accept_compression_encodings: EnabledCompressionEncodings, send_compression_encodings: EnabledCompressionEncodings, max_decoding_message_size: Option, max_encoding_message_size: Option, } - impl SchedulerStorageServiceServer { + impl InboundQueueServiceServer { pub fn new(inner: T) -> Self { Self::from_arc(Arc::new(inner)) } @@ -2858,10 +2771,9 @@ pub mod scheduler_storage_service_server { self } } - impl tonic::codegen::Service> - for SchedulerStorageServiceServer + impl tonic::codegen::Service> for InboundQueueServiceServer where - T: SchedulerStorageService, + T: InboundQueueService, B: Body + std::marker::Send + 'static, B::Error: Into + std::marker::Send + 'static, { @@ -2876,11 +2788,11 @@ pub mod scheduler_storage_service_server { } fn call(&mut self, req: http::Request) -> Self::Future { match req.uri().path() { - "/storage.SchedulerStorageService/PollReadyTasks" => { + "/storage.InboundQueueService/PollReadyTasks" => { #[allow(non_camel_case_types)] - struct PollReadyTasksSvc(pub Arc); + struct PollReadyTasksSvc(pub Arc); impl< - T: SchedulerStorageService, + T: InboundQueueService, > tonic::server::UnaryService for PollReadyTasksSvc { type Response = super::PollReadyTasksResponse; @@ -2894,7 +2806,7 @@ pub mod scheduler_storage_service_server { ) -> Self::Future { let inner = Arc::clone(&self.0); let fut = async move { - ::poll_ready_tasks( + ::poll_ready_tasks( &inner, request, ) @@ -2925,13 +2837,11 @@ pub mod scheduler_storage_service_server { }; Box::pin(fut) } - "/storage.SchedulerStorageService/PollReadyCommitTasks" => { + "/storage.InboundQueueService/PollReadyCommitTasks" => { #[allow(non_camel_case_types)] - struct PollReadyCommitTasksSvc( - pub Arc, - ); + struct PollReadyCommitTasksSvc(pub Arc); impl< - T: SchedulerStorageService, + T: InboundQueueService, > tonic::server::UnaryService for PollReadyCommitTasksSvc { type Response = super::PollReadyTasksResponse; @@ -2945,7 +2855,7 @@ pub mod scheduler_storage_service_server { ) -> Self::Future { let inner = Arc::clone(&self.0); let fut = async move { - ::poll_ready_commit_tasks( + ::poll_ready_commit_tasks( &inner, request, ) @@ -2976,13 +2886,11 @@ pub mod scheduler_storage_service_server { }; Box::pin(fut) } - "/storage.SchedulerStorageService/PollReadyCleanupTasks" => { + "/storage.InboundQueueService/PollReadyCleanupTasks" => { #[allow(non_camel_case_types)] - struct PollReadyCleanupTasksSvc( - pub Arc, - ); + struct PollReadyCleanupTasksSvc(pub Arc); impl< - T: SchedulerStorageService, + T: InboundQueueService, > tonic::server::UnaryService for PollReadyCleanupTasksSvc { type Response = super::PollReadyTasksResponse; @@ -2996,7 +2904,7 @@ pub mod scheduler_storage_service_server { ) -> Self::Future { let inner = Arc::clone(&self.0); let fut = async move { - ::poll_ready_cleanup_tasks( + ::poll_ready_cleanup_tasks( &inner, request, ) @@ -3047,7 +2955,7 @@ pub mod scheduler_storage_service_server { } } } - impl Clone for SchedulerStorageServiceServer { + impl Clone for InboundQueueServiceServer { fn clone(&self) -> Self { let inner = self.inner.clone(); Self { @@ -3060,8 +2968,8 @@ pub mod scheduler_storage_service_server { } } /// Generated gRPC service name - pub const SERVICE_NAME: &str = "storage.SchedulerStorageService"; - impl tonic::server::NamedService for SchedulerStorageServiceServer { + pub const SERVICE_NAME: &str = "storage.InboundQueueService"; + impl tonic::server::NamedService for InboundQueueServiceServer { const NAME: &'static str = SERVICE_NAME; } } @@ -3395,20 +3303,6 @@ pub mod execution_manager_liveness_service_server { tonic::Response, tonic::Status, >; - async fn is_execution_manager_alive( - &self, - request: tonic::Request, - ) -> std::result::Result< - tonic::Response, - tonic::Status, - >; - async fn get_dead_execution_managers( - &self, - request: tonic::Request, - ) -> std::result::Result< - tonic::Response, - tonic::Status, - >; } #[derive(Debug)] pub struct ExecutionManagerLivenessServiceServer { @@ -3595,114 +3489,6 @@ pub mod execution_manager_liveness_service_server { }; Box::pin(fut) } - "/storage.ExecutionManagerLivenessService/IsExecutionManagerAlive" => { - #[allow(non_camel_case_types)] - struct IsExecutionManagerAliveSvc< - T: ExecutionManagerLivenessService, - >( - pub Arc, - ); - impl< - T: ExecutionManagerLivenessService, - > tonic::server::UnaryService - for IsExecutionManagerAliveSvc { - type Response = super::IsExecutionManagerAliveResponse; - 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 { - ::is_execution_manager_alive( - &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 = IsExecutionManagerAliveSvc(inner); - let codec = tonic::codec::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) - } - "/storage.ExecutionManagerLivenessService/GetDeadExecutionManagers" => { - #[allow(non_camel_case_types)] - struct GetDeadExecutionManagersSvc< - T: ExecutionManagerLivenessService, - >( - pub Arc, - ); - impl< - T: ExecutionManagerLivenessService, - > tonic::server::UnaryService - for GetDeadExecutionManagersSvc { - type Response = super::GetDeadExecutionManagersResponse; - type Future = BoxFuture< - tonic::Response, - tonic::Status, - >; - fn call( - &mut self, - request: tonic::Request< - super::GetDeadExecutionManagersRequest, - >, - ) -> Self::Future { - let inner = Arc::clone(&self.0); - let fut = async move { - ::get_dead_execution_managers( - &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 = GetDeadExecutionManagersSvc(inner); - let codec = tonic::codec::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(empty_body()); diff --git a/components/spider-proto/storage/storage.proto b/components/spider-proto/storage/storage.proto index 7de2defd..d28048c5 100644 --- a/components/spider-proto/storage/storage.proto +++ b/components/spider-proto/storage/storage.proto @@ -2,7 +2,7 @@ syntax = "proto3"; package storage; -service JobManagementService { +service JobOrchestrationService { rpc SubmitJob(SubmitJobRequest) returns (SubmitJobResponse); rpc StartJob(JobIdRequest) returns (JobStateResponse); rpc CancelJob(JobIdRequest) returns (JobStateResponse); @@ -20,7 +20,7 @@ service TaskInstanceManagementService { rpc ReportTaskFailure(ReportTaskFailureRequest) returns (TaskInstanceOperationResponse); } -service SchedulerStorageService { +service InboundQueueService { rpc PollReadyTasks(PollReadyTasksRequest) returns (PollReadyTasksResponse); rpc PollReadyCommitTasks(PollReadyTasksRequest) returns (PollReadyTasksResponse); rpc PollReadyCleanupTasks(PollReadyTasksRequest) returns (PollReadyTasksResponse); @@ -37,10 +37,6 @@ service ExecutionManagerLivenessService { returns (RegisterExecutionManagerResponse); rpc UpdateExecutionManagerHeartbeat(ExecutionManagerIdRequest) returns (UpdateExecutionManagerHeartbeatResponse); - rpc IsExecutionManagerAlive(ExecutionManagerIdRequest) - returns (IsExecutionManagerAliveResponse); - rpc GetDeadExecutionManagers(GetDeadExecutionManagersRequest) - returns (GetDeadExecutionManagersResponse); } service SessionManagementService { @@ -50,14 +46,14 @@ service SessionManagementService { message SubmitJobRequest { uint64 resource_group_id = 1; bytes serialized_task_graph = 2; - repeated bytes serialized_inputs = 3; + bytes serialized_inputs = 3; uint64 session_id = 4; } message SubmitJobResponse { oneof result { uint64 job_id = 1; - JobManagementError error = 2; + JobOrchestrationError error = 2; } } @@ -69,14 +65,14 @@ message JobIdRequest { message JobStateResponse { oneof result { JobState state = 1; - JobManagementError error = 2; + JobOrchestrationError error = 2; } } message JobOutputsResponse { oneof result { JobOutputs outputs = 1; - JobManagementError error = 2; + JobOrchestrationError error = 2; } } @@ -87,7 +83,7 @@ message JobOutputs { message JobErrorResponse { oneof result { string error_message = 1; - JobManagementError error = 2; + JobOrchestrationError error = 2; } } @@ -99,7 +95,7 @@ message DeleteExpiredTerminatedJobsRequest { message DeleteExpiredTerminatedJobsResponse { oneof result { DeletedJobs deleted_jobs = 1; - JobManagementError error = 2; + JobOrchestrationError error = 2; } } @@ -120,7 +116,7 @@ message PollReadyTasksRequest { message PollReadyTasksResponse { oneof result { ReadyTasks tasks = 1; - SchedulerStorageError error = 2; + InboundQueueResponseError error = 2; } } @@ -145,7 +141,7 @@ message RegisterTaskInstanceRequest { message RegisterTaskInstanceResponse { oneof result { bytes execution_context = 1; - TaskInstanceError error = 2; + TaskInstanceManagementError error = 2; } } @@ -179,7 +175,7 @@ message ResourceGroupIdRequest { message ResourceGroupIdResponse { oneof result { uint64 resource_group_id = 1; - ResourceGroupError error = 2; + ResourceGroupManagementError error = 2; } } @@ -216,29 +212,6 @@ message UpdateExecutionManagerHeartbeatResponse { } } -message IsExecutionManagerAliveResponse { - oneof result { - bool alive = 1; - ExecutionManagerLivenessError error = 2; - } -} - -message GetDeadExecutionManagersRequest { - uint64 stale_after_sec = 1; - uint64 session_id = 2; -} - -message GetDeadExecutionManagersResponse { - oneof result { - DeadExecutionManagers dead_execution_managers = 1; - ExecutionManagerLivenessError error = 2; - } -} - -message DeadExecutionManagers { - repeated uint64 execution_manager_ids = 1; -} - message GetSessionResponse { uint64 session_id = 1; } @@ -265,27 +238,27 @@ enum JobState { message JobManagementOperationResponse { oneof result { Void ok = 1; - JobManagementError error = 2; + JobOrchestrationError error = 2; } } message TaskInstanceOperationResponse { oneof result { Void ok = 1; - TaskInstanceError error = 2; + TaskInstanceManagementError error = 2; } } message ResourceGroupOperationResponse { oneof result { Void ok = 1; - ResourceGroupError error = 2; + ResourceGroupManagementError error = 2; } } message Void {} -message JobManagementError { +message JobOrchestrationError { enum ErrCode { ERR_CODE_UNSPECIFIED = 0; STALE_SESSION = 1; @@ -299,7 +272,7 @@ message JobManagementError { uint64 storage_session = 3; } -message TaskInstanceError { +message TaskInstanceManagementError { enum ErrCode { ERR_CODE_UNSPECIFIED = 0; STALE_SESSION = 1; @@ -313,7 +286,7 @@ message TaskInstanceError { uint64 storage_session = 3; } -message SchedulerStorageError { +message InboundQueueResponseError { enum ErrCode { ERR_CODE_UNSPECIFIED = 0; INBOUND_CLOSED = 1; @@ -325,7 +298,7 @@ message SchedulerStorageError { string message = 2; } -message ResourceGroupError { +message ResourceGroupManagementError { enum ErrCode { ERR_CODE_UNSPECIFIED = 0; STALE_SESSION = 1; diff --git a/components/spider-scheduler/src/grpc/storage_client.rs b/components/spider-scheduler/src/grpc/storage_client.rs index a3fdc5be..5d898f8e 100644 --- a/components/spider-scheduler/src/grpc/storage_client.rs +++ b/components/spider-scheduler/src/grpc/storage_client.rs @@ -10,12 +10,12 @@ use spider_core::{ }; use spider_proto_rust::storage::{ self, - job_management_error, - job_management_service_client::JobManagementServiceClient, + inbound_queue_response_error, + inbound_queue_service_client::InboundQueueServiceClient, + job_orchestration_error, + job_orchestration_service_client::JobOrchestrationServiceClient, job_state_response, poll_ready_tasks_response, - scheduler_storage_error, - scheduler_storage_service_client::SchedulerStorageServiceClient, session_management_service_client::SessionManagementServiceClient, }; use tonic::transport::{Channel, Endpoint}; @@ -29,8 +29,8 @@ use crate::{ /// gRPC-backed [`SchedulerStorageClient`] implementation. #[derive(Debug, Clone)] pub struct GrpcSchedulerStorageClient { - scheduler_client: SchedulerStorageServiceClient, - job_client: JobManagementServiceClient, + scheduler_client: InboundQueueServiceClient, + job_client: JobOrchestrationServiceClient, session_tracker: SessionTracker, } @@ -52,8 +52,8 @@ impl GrpcSchedulerStorageClient { let session_id = Self::get_initial_session(channel.clone()).await?; Ok(Self { - scheduler_client: SchedulerStorageServiceClient::new(channel.clone()), - job_client: JobManagementServiceClient::new(channel), + scheduler_client: InboundQueueServiceClient::new(channel.clone()), + job_client: JobOrchestrationServiceClient::new(channel), session_tracker: SessionTracker::new(session_id), }) } @@ -144,9 +144,9 @@ impl SchedulerStorageClient for GrpcSchedulerStorageClient { } } -impl From for StorageClientError { - fn from(error: storage::SchedulerStorageError) -> Self { - scheduler_storage_error_to_client_error(error) +impl From for StorageClientError { + fn from(error: storage::InboundQueueResponseError) -> Self { + inbound_queue_response_error_to_client_error(error) } } @@ -254,7 +254,7 @@ fn job_state_response_to_result( .map_err(|error| StorageClientError::Transport(error.to_string())) .and_then(|state| JobState::try_from(state).map_err(StorageClientError::Transport)), Some(job_state_response::Result::Error(error)) => { - Err(job_management_error_to_client_error(error, job_id)) + Err(job_orchestration_error_to_client_error(error, job_id)) } None => Err(StorageClientError::Transport( "job state response missing `result` message".to_owned(), @@ -264,18 +264,20 @@ fn job_state_response_to_result( /// # Returns /// -/// [`storage::SchedulerStorageError`] converted into [`StorageClientError`]. -fn scheduler_storage_error_to_client_error( - error: storage::SchedulerStorageError, +/// [`storage::InboundQueueResponseError`] converted into [`StorageClientError`]. +fn inbound_queue_response_error_to_client_error( + error: storage::InboundQueueResponseError, ) -> StorageClientError { - match scheduler_storage_error::ErrCode::try_from(error.err_code) { - Ok(scheduler_storage_error::ErrCode::InboundClosed) => StorageClientError::InboundClosed, - Ok(scheduler_storage_error::ErrCode::InvalidInput) => { + match inbound_queue_response_error::ErrCode::try_from(error.err_code) { + Ok(inbound_queue_response_error::ErrCode::InboundClosed) => { + StorageClientError::InboundClosed + } + Ok(inbound_queue_response_error::ErrCode::InvalidInput) => { StorageClientError::InvalidInput(error.message) } Ok( - scheduler_storage_error::ErrCode::Server - | scheduler_storage_error::ErrCode::Unspecified, + inbound_queue_response_error::ErrCode::Server + | inbound_queue_response_error::ErrCode::Unspecified, ) => StorageClientError::Server(error.message), Err(error) => { StorageClientError::Transport(format!("unknown scheduler storage error kind: {error}")) @@ -285,24 +287,25 @@ fn scheduler_storage_error_to_client_error( /// # Returns /// -/// [`storage::JobManagementError`] converted into [`StorageClientError`]. -fn job_management_error_to_client_error( - error: storage::JobManagementError, +/// [`storage::JobOrchestrationError`] converted into [`StorageClientError`]. +fn job_orchestration_error_to_client_error( + error: storage::JobOrchestrationError, requested_job_id: JobId, ) -> StorageClientError { - match job_management_error::ErrCode::try_from(error.err_code) { - Ok(job_management_error::ErrCode::JobNotFound) => { + match job_orchestration_error::ErrCode::try_from(error.err_code) { + Ok(job_orchestration_error::ErrCode::JobNotFound) => { StorageClientError::JobNotFound(requested_job_id) } - Ok(job_management_error::ErrCode::StaleSession) => StorageClientError::StaleSession { + Ok(job_orchestration_error::ErrCode::StaleSession) => StorageClientError::StaleSession { storage_session: error.storage_session, }, - Ok(job_management_error::ErrCode::InvalidInput) => { + Ok(job_orchestration_error::ErrCode::InvalidInput) => { StorageClientError::InvalidInput(error.message) } - Ok(job_management_error::ErrCode::Server | job_management_error::ErrCode::Unspecified) => { - StorageClientError::Server(error.message) - } + Ok( + job_orchestration_error::ErrCode::Server + | job_orchestration_error::ErrCode::Unspecified, + ) => StorageClientError::Server(error.message), Err(error) => { StorageClientError::Transport(format!("unknown job management error kind: {error}")) } @@ -321,7 +324,11 @@ fn to_transport_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, poll_ready_tasks_response, scheduler_storage_error}; + use spider_proto_rust::storage::{ + self, + inbound_queue_response_error, + poll_ready_tasks_response, + }; use super::*; @@ -378,9 +385,9 @@ mod tests { } #[test] - fn scheduler_storage_error_maps_inbound_closed() { - let error = storage::SchedulerStorageError { - err_code: scheduler_storage_error::ErrCode::InboundClosed.into(), + fn inbound_queue_response_error_maps_inbound_closed() { + let error = storage::InboundQueueResponseError { + err_code: inbound_queue_response_error::ErrCode::InboundClosed.into(), message: "closed".to_owned(), }; From 89fac75cbfe0783f2a12575aa6fcff11b96e6a10 Mon Sep 17 00:00:00 2001 From: LinZhihao-723 Date: Fri, 12 Jun 2026 15:51:17 -0400 Subject: [PATCH 13/16] Review and polish . --- Cargo.lock | 1 + components/spider-proto-rust/Cargo.toml | 1 + components/spider-proto-rust/src/error.rs | 17 +++++++++++++++++ components/spider-proto-rust/src/id.rs | 17 ++++++++++------- components/spider-proto-rust/src/job.rs | 10 +++++----- components/spider-proto-rust/src/lib.rs | 1 + .../spider-scheduler/src/grpc/storage_client.rs | 10 ++++++++-- 7 files changed, 43 insertions(+), 14 deletions(-) create mode 100644 components/spider-proto-rust/src/error.rs diff --git a/Cargo.lock b/Cargo.lock index 9b5cbe34..a6fbc043 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1859,6 +1859,7 @@ version = "0.1.0" dependencies = [ "prost", "spider-core", + "thiserror", "tonic", "tonic-build", ] diff --git a/components/spider-proto-rust/Cargo.toml b/components/spider-proto-rust/Cargo.toml index 6a5e53db..dc5ec267 100644 --- a/components/spider-proto-rust/Cargo.toml +++ b/components/spider-proto-rust/Cargo.toml @@ -10,6 +10,7 @@ path = "src/lib.rs" [dependencies] prost = "0.13.5" spider-core = { path = "../spider-core" } +thiserror = "2.0.18" tonic = "0.12.3" [build-dependencies] diff --git a/components/spider-proto-rust/src/error.rs b/components/spider-proto-rust/src/error.rs new file mode 100644 index 00000000..f9014e7b --- /dev/null +++ b/components/spider-proto-rust/src/error.rs @@ -0,0 +1,17 @@ +//! Error types for converting protobuf wire values into Spider core types. + +/// Errors produced when converting a protobuf message into its Spider core representation. +#[derive(Debug, thiserror::Error)] +pub enum Error { + /// A protobuf task index could not be represented as a [`usize`]. + #[error("task index does not fit in `usize`: {0}")] + TaskIndexOutOfRange(u64), + + /// A protobuf [`crate::storage::TaskId`] carried no `kind`. + #[error("task id missing kind")] + TaskIdKindMissing, + + /// A protobuf [`crate::storage::JobState`] was left unspecified. + #[error("job state is unspecified")] + JobStateUnspecified, +} diff --git a/components/spider-proto-rust/src/id.rs b/components/spider-proto-rust/src/id.rs index 1da62ca5..bd21a55a 100644 --- a/components/spider-proto-rust/src/id.rs +++ b/components/spider-proto-rust/src/id.rs @@ -2,7 +2,10 @@ use spider_core::types::id::TaskId; -use crate::storage::{self, task_id}; +use crate::{ + error::Error, + storage::{self, task_id}, +}; impl From for storage::TaskId { fn from(task_id: TaskId) -> Self { @@ -18,16 +21,16 @@ impl From for storage::TaskId { } impl TryFrom for TaskId { - type Error = String; + type Error = Error; fn try_from(task_id: storage::TaskId) -> Result { match task_id.kind { - Some(task_id::Kind::Index(task_index)) => usize::try_from(task_index) - .map(TaskId::Index) - .map_err(|error| format!("task index does not fit in usize: {error}")), + Some(task_id::Kind::Index(task_index)) => Ok(Self::Index( + usize::try_from(task_index).map_err(|_| Error::TaskIndexOutOfRange(task_index))?, + )), Some(task_id::Kind::Commit(_)) => Ok(Self::Commit), Some(task_id::Kind::Cleanup(_)) => Ok(Self::Cleanup), - None => Err("task id missing kind".to_owned()), + None => Err(Error::TaskIdKindMissing), } } } @@ -72,6 +75,6 @@ mod tests { let error = TaskId::try_from(storage::TaskId { kind: None }) .expect_err("missing task id kind should fail"); - assert!(error.contains("missing kind")); + assert!(matches!(error, Error::TaskIdKindMissing)); } } diff --git a/components/spider-proto-rust/src/job.rs b/components/spider-proto-rust/src/job.rs index f8d62494..93d167cb 100644 --- a/components/spider-proto-rust/src/job.rs +++ b/components/spider-proto-rust/src/job.rs @@ -2,7 +2,7 @@ use spider_core::job::JobState; -use crate::storage; +use crate::{error::Error, storage}; impl From for storage::JobState { fn from(state: JobState) -> Self { @@ -19,11 +19,11 @@ impl From for storage::JobState { } impl TryFrom for JobState { - type Error = String; + type Error = Error; fn try_from(state: storage::JobState) -> Result { match state { - storage::JobState::Unspecified => Err("job state is unspecified".to_owned()), + storage::JobState::Unspecified => Err(Error::JobStateUnspecified), storage::JobState::Ready => Ok(Self::Ready), storage::JobState::Running => Ok(Self::Running), storage::JobState::CommitReady => Ok(Self::CommitReady), @@ -39,7 +39,7 @@ impl TryFrom for JobState { mod tests { use spider_core::job::JobState; - use crate::storage::JobState as ProtocolJobState; + use crate::{error::Error, storage::JobState as ProtocolJobState}; #[test] fn job_state_to_protocol_converts_succeeded() { @@ -61,6 +61,6 @@ mod tests { let error = JobState::try_from(ProtocolJobState::Unspecified) .expect_err("unspecified job state should fail"); - assert!(error.contains("unspecified")); + assert!(matches!(error, Error::JobStateUnspecified)); } } diff --git a/components/spider-proto-rust/src/lib.rs b/components/spider-proto-rust/src/lib.rs index 06f1c0ba..9dcff1c9 100644 --- a/components/spider-proto-rust/src/lib.rs +++ b/components/spider-proto-rust/src/lib.rs @@ -1,5 +1,6 @@ //! Rust gRPC protocol definitions generated from Spider protobuf files. +pub mod error; pub mod id; pub mod job; diff --git a/components/spider-scheduler/src/grpc/storage_client.rs b/components/spider-scheduler/src/grpc/storage_client.rs index 5d898f8e..c228ff79 100644 --- a/components/spider-scheduler/src/grpc/storage_client.rs +++ b/components/spider-scheduler/src/grpc/storage_client.rs @@ -227,7 +227,10 @@ fn ready_task_to_inbound_entry( let task_id = task .task_id .ok_or_else(|| StorageClientError::Transport("ready task missing task ID".to_owned())) - .and_then(|task_id| TaskId::try_from(task_id).map_err(StorageClientError::Transport))?; + .and_then(|task_id| { + TaskId::try_from(task_id) + .map_err(|error| StorageClientError::Transport(error.to_string())) + })?; Ok(InboundEntry { resource_group_id: ResourceGroupId::from(task.resource_group_id), job_id: JobId::from(task.job_id), @@ -252,7 +255,10 @@ fn job_state_response_to_result( match response.result { Some(job_state_response::Result::State(state)) => storage::JobState::try_from(state) .map_err(|error| StorageClientError::Transport(error.to_string())) - .and_then(|state| JobState::try_from(state).map_err(StorageClientError::Transport)), + .and_then(|state| { + JobState::try_from(state) + .map_err(|error| StorageClientError::Transport(error.to_string())) + }), Some(job_state_response::Result::Error(error)) => { Err(job_orchestration_error_to_client_error(error, job_id)) } From a503e10d10d704b7056b87c1baafda251881f575 Mon Sep 17 00:00:00 2001 From: LinZhihao-723 Date: Fri, 12 Jun 2026 16:15:18 -0400 Subject: [PATCH 14/16] Review spider-execution-manager's client implementations. --- .../src/client/grpc/liveness.rs | 78 +++++++++---------- .../src/client/grpc/storage.rs | 2 +- 2 files changed, 40 insertions(+), 40 deletions(-) diff --git a/components/spider-execution-manager/src/client/grpc/liveness.rs b/components/spider-execution-manager/src/client/grpc/liveness.rs index 8292cef8..af5c0fb1 100644 --- a/components/spider-execution-manager/src/client/grpc/liveness.rs +++ b/components/spider-execution-manager/src/client/grpc/liveness.rs @@ -36,7 +36,7 @@ impl GrpcLivenessClient { pub async fn connect(endpoint: Endpoint) -> Result { ExecutionManagerLivenessServiceClient::connect(endpoint) .await - .map(|inner| Self { client: inner }) + .map(|client| Self { client }) .map_err(to_transport_error) } } @@ -77,6 +77,24 @@ impl LivenessClient for GrpcLivenessClient { } } +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}" + )), + } + } +} + /// # Returns /// /// [`storage::RegisterExecutionManagerResponse`] converted into @@ -91,9 +109,7 @@ fn register_response_to_result( session_id: registration.session_id, }) } - Some(register_execution_manager_response::Result::Error(error)) => { - Err(storage_error_to_liveness_error(error)) - } + Some(register_execution_manager_response::Result::Error(error)) => Err(error.into()), None => Err(LivenessResponseError::Transport( "register execution manager response missing result".to_owned(), )), @@ -112,7 +128,7 @@ fn heartbeat_response_to_result( Ok(session_id) } Some(update_execution_manager_heartbeat_response::Result::Error(error)) => { - Err(storage_error_to_liveness_error(error)) + Err(error.into()) } None => Err(LivenessResponseError::Transport( "update execution manager heartbeat response missing result".to_owned(), @@ -120,31 +136,6 @@ fn heartbeat_response_to_result( } } -/// Converts a protobuf storage error into a liveness client error. -/// -/// # Returns -/// -/// The corresponding [`LivenessResponseError`]. -fn storage_error_to_liveness_error( - error: storage::ExecutionManagerLivenessError, -) -> LivenessResponseError { - match execution_manager_liveness_error::ErrCode::try_from(error.err_code) { - Ok(execution_manager_liveness_error::ErrCode::MarkedDead) => { - LivenessResponseError::MarkedDead - } - Ok(execution_manager_liveness_error::ErrCode::InvalidInput) => { - LivenessResponseError::IllegalId(error.message) - } - Ok( - execution_manager_liveness_error::ErrCode::Server - | execution_manager_liveness_error::ErrCode::Unspecified, - ) => LivenessResponseError::Transport(error.message), - Err(error) => LivenessResponseError::Transport(format!( - "unknown execution manager liveness error kind: {error}" - )), - } -} - /// Converts a displayable transport-layer error into [`LivenessResponseError::Transport`]. /// /// # Returns @@ -163,11 +154,14 @@ mod tests { #[test] fn register_response_to_result_returns_registration() { + const SESSION_ID: SessionId = 7; + const EM_ID: ExecutionManagerId = ExecutionManagerId::from(5); + let response = storage::RegisterExecutionManagerResponse { result: Some(register_execution_manager_response::Result::Registration( storage::ExecutionManagerRegistration { - execution_manager_id: 5, - session_id: 7, + execution_manager_id: EM_ID.get(), + session_id: SESSION_ID, }, )), }; @@ -178,34 +172,40 @@ mod tests { assert_eq!( registration, RegistrationResponse { - em_id: ExecutionManagerId::from(5), - session_id: 7, + em_id: EM_ID, + session_id: SESSION_ID, } ); } #[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(9)), + result: Some( + update_execution_manager_heartbeat_response::Result::SessionId(SESSION_ID), + ), }; let session_id = heartbeat_response_to_result(response) .expect("heartbeat response conversion should succeed"); - assert_eq!(session_id, 9); + assert_eq!(session_id, SESSION_ID); } #[test] fn liveness_storage_error_maps_invalid_input_to_illegal_id() { + const ERROR_MSG: &str = "bad em id"; + let error = storage::ExecutionManagerLivenessError { err_code: execution_manager_liveness_error::ErrCode::InvalidInput.into(), - message: "bad em id".to_owned(), + message: ERROR_MSG.to_owned(), }; - match storage_error_to_liveness_error(error) { + match LivenessResponseError::from(error) { LivenessResponseError::IllegalId(message) => { - assert_eq!(message, "bad em id"); + assert_eq!(message, ERROR_MSG); } error => panic!("unexpected liveness response error: {error:?}"), } diff --git a/components/spider-execution-manager/src/client/grpc/storage.rs b/components/spider-execution-manager/src/client/grpc/storage.rs index 1a7df78a..378de5d4 100644 --- a/components/spider-execution-manager/src/client/grpc/storage.rs +++ b/components/spider-execution-manager/src/client/grpc/storage.rs @@ -40,7 +40,7 @@ impl GrpcStorageClient { pub async fn connect(endpoint: Endpoint) -> Result { TaskInstanceManagementServiceClient::connect(endpoint) .await - .map(|inner| Self { client: inner }) + .map(|client| Self { client }) .map_err(to_transport_error) } } From 27a8b28e036561d84e23d7acabc026a1c1e41960 Mon Sep 17 00:00:00 2001 From: LinZhihao-723 Date: Fri, 12 Jun 2026 17:19:39 -0400 Subject: [PATCH 15/16] Review . --- .../src/generated/storage.rs | 4 +- components/spider-proto/storage/storage.proto | 3 +- components/spider-scheduler/src/grpc/mod.rs | 5 - components/spider-scheduler/src/lib.rs | 4 +- .../grpc.rs} | 139 +++++++----------- .../mod.rs} | 4 + 6 files changed, 57 insertions(+), 102 deletions(-) delete mode 100644 components/spider-scheduler/src/grpc/mod.rs rename components/spider-scheduler/src/{grpc/storage_client.rs => storage_client/grpc.rs} (74%) rename components/spider-scheduler/src/{storage_client.rs => storage_client/mod.rs} (98%) diff --git a/components/spider-proto-rust/src/generated/storage.rs b/components/spider-proto-rust/src/generated/storage.rs index 4323ce13..970dc141 100644 --- a/components/spider-proto-rust/src/generated/storage.rs +++ b/components/spider-proto-rust/src/generated/storage.rs @@ -29,8 +29,6 @@ pub mod submit_job_response { pub struct JobIdRequest { #[prost(uint64, tag = "1")] pub job_id: u64, - #[prost(uint64, tag = "2")] - pub session_id: u64, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct JobStateResponse { @@ -121,7 +119,7 @@ pub struct PollReadyTasksRequest { #[prost(uint64, tag = "1")] pub max_items: u64, #[prost(uint64, tag = "2")] - pub wait_ns: u64, + pub wait_ms: u64, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct PollReadyTasksResponse { diff --git a/components/spider-proto/storage/storage.proto b/components/spider-proto/storage/storage.proto index d28048c5..df771fb7 100644 --- a/components/spider-proto/storage/storage.proto +++ b/components/spider-proto/storage/storage.proto @@ -59,7 +59,6 @@ message SubmitJobResponse { message JobIdRequest { uint64 job_id = 1; - uint64 session_id = 2; } message JobStateResponse { @@ -110,7 +109,7 @@ message ResendReadyTasksRequest { message PollReadyTasksRequest { uint64 max_items = 1; - uint64 wait_ns = 2; + uint64 wait_ms = 2; } message PollReadyTasksResponse { diff --git a/components/spider-scheduler/src/grpc/mod.rs b/components/spider-scheduler/src/grpc/mod.rs deleted file mode 100644 index 84954f46..00000000 --- a/components/spider-scheduler/src/grpc/mod.rs +++ /dev/null @@ -1,5 +0,0 @@ -//! gRPC-backed scheduler client implementations. - -mod storage_client; - -pub use storage_client::GrpcSchedulerStorageClient; diff --git a/components/spider-scheduler/src/lib.rs b/components/spider-scheduler/src/lib.rs index 00037216..9a16cd97 100644 --- a/components/spider-scheduler/src/lib.rs +++ b/components/spider-scheduler/src/lib.rs @@ -34,7 +34,6 @@ pub mod core; pub mod dispatch_queue; pub mod error; -pub mod grpc; pub mod storage_client; pub mod types; @@ -42,7 +41,6 @@ pub use crate::{ core::SchedulerCore, dispatch_queue::{DispatchQueueSink, DispatchQueueSource}, error::{SchedulerError, StorageClientError}, - grpc::GrpcSchedulerStorageClient, - storage_client::SchedulerStorageClient, + storage_client::{GrpcSchedulerStorageClient, SchedulerStorageClient}, types::{InboundEntry, TaskAssignment}, }; diff --git a/components/spider-scheduler/src/grpc/storage_client.rs b/components/spider-scheduler/src/storage_client/grpc.rs similarity index 74% rename from components/spider-scheduler/src/grpc/storage_client.rs rename to components/spider-scheduler/src/storage_client/grpc.rs index c228ff79..4ee3f464 100644 --- a/components/spider-scheduler/src/grpc/storage_client.rs +++ b/components/spider-scheduler/src/storage_client/grpc.rs @@ -5,7 +5,6 @@ use std::time::Duration; use async_trait::async_trait; use spider_core::{ job::JobState, - session::SessionTracker, types::id::{JobId, ResourceGroupId, SessionId, TaskId}, }; use spider_proto_rust::storage::{ @@ -16,7 +15,6 @@ use spider_proto_rust::storage::{ job_orchestration_service_client::JobOrchestrationServiceClient, job_state_response, poll_ready_tasks_response, - session_management_service_client::SessionManagementServiceClient, }; use tonic::transport::{Channel, Endpoint}; @@ -31,7 +29,6 @@ use crate::{ pub struct GrpcSchedulerStorageClient { scheduler_client: InboundQueueServiceClient, job_client: JobOrchestrationServiceClient, - session_tracker: SessionTracker, } impl GrpcSchedulerStorageClient { @@ -46,36 +43,14 @@ impl GrpcSchedulerStorageClient { /// Returns an error if: /// /// * [`StorageClientError::Transport`] if tonic cannot create or connect to the endpoint. - /// * Forwards [`Self::get_initial_session`]'s return values on failure. pub async fn connect(endpoint: Endpoint) -> Result { let channel = endpoint.connect().await.map_err(to_transport_error)?; - let session_id = Self::get_initial_session(channel.clone()).await?; Ok(Self { scheduler_client: InboundQueueServiceClient::new(channel.clone()), job_client: JobOrchestrationServiceClient::new(channel), - session_tracker: SessionTracker::new(session_id), }) } - - /// Fetches storage's current session ID. - /// - /// # Returns - /// - /// The current storage session ID on success. - /// - /// # Errors - /// - /// Returns an error if: - /// - /// * [`StorageClientError::Transport`] if the request fails. - async fn get_initial_session(channel: Channel) -> Result { - SessionManagementServiceClient::new(channel) - .get_session(storage::Void {}) - .await - .map(|response| response.into_inner().session_id) - .map_err(to_transport_error) - } } #[async_trait] @@ -131,7 +106,6 @@ impl SchedulerStorageClient for GrpcSchedulerStorageClient { async fn job_state(&self, job_id: JobId) -> Result { let request = storage::JobIdRequest { job_id: job_id.get(), - session_id: self.session_tracker.current(), }; let response = self .job_client @@ -146,7 +120,17 @@ impl SchedulerStorageClient for GrpcSchedulerStorageClient { impl From for StorageClientError { fn from(error: storage::InboundQueueResponseError) -> Self { - inbound_queue_response_error_to_client_error(error) + 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}")), + } } } @@ -158,27 +142,21 @@ impl From for StorageClientError { /// /// Returns an error if: /// -/// * [`StorageClientError::Transport`] if either value cannot fit in the protobuf field type. +/// * [`StorageClientError::InvalidInput`] if either value cannot fit in the protobuf field type. fn poll_ready_tasks_request( max_items: usize, wait: Duration, ) -> Result { Ok(storage::PollReadyTasksRequest { - max_items: u64::try_from(max_items).map_err(to_transport_error)?, - wait_ns: u64::try_from(wait.as_nanos()).map_err(to_transport_error)?, + max_items: u64::try_from(max_items).map_err(to_invalid_input_error)?, + wait_ms: u64::try_from(wait.as_millis()).map_err(to_invalid_input_error)?, }) } /// # Returns /// -/// [`storage::PollReadyTasksResponse`] converted into scheduler entries on success. -/// -/// # Errors -/// -/// Returns an error if: -/// -/// * [`StorageClientError::Transport`] if the response is malformed. -/// * Forwards [`StorageClientError::from`]'s return values on failure. +/// [`storage::PollReadyTasksResponse`] converted into +/// [`Result<(SessionId, Vec), StorageClientError>`]. fn poll_ready_tasks_response_to_result( response: storage::PollReadyTasksResponse, ) -> Result<(SessionId, Vec), StorageClientError> { @@ -193,13 +171,8 @@ fn poll_ready_tasks_response_to_result( /// # Returns /// -/// [`storage::ReadyTasks`] converted into scheduler entries on success. -/// -/// # Errors -/// -/// Returns an error if: -/// -/// * [`StorageClientError::Transport`] if a ready task is missing or has an invalid task ID. +/// [`storage::ReadyTasks`] converted into +/// [`Result<(SessionId, Vec), StorageClientError>`]. fn ready_tasks_to_result( tasks: storage::ReadyTasks, ) -> Result<(SessionId, Vec), StorageClientError> { @@ -240,14 +213,7 @@ fn ready_task_to_inbound_entry( /// # Returns /// -/// [`storage::JobStateResponse`] converted into [`JobState`] on success. -/// -/// # Errors -/// -/// Returns an error if: -/// -/// * [`StorageClientError::Transport`] if the response is malformed. -/// * Forwards [`StorageClientError::from`]'s return values on failure. +/// [`storage::JobStateResponse`] converted into [`Result`]. fn job_state_response_to_result( response: storage::JobStateResponse, job_id: JobId, @@ -268,29 +234,6 @@ fn job_state_response_to_result( } } -/// # Returns -/// -/// [`storage::InboundQueueResponseError`] converted into [`StorageClientError`]. -fn inbound_queue_response_error_to_client_error( - error: storage::InboundQueueResponseError, -) -> StorageClientError { - match inbound_queue_response_error::ErrCode::try_from(error.err_code) { - Ok(inbound_queue_response_error::ErrCode::InboundClosed) => { - StorageClientError::InboundClosed - } - Ok(inbound_queue_response_error::ErrCode::InvalidInput) => { - StorageClientError::InvalidInput(error.message) - } - Ok( - inbound_queue_response_error::ErrCode::Server - | inbound_queue_response_error::ErrCode::Unspecified, - ) => StorageClientError::Server(error.message), - Err(error) => { - StorageClientError::Transport(format!("unknown scheduler storage error kind: {error}")) - } - } -} - /// # Returns /// /// [`storage::JobOrchestrationError`] converted into [`StorageClientError`]. @@ -327,6 +270,15 @@ fn to_transport_error(error: impl std::fmt::Display) -> StorageClientError { StorageClientError::Transport(error.to_string()) } +/// Converts a displayable out-of-range error into [`StorageClientError::InvalidInput`]. +/// +/// # Returns +/// +/// A [`StorageClientError::InvalidInput`] containing `error`'s display string. +fn to_invalid_input_error(error: impl std::fmt::Display) -> StorageClientError { + StorageClientError::InvalidInput(error.to_string()) +} + #[cfg(test)] mod tests { use spider_core::types::id::{JobId, ResourceGroupId, TaskId}; @@ -338,16 +290,21 @@ mod tests { use super::*; + const SESSION_ID: SessionId = 11; + const RESOURCE_GROUP_ID: u64 = 3; + const JOB_ID: u64 = 5; + const TASK_INDEX: usize = 7; + #[test] fn poll_ready_tasks_response_converts_entries() { let response = storage::PollReadyTasksResponse { result: Some(poll_ready_tasks_response::Result::Tasks( storage::ReadyTasks { - session_id: 11, + session_id: SESSION_ID, tasks: vec![storage::ReadyTask { - resource_group_id: 3, - job_id: 5, - task_id: Some(storage::TaskId::from(TaskId::Index(7))), + resource_group_id: RESOURCE_GROUP_ID, + job_id: JOB_ID, + task_id: Some(storage::TaskId::from(TaskId::Index(TASK_INDEX))), }], }, )), @@ -356,26 +313,28 @@ mod tests { let (session_id, entries) = poll_ready_tasks_response_to_result(response) .expect("poll response conversion should succeed"); - assert_eq!(session_id, 11); + assert_eq!(session_id, SESSION_ID); assert_eq!( entries, vec![InboundEntry { - resource_group_id: ResourceGroupId::from(3), - job_id: JobId::from(5), - task_id: TaskId::Index(7), + resource_group_id: ResourceGroupId::from(RESOURCE_GROUP_ID), + job_id: JobId::from(JOB_ID), + task_id: TaskId::Index(TASK_INDEX), }] ); } #[test] fn poll_ready_tasks_response_rejects_missing_task_id() { + 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: 11, + session_id: SESSION_ID, tasks: vec![storage::ReadyTask { - resource_group_id: 3, - job_id: 5, + resource_group_id: RESOURCE_GROUP_ID, + job_id: JOB_ID, task_id: None, }], }, @@ -384,7 +343,7 @@ mod tests { match poll_ready_tasks_response_to_result(response) { Err(StorageClientError::Transport(message)) => { - assert!(message.contains("missing task ID")); + assert!(message.contains(MISSING_TASK_ID_MESSAGE)); } result => panic!("unexpected poll response conversion result: {result:?}"), } @@ -392,9 +351,11 @@ mod tests { #[test] fn inbound_queue_response_error_maps_inbound_closed() { + const ERROR_MESSAGE: &str = "closed"; + let error = storage::InboundQueueResponseError { err_code: inbound_queue_response_error::ErrCode::InboundClosed.into(), - message: "closed".to_owned(), + message: ERROR_MESSAGE.to_owned(), }; assert!(matches!( diff --git a/components/spider-scheduler/src/storage_client.rs b/components/spider-scheduler/src/storage_client/mod.rs similarity index 98% rename from components/spider-scheduler/src/storage_client.rs rename to components/spider-scheduler/src/storage_client/mod.rs index 9f7adaf4..e8a607bd 100644 --- a/components/spider-scheduler/src/storage_client.rs +++ b/components/spider-scheduler/src/storage_client/mod.rs @@ -10,6 +10,10 @@ use spider_core::{ use crate::{error::StorageClientError, types::InboundEntry}; +pub mod grpc; + +pub use grpc::GrpcSchedulerStorageClient; + /// The scheduler's view of the storage layer. /// /// Abstracts the storage-owned inbound queue and the read-only queries a scheduling algorithm From 52c808a03571e9d5b1b1add2062fcda1da957413 Mon Sep 17 00:00:00 2001 From: LinZhihao-723 Date: Fri, 12 Jun 2026 19:04:40 -0400 Subject: [PATCH 16/16] Toml lint. --- components/spider-scheduler/Cargo.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/components/spider-scheduler/Cargo.toml b/components/spider-scheduler/Cargo.toml index bf2de66d..3254c747 100644 --- a/components/spider-scheduler/Cargo.toml +++ b/components/spider-scheduler/Cargo.toml @@ -17,7 +17,6 @@ tokio = { version = "1.52.3", features = ["sync", "time"] } tokio-util = "0.7.18" tonic = "0.12.3" - [dev-dependencies] anyhow = "1.0.102" dashmap = "6.1.0"