feat(spider-utils): Add a gRPC connection pool and route the execution-manager and scheduler gRPC clients through it.

components/spider-execution-manager/src/client/grpc/liveness.rs

@@ -1,6 +1,6 @@
 //! gRPC-backed [`LivenessClient`] implementation.
 
-use std::net::IpAddr;
+use std::{net::IpAddr, num::NonZeroUsize};
 
 use async_trait::async_trait;
 use spider_core::types::id::{ExecutionManagerId, SessionId};
@@ -11,18 +11,19 @@ use spider_proto_rust::storage::{
     register_execution_manager_response,
     update_execution_manager_heartbeat_response,
 };
+use spider_utils::grpc::client::ConnectionPool;
 use tonic::transport::{Channel, Endpoint};
 
 use crate::client::liveness::{LivenessClient, LivenessResponseError, RegistrationResponse};
 
 /// gRPC-backed [`LivenessClient`] implementation.
 #[derive(Debug, Clone)]
 pub struct GrpcLivenessClient {
-    client: ExecutionManagerLivenessServiceClient<Channel>,
+    connection_pool: ConnectionPool<ExecutionManagerLivenessServiceClient<Channel>>,
 }
 
 impl GrpcLivenessClient {
-    /// Connects to the storage gRPC endpoint.
+    /// Connects a pool of `pool_size` connections to the liveness gRPC endpoint.
     ///
     /// # Returns
     ///
@@ -33,11 +34,17 @@ impl GrpcLivenessClient {
     /// Returns an error if:
     ///
     /// * [`LivenessResponseError::Transport`] if tonic cannot create or connect to the endpoint.
-    pub async fn connect(endpoint: Endpoint) -> Result<Self, LivenessResponseError> {
-        ExecutionManagerLivenessServiceClient::connect(endpoint)
-            .await
-            .map(|client| Self { client })
-            .map_err(to_transport_error)
+    pub async fn connect(
+        endpoint: Endpoint,
+        pool_size: NonZeroUsize,
+    ) -> Result<Self, LivenessResponseError> {
+        let connection_pool = ConnectionPool::connect(endpoint, pool_size, |channel| {
+            ExecutionManagerLivenessServiceClient::new(channel)
+        })
+        .await
+        .map_err(to_transport_error)?;
+
+        Ok(Self { connection_pool })
     }
 }
 
@@ -48,8 +55,8 @@ impl LivenessClient for GrpcLivenessClient {
             ip_address: ip.to_string(),
         };
         let response = self
-            .client
-            .clone()
+            .connection_pool
+            .get_client()
             .register_execution_manager(request)
             .await
             .map_err(to_transport_error)?
@@ -66,8 +73,8 @@ impl LivenessClient for GrpcLivenessClient {
             execution_manager_id: em_id.get(),
         };
         let response = self
-            .client
-            .clone()
+            .connection_pool
+            .get_client()
             .update_execution_manager_heartbeat(request)
             .await
             .map_err(to_transport_error)?

components/spider-execution-manager/src/client/grpc/scheduler.rs

@@ -1,20 +1,23 @@
 //! gRPC-backed [`SchedulerClient`] implementation.
 
+use std::num::NonZeroUsize;
+
 use async_trait::async_trait;
 use spider_core::types::{id::ExecutionManagerId, scheduler::TaskAssignmentRecord};
 use spider_proto_rust::scheduler::{self, scheduler_service_client::SchedulerServiceClient};
+use spider_utils::grpc::client::ConnectionPool;
 use tonic::transport::{Channel, Endpoint};
 
 use crate::client::{SchedulerClient, SchedulerError, SchedulerResponse};
 
 /// gRPC-backed [`SchedulerClient`] implementation.
 #[derive(Debug, Clone)]
 pub struct GrpcSchedulerClient {
-    client: SchedulerServiceClient<Channel>,
+    connection_pool: ConnectionPool<SchedulerServiceClient<Channel>>,
 }
 
 impl GrpcSchedulerClient {
-    /// Connects to the scheduler gRPC endpoint.
+    /// Connects a pool of `pool_size` connections to the scheduler gRPC endpoint.
     ///
     /// # Returns
     ///
@@ -25,11 +28,17 @@ impl GrpcSchedulerClient {
     /// Returns an error if:
     ///
     /// * [`SchedulerError::Transport`] if tonic cannot create or connect to the endpoint.
-    pub async fn connect(endpoint: Endpoint) -> Result<Self, SchedulerError> {
-        SchedulerServiceClient::connect(endpoint)
-            .await
-            .map(|client| Self { client })
-            .map_err(to_transport_error)
+    pub async fn connect(
+        endpoint: Endpoint,
+        pool_size: NonZeroUsize,
+    ) -> Result<Self, SchedulerError> {
+        let connection_pool = ConnectionPool::connect(endpoint, pool_size, |channel| {
+            SchedulerServiceClient::new(channel)
+        })
+        .await
+        .map_err(to_transport_error)?;
+
+        Ok(Self { connection_pool })
     }
 }
 
@@ -43,8 +52,8 @@ impl SchedulerClient for GrpcSchedulerClient {
     ) -> Result<SchedulerResponse, SchedulerError> {
         loop {
             let response = self
-                .client
-                .clone()
+                .connection_pool
+                .get_client()
                 .next_task(scheduler::NextTaskRequest {
                     execution_manager_id: em_id.get(),
                     prev_assignment: prev_assignment.map(Into::into),
@@ -63,8 +72,8 @@ impl SchedulerClient for GrpcSchedulerClient {
     }
 
     async fn heartbeat(&self, em_id: ExecutionManagerId) -> Result<(), SchedulerError> {
-        self.client
-            .clone()
+        self.connection_pool
+            .get_client()
             .heartbeat(scheduler::HeartbeatRequest {
                 execution_manager_id: em_id.get(),
             })
@@ -79,8 +88,8 @@ impl SchedulerClient for GrpcSchedulerClient {
         prev_assignments: Vec<TaskAssignmentRecord>,
     ) {
         if let Err(error) = self
-            .client
-            .clone()
+            .connection_pool
+            .get_client()
             .shutdown(scheduler::ShutdownRequest {
                 execution_manager_id: em_id.get(),
                 prev_assignments: prev_assignments.into_iter().map(Into::into).collect(),

components/spider-execution-manager/src/client/grpc/storage.rs

@@ -3,6 +3,8 @@
 //! Wraps the generated [`TaskInstanceManagementServiceClient`] and adapts its protobuf
 //! request/response types to the transport-agnostic [`StorageClient`] trait.
 
+use std::num::NonZeroUsize;
+
 use async_trait::async_trait;
 use spider_core::types::{
     id::{ExecutionManagerId, JobId, SessionId, TaskId, TaskInstanceId},
@@ -12,6 +14,7 @@ use spider_proto_rust::{
     common,
     storage::{self, task_instance_management_service_client::TaskInstanceManagementServiceClient},
 };
+use spider_utils::grpc::client::ConnectionPool;
 use tonic::{
     Code,
     Status,
@@ -23,11 +26,11 @@ use crate::client::storage::{StorageClient, StorageResponseError};
 /// gRPC-backed [`StorageClient`] implementation.
 #[derive(Debug, Clone)]
 pub struct GrpcStorageClient {
-    client: TaskInstanceManagementServiceClient<Channel>,
+    connection_pool: ConnectionPool<TaskInstanceManagementServiceClient<Channel>>,
 }
 
 impl GrpcStorageClient {
-    /// Connects to the storage gRPC endpoint.
+    /// Connects a pool of `pool_size` connections to the storage gRPC endpoint.
     ///
     /// # Returns
     ///
@@ -38,11 +41,17 @@ impl GrpcStorageClient {
     /// Returns an error if:
     ///
     /// * [`StorageResponseError::Transport`] if tonic cannot create or connect to the endpoint.
-    pub async fn connect(endpoint: Endpoint) -> Result<Self, StorageResponseError> {
-        TaskInstanceManagementServiceClient::connect(endpoint)
-            .await
-            .map(|client| Self { client })
-            .map_err(to_transport_error)
+    pub async fn connect(
+        endpoint: Endpoint,
+        pool_size: NonZeroUsize,
+    ) -> Result<Self, StorageResponseError> {
+        let connection_pool = ConnectionPool::connect(endpoint, pool_size, |channel| {
+            TaskInstanceManagementServiceClient::new(channel)
+        })
+        .await
+        .map_err(to_transport_error)?;
+
+        Ok(Self { connection_pool })
     }
 }
 
@@ -62,8 +71,8 @@ impl StorageClient for GrpcStorageClient {
             session_id,
         };
         let response = self
-            .client
-            .clone()
+            .connection_pool
+            .get_client()
             .register_task_instance(request)
             .await
             .map_err(|status| status_to_error(&status))?
@@ -95,8 +104,8 @@ impl StorageClient for GrpcStorageClient {
             serialized_outputs: serialized_outputs.unwrap_or_default(),
             task_instance_id,
         };
-        self.client
-            .clone()
+        self.connection_pool
+            .get_client()
             .report_task_success(request)
             .await
             .map_err(|status| status_to_error(&status))?;
@@ -120,8 +129,8 @@ impl StorageClient for GrpcStorageClient {
             error_message,
             task_instance_id,
         };
-        self.client
-            .clone()
+        self.connection_pool
+            .get_client()
             .report_task_failure(request)
             .await
             .map_err(|status| status_to_error(&status))?;

components/spider-scheduler/Cargo.toml

@@ -13,6 +13,7 @@ async-trait = "0.1.89"
 serde = { version = "1.0.228", features = ["derive"] }
 spider-core = { path = "../spider-core" }
 spider-proto-rust = { path = "../spider-proto-rust" }
+spider-utils = { path = "../spider-utils" }
 thiserror = "2.0.18"
 tokio = { version = "1.52.3", features = ["macros", "rt", "sync", "time"] }
 tokio-util = "0.7.18"

components/spider-scheduler/src/storage_client/grpc.rs

@@ -1,6 +1,6 @@
 //! gRPC-backed [`SchedulerStorageClient`] implementation.
 
-use std::time::Duration;
+use std::{num::NonZeroUsize, time::Duration};
 
 use async_trait::async_trait;
 use spider_core::{
@@ -14,6 +14,7 @@ use spider_proto_rust::storage::{
     job_orchestration_service_client::JobOrchestrationServiceClient,
     poll_ready_tasks_response,
 };
+use spider_utils::grpc::client::ConnectionPool;
 use tonic::{
     Code,
     transport::{Channel, Endpoint},
@@ -28,12 +29,12 @@ use crate::{
 /// gRPC-backed [`SchedulerStorageClient`] implementation.
 #[derive(Debug, Clone)]
 pub struct GrpcSchedulerStorageClient {
-    scheduler_client: InboundQueueServiceClient<Channel>,
-    job_client: JobOrchestrationServiceClient<Channel>,
+    inbound_queue_connection_pool: ConnectionPool<InboundQueueServiceClient<Channel>>,
+    job_orchestration_connection_pool: ConnectionPool<JobOrchestrationServiceClient<Channel>>,
 }
 
 impl GrpcSchedulerStorageClient {
-    /// Connects to the storage gRPC endpoint.
+    /// Connects pools of `pool_size` connections to the storage gRPC endpoint.
     ///
     /// # Returns
     ///
@@ -44,12 +45,26 @@ impl GrpcSchedulerStorageClient {
     /// Returns an error if:
     ///
     /// * [`StorageClientError::Transport`] if tonic cannot create or connect to the endpoint.
-    pub async fn connect(endpoint: Endpoint) -> Result<Self, StorageClientError> {
-        let channel = endpoint.connect().await.map_err(to_transport_error)?;
+    pub async fn connect(
+        endpoint: Endpoint,
+        pool_size: NonZeroUsize,
+    ) -> Result<Self, StorageClientError> {
+        let inbound_queue_connection_pool =
+            ConnectionPool::connect(endpoint.clone(), pool_size, |channel| {
+                InboundQueueServiceClient::new(channel)
+            })
+            .await
+            .map_err(to_transport_error)?;
+        let job_orchestration_connection_pool =
+            ConnectionPool::connect(endpoint, pool_size, |channel| {
+                JobOrchestrationServiceClient::new(channel)
+            })
+            .await
+            .map_err(to_transport_error)?;
 
         Ok(Self {
-            scheduler_client: InboundQueueServiceClient::new(channel.clone()),
-            job_client: JobOrchestrationServiceClient::new(channel),
+            inbound_queue_connection_pool,
+            job_orchestration_connection_pool,
         })
     }
 }
@@ -63,8 +78,8 @@ impl SchedulerStorageClient for GrpcSchedulerStorageClient {
     ) -> Result<(SessionId, Vec<InboundEntry>), StorageClientError> {
         let request = poll_ready_tasks_request(max_items, wait)?;
         let response = self
-            .scheduler_client
-            .clone()
+            .inbound_queue_connection_pool
+            .get_client()
             .poll_ready_tasks(request)
             .await
             .map_err(to_transport_error)?
@@ -79,8 +94,8 @@ impl SchedulerStorageClient for GrpcSchedulerStorageClient {
     ) -> Result<(SessionId, Vec<InboundEntry>), StorageClientError> {
         let request = poll_ready_tasks_request(max_items, wait)?;
         let response = self
-            .scheduler_client
-            .clone()
+            .inbound_queue_connection_pool
+            .get_client()
             .poll_ready_commit_tasks(request)
             .await
             .map_err(to_transport_error)?
@@ -95,8 +110,8 @@ impl SchedulerStorageClient for GrpcSchedulerStorageClient {
     ) -> Result<(SessionId, Vec<InboundEntry>), StorageClientError> {
         let request = poll_ready_tasks_request(max_items, wait)?;
         let response = self
-            .scheduler_client
-            .clone()
+            .inbound_queue_connection_pool
+            .get_client()
             .poll_ready_cleanup_tasks(request)
             .await
             .map_err(to_transport_error)?
@@ -109,8 +124,8 @@ impl SchedulerStorageClient for GrpcSchedulerStorageClient {
             job_id: job_id.get(),
         };
         let response = self
-            .job_client
-            .clone()
+            .job_orchestration_connection_pool
+            .get_client()
             .get_job_state(request)
             .await
             .map_err(|status| match status.code() {

components/spider-utils/Cargo.toml

@@ -11,6 +11,7 @@ path = "src/lib.rs"
 rmp-serde = "1.3.1"
 serde = { version = "1.0.228", features = ["derive"] }
 thiserror = "2.0.18"
+tonic = "0.12.3"
 tracing-appender = "0.2.5"
 tracing-subscriber = {
   version = "0.3.23",