From 4d492223adb9c950cbcc1052933217792fa8ed94 Mon Sep 17 00:00:00 2001 From: LinZhihao-723 Date: Mon, 13 Jul 2026 22:12:57 -0400 Subject: [PATCH 1/3] Implementation done. --- Cargo.lock | 2 + components/spider-client/src/client.rs | 103 +++++- components/spider-client/src/grpc/job.rs | 104 +++--- .../spider-client/src/grpc/resource_group.rs | 43 ++- components/spider-client/src/lib.rs | 2 + components/spider-utils/Cargo.toml | 5 + components/spider-utils/src/grpc/mod.rs | 1 + components/spider-utils/src/grpc/retry.rs | 318 ++++++++++++++++++ tests/huntsman/e2e/src/test_driver.rs | 7 +- 9 files changed, 507 insertions(+), 78 deletions(-) create mode 100644 components/spider-utils/src/grpc/retry.rs diff --git a/Cargo.lock b/Cargo.lock index fb506b56d..4e8020fe2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2112,9 +2112,11 @@ dependencies = [ name = "spider-utils" version = "0.1.0" dependencies = [ + "rand 0.9.4", "rmp-serde", "serde", "thiserror", + "tokio", "tonic", "tracing-appender", "tracing-subscriber", diff --git a/components/spider-client/src/client.rs b/components/spider-client/src/client.rs index e54820f1a..0f359f510 100644 --- a/components/spider-client/src/client.rs +++ b/components/spider-client/src/client.rs @@ -1,6 +1,7 @@ //! [`SpiderClient`] — the top-level handle holding the gRPC connection pools. use std::num::NonZeroUsize; +use std::time::Duration; use spider_core::job::JobState; use spider_core::task::TaskGraph; @@ -8,6 +9,7 @@ use spider_core::types::id::JobId; use spider_core::types::id::ResourceGroupId; use spider_core::types::io::TaskInput; use spider_core::types::io::TaskOutput; +use spider_utils::grpc::retry::RetryConfig; use tonic::transport::Endpoint; use crate::error::ClientError; @@ -22,27 +24,17 @@ pub struct SpiderClient { } impl SpiderClient { - /// Connects pools of `pool_size` connections to the storage gRPC endpoint. + /// Creates a builder for connecting a [`SpiderClient`] to `endpoint`. /// /// # Returns /// - /// A new [`SpiderClient`] connected to `endpoint` on success. - /// - /// # Errors - /// - /// Returns an error if: - /// - /// * [`ClientError::Transport`] if tonic cannot create or connect to the endpoint. - pub async fn connect(endpoint: Endpoint, pool_size: NonZeroUsize) -> Result { - let (job_orchestration, resource_group) = tokio::try_join!( - JobOrchestrationClient::connect(endpoint.clone(), pool_size), - ResourceGroupManagementClient::connect(endpoint, pool_size), - )?; - - Ok(Self { - job_orchestration, - resource_group, - }) + /// A [`SpiderClientBuilder`] for `endpoint` with default pool size and retry configuration. + pub fn builder(endpoint: Endpoint) -> SpiderClientBuilder { + SpiderClientBuilder { + endpoint, + pool_size: DEFAULT_POOL_SIZE, + retry_config: RetryConfig::default(), + } } /// Serializes and zstd-compresses the task graph and inputs, registers the job, and returns its @@ -222,3 +214,78 @@ impl SpiderClient { .await } } + +/// Builder for configuring and connecting a [`SpiderClient`]. +pub struct SpiderClientBuilder { + endpoint: Endpoint, + pool_size: NonZeroUsize, + retry_config: RetryConfig, +} + +impl SpiderClientBuilder { + /// Sets the size of each gRPC connection pool. + /// + /// # Returns + /// + /// The builder with `pool_size` set. + #[must_use] + pub const fn pool_size(mut self, pool_size: NonZeroUsize) -> Self { + self.pool_size = pool_size; + self + } + + /// Sets the number of retries allowed after the initial attempt. + /// + /// # Returns + /// + /// The builder with `max_retries` set. + #[must_use] + pub const fn max_retries(mut self, max_retries: usize) -> Self { + self.retry_config.max_retries = max_retries; + self + } + + /// Sets the upper bound on the exponential backoff between attempts. + /// + /// # Returns + /// + /// The builder with `max_backoff` set. + #[must_use] + pub const fn max_backoff(mut self, max_backoff: Duration) -> Self { + self.retry_config.max_backoff = max_backoff; + self + } + + /// Connects pools of the configured size to the storage gRPC endpoint. + /// + /// # Returns + /// + /// A new [`SpiderClient`] connected to the configured endpoint on success. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`ClientError::Transport`] if tonic cannot create or connect to the endpoint. + pub async fn connect(self) -> Result { + let (job_orchestration, resource_group) = tokio::try_join!( + JobOrchestrationClient::connect( + self.endpoint.clone(), + self.pool_size, + self.retry_config + ), + ResourceGroupManagementClient::connect( + self.endpoint, + self.pool_size, + self.retry_config + ), + )?; + + Ok(SpiderClient { + job_orchestration, + resource_group, + }) + } +} + +const DEFAULT_POOL_SIZE: NonZeroUsize = NonZeroUsize::new(8).unwrap(); diff --git a/components/spider-client/src/grpc/job.rs b/components/spider-client/src/grpc/job.rs index d6b353cfe..e45cc5537 100644 --- a/components/spider-client/src/grpc/job.rs +++ b/components/spider-client/src/grpc/job.rs @@ -15,6 +15,8 @@ use spider_proto_rust::error::Error as ProtoError; use spider_proto_rust::storage::JobOrchestrationServiceClient; use spider_proto_rust::storage::{self}; use spider_utils::grpc::client::ConnectionPool; +use spider_utils::grpc::retry::RetryConfig; +use spider_utils::grpc::retry::call_with_retry; use tonic::Code; use tonic::Status; use tonic::transport::Channel; @@ -27,6 +29,7 @@ use crate::error::to_transport_error; #[derive(Debug, Clone)] pub struct JobOrchestrationClient { connection_pool: ConnectionPool>, + retry_config: RetryConfig, } impl JobOrchestrationClient { @@ -41,14 +44,21 @@ impl JobOrchestrationClient { /// Returns an error if: /// /// * [`ClientError::Transport`] if tonic cannot create or connect to the endpoint. - pub async fn connect(endpoint: Endpoint, pool_size: NonZeroUsize) -> Result { + pub async fn connect( + endpoint: Endpoint, + pool_size: NonZeroUsize, + retry_config: RetryConfig, + ) -> Result { let connection_pool = ConnectionPool::connect(endpoint, pool_size, |channel| { JobOrchestrationServiceClient::new(channel) }) .await .map_err(to_transport_error)?; - Ok(Self { connection_pool }) + Ok(Self { + connection_pool, + retry_config, + }) } /// Serializes and zstd-compresses the task graph and inputs, registers the job, and returns @@ -80,13 +90,15 @@ impl JobOrchestrationClient { compressed_serialized_task_graph, compressed_serialized_inputs, }; - let response = self - .connection_pool - .get_client() - .register_job(request) - .await - .map_err(|status| job_status_to_error(&status))? - .into_inner(); + let response = call_with_retry(self.retry_config, async || { + self.connection_pool + .get_client() + .register_job(request.clone()) + .await + }) + .await + .map_err(|status| job_status_to_error(&status))? + .into_inner(); Ok(JobId::from(response.job_id)) } @@ -107,13 +119,12 @@ impl JobOrchestrationClient { let request = storage::JobIdRequest { job_id: job_id.get(), }; - let response = self - .connection_pool - .get_client() - .start_job(request) - .await - .map_err(|status| job_status_to_error(&status))? - .into_inner(); + let response = call_with_retry(self.retry_config, async || { + self.connection_pool.get_client().start_job(request).await + }) + .await + .map_err(|status| job_status_to_error(&status))? + .into_inner(); job_state_response_to_result(response) } @@ -134,13 +145,12 @@ impl JobOrchestrationClient { let request = storage::JobIdRequest { job_id: job_id.get(), }; - let response = self - .connection_pool - .get_client() - .cancel_job(request) - .await - .map_err(|status| job_status_to_error(&status))? - .into_inner(); + let response = call_with_retry(self.retry_config, async || { + self.connection_pool.get_client().cancel_job(request).await + }) + .await + .map_err(|status| job_status_to_error(&status))? + .into_inner(); job_state_response_to_result(response) } @@ -161,13 +171,15 @@ impl JobOrchestrationClient { let request = storage::JobIdRequest { job_id: job_id.get(), }; - let response = self - .connection_pool - .get_client() - .get_job_state(request) - .await - .map_err(|status| job_status_to_error(&status))? - .into_inner(); + let response = call_with_retry(self.retry_config, async || { + self.connection_pool + .get_client() + .get_job_state(request) + .await + }) + .await + .map_err(|status| job_status_to_error(&status))? + .into_inner(); job_state_response_to_result(response) } @@ -190,13 +202,15 @@ impl JobOrchestrationClient { let request = storage::JobIdRequest { job_id: job_id.get(), }; - let response = self - .connection_pool - .get_client() - .get_job_outputs(request) - .await - .map_err(|status| job_status_to_error(&status))? - .into_inner(); + let response = call_with_retry(self.retry_config, async || { + self.connection_pool + .get_client() + .get_job_outputs(request) + .await + }) + .await + .map_err(|status| job_status_to_error(&status))? + .into_inner(); SerializedTaskOutputs::deserialize_from_raw(&response.serialized_outputs) .map_err(|error| ClientError::Deserialization(error.to_string())) @@ -217,13 +231,15 @@ impl JobOrchestrationClient { let request = storage::JobIdRequest { job_id: job_id.get(), }; - let response = self - .connection_pool - .get_client() - .get_job_error(request) - .await - .map_err(|status| job_status_to_error(&status))? - .into_inner(); + let response = call_with_retry(self.retry_config, async || { + self.connection_pool + .get_client() + .get_job_error(request) + .await + }) + .await + .map_err(|status| job_status_to_error(&status))? + .into_inner(); Ok(response.error_message) } diff --git a/components/spider-client/src/grpc/resource_group.rs b/components/spider-client/src/grpc/resource_group.rs index 102f6151b..2aaad9b03 100644 --- a/components/spider-client/src/grpc/resource_group.rs +++ b/components/spider-client/src/grpc/resource_group.rs @@ -6,6 +6,8 @@ use spider_core::types::id::ResourceGroupId; use spider_proto_rust::storage::ResourceGroupManagementServiceClient; use spider_proto_rust::storage::{self}; use spider_utils::grpc::client::ConnectionPool; +use spider_utils::grpc::retry::RetryConfig; +use spider_utils::grpc::retry::call_with_retry; use tonic::Code; use tonic::Status; use tonic::transport::Channel; @@ -18,6 +20,7 @@ use crate::error::to_transport_error; #[derive(Debug, Clone)] pub struct ResourceGroupManagementClient { connection_pool: ConnectionPool>, + retry_config: RetryConfig, } impl ResourceGroupManagementClient { @@ -32,14 +35,21 @@ impl ResourceGroupManagementClient { /// Returns an error if: /// /// * [`ClientError::Transport`] if tonic cannot create or connect to the endpoint. - pub async fn connect(endpoint: Endpoint, pool_size: NonZeroUsize) -> Result { + pub async fn connect( + endpoint: Endpoint, + pool_size: NonZeroUsize, + retry_config: RetryConfig, + ) -> Result { let connection_pool = ConnectionPool::connect(endpoint, pool_size, |channel| { ResourceGroupManagementServiceClient::new(channel) }) .await .map_err(to_transport_error)?; - Ok(Self { connection_pool }) + Ok(Self { + connection_pool, + retry_config, + }) } /// Registers an external resource group and returns its server-assigned id. @@ -63,13 +73,15 @@ impl ResourceGroupManagementClient { external_resource_group_id, password, }; - let response = self - .connection_pool - .get_client() - .add_resource_group(request) - .await - .map_err(|status| resource_group_status_to_error(&status))? - .into_inner(); + let response = call_with_retry(self.retry_config, async || { + self.connection_pool + .get_client() + .add_resource_group(request.clone()) + .await + }) + .await + .map_err(|status| resource_group_status_to_error(&status))? + .into_inner(); Ok(ResourceGroupId::from(response.resource_group_id)) } @@ -95,11 +107,14 @@ impl ResourceGroupManagementClient { resource_group_id: resource_group_id.get(), password, }; - self.connection_pool - .get_client() - .verify_resource_group(request) - .await - .map_err(|status| resource_group_status_to_error(&status))?; + call_with_retry(self.retry_config, async || { + self.connection_pool + .get_client() + .verify_resource_group(request.clone()) + .await + }) + .await + .map_err(|status| resource_group_status_to_error(&status))?; Ok(()) } } diff --git a/components/spider-client/src/lib.rs b/components/spider-client/src/lib.rs index 61e3d2f93..ed682fd85 100644 --- a/components/spider-client/src/lib.rs +++ b/components/spider-client/src/lib.rs @@ -12,3 +12,5 @@ pub mod error; pub(crate) mod grpc; pub use client::SpiderClient; +pub use client::SpiderClientBuilder; +pub use spider_utils::grpc::retry::RetryConfig; diff --git a/components/spider-utils/Cargo.toml b/components/spider-utils/Cargo.toml index 91dab6c4a..079a30a03 100644 --- a/components/spider-utils/Cargo.toml +++ b/components/spider-utils/Cargo.toml @@ -8,9 +8,11 @@ name = "spider_utils" path = "src/lib.rs" [dependencies] +rand = "0.9.1" rmp-serde = "1.3.1" serde = { version = "1.0.228", features = ["derive"] } thiserror = "2.0.18" +tokio = { version = "1.50.0", features = ["time"] } tonic = "0.14.6" tracing-appender = "0.2.5" tracing-subscriber = { @@ -19,3 +21,6 @@ tracing-subscriber = { features = ["env-filter", "fmt", "json"] } yaml_serde = "0.10.4" + +[dev-dependencies] +tokio = { version = "1.50.0", features = ["macros", "rt", "time"] } diff --git a/components/spider-utils/src/grpc/mod.rs b/components/spider-utils/src/grpc/mod.rs index 87de00150..1884e1f99 100644 --- a/components/spider-utils/src/grpc/mod.rs +++ b/components/spider-utils/src/grpc/mod.rs @@ -1,6 +1,7 @@ //! gRPC-related utilities. pub mod client; +pub mod retry; #[derive(thiserror::Error, Debug)] pub enum Error { diff --git a/components/spider-utils/src/grpc/retry.rs b/components/spider-utils/src/grpc/retry.rs new file mode 100644 index 000000000..d3a49f1a8 --- /dev/null +++ b/components/spider-utils/src/grpc/retry.rs @@ -0,0 +1,318 @@ +//! An async retry helper for transient gRPC call failures. + +use std::time::Duration; + +use rand::Rng; +use tonic::Code; +use tonic::Status; + +/// Tunable retry policy for transient gRPC call failures. +#[derive(Clone, Copy, Debug)] +pub struct RetryConfig { + /// The number of retries allowed after the initial attempt. + pub max_retries: usize, + /// The upper bound on the exponential backoff between attempts. + pub max_backoff: Duration, +} + +impl Default for RetryConfig { + /// # Returns + /// + /// A [`RetryConfig`] with [`DEFAULT_MAX_RETRIES`] retries and a [`DEFAULT_MAX_BACKOFF`] backoff + /// cap. + fn default() -> Self { + Self { + max_retries: DEFAULT_MAX_RETRIES, + max_backoff: DEFAULT_MAX_BACKOFF, + } + } +} + +/// Repeatedly invokes an async gRPC call until it succeeds, a non-retriable error occurs, or the +/// retry budget is exhausted. +/// +/// Between attempts, the helper sleeps for an exponentially increasing backoff that doubles on each +/// retry, capped at `max_backoff`, plus a small random jitter, so the actual wait may slightly +/// exceed `max_backoff`. +/// +/// `max_retries` counts the retries allowed *after* the initial attempt, so `grpc_call` is invoked +/// at most `max_retries + 1` times. +/// +/// # Type Parameters +/// +/// * `ResponseType` - The success value produced by `grpc_call`. +/// * `ErrorType` - The error produced by `grpc_call`. +/// * `GrpcCall` - The async closure performing the gRPC call. +/// * `RetriableCheck` - Classifies an error as retriable or not. +/// +/// # Returns +/// +/// The first `Ok` value returned by `grpc_call`. +/// +/// # Errors +/// +/// Returns an error if: +/// +/// * Forwards `grpc_call`'s return values on failure if: +/// * The error is rejected by `is_retriable`. +/// * The retry budget is exhausted. +pub async fn execute_with_retry< + ResponseType, + ErrorType, + GrpcCall: AsyncFnMut() -> Result, + RetriableCheck: Fn(&ErrorType) -> bool, +>( + max_retries: usize, + max_backoff: Duration, + mut grpc_call: GrpcCall, + is_retriable: RetriableCheck, +) -> Result { + let mut retry = 0usize; + loop { + let error = match grpc_call().await { + Ok(value) => return Ok(value), + Err(error) => error, + }; + if !is_retriable(&error) || retry >= max_retries { + return Err(error); + } + tokio::time::sleep(backoff(retry, max_backoff)).await; + retry += 1; + } +} + +/// Executes a gRPC call under the retry policy configured by `retry_config`, retrying only on +/// [`Code::Unavailable`]. +/// +/// # Type Parameters +/// +/// * `ResponseType` - The success value produced by `grpc_call`. +/// * `GrpcCall` - The async closure performing the gRPC round-trip. +/// +/// # Returns +/// +/// The first successful response returned by `grpc_call`. +/// +/// # Errors +/// +/// Returns an error if: +/// +/// * Forwards `grpc_call`'s [`Status`] on failure if the status is non-retriable or the retry +/// budget is exhausted. +pub async fn call_with_retry< + ResponseType, + GrpcCall: AsyncFnMut() -> Result, +>( + retry_config: RetryConfig, + grpc_call: GrpcCall, +) -> Result { + execute_with_retry( + retry_config.max_retries, + retry_config.max_backoff, + grpc_call, + is_retriable_status, + ) + .await +} + +/// Classifies a gRPC [`Status`] as retriable. +/// +/// # Returns +/// +/// Whether `status`'s code is [`Code::Unavailable`], the only code treated as retriable. +fn is_retriable_status(status: &Status) -> bool { + matches!(status.code(), Code::Unavailable) +} + +/// The default number of retries allowed after the initial attempt. +const DEFAULT_MAX_RETRIES: usize = 10; +/// The default upper bound on the exponential backoff between attempts. +const DEFAULT_MAX_BACKOFF: Duration = Duration::from_secs(3); + +/// Computes the delay before the `retry`-th retry. +/// +/// # Returns +/// +/// An exponentially increasing delay capped at `max_backoff`, plus a small random jitter; the +/// returned value may slightly exceed `max_backoff`. +fn backoff(retry: usize, max_backoff: Duration) -> Duration { + /// The backoff applied before the first retry, doubled on each subsequent retry. + const INITIAL_BACKOFF: Duration = Duration::from_millis(100); + /// The maximum random jitter, in milliseconds, added on top of the capped backoff. + const MAX_JITTER_MILLIS: u64 = 20; + + let capped = u32::try_from(retry) + .ok() + .and_then(|shift| 1u32.checked_shl(shift)) + .and_then(|multiplier| INITIAL_BACKOFF.checked_mul(multiplier)) + .unwrap_or(max_backoff) + .min(max_backoff); + let jitter = Duration::from_millis(rand::rng().random_range(0..=MAX_JITTER_MILLIS)); + capped.saturating_add(jitter) +} + +#[cfg(test)] +mod tests { + use std::cell::Cell; + use std::time::Duration; + + use tonic::Code; + use tonic::Status; + + use super::RetryConfig; + use super::backoff; + use super::call_with_retry; + use super::execute_with_retry; + + /// A negligible cap so the tests never sleep for a meaningful amount of time. + const TEST_MAX_BACKOFF: Duration = Duration::from_millis(1); + + #[tokio::test] + async fn succeeds_on_first_attempt() { + let calls = Cell::new(0usize); + let result: Result = execute_with_retry( + 3, + TEST_MAX_BACKOFF, + async || { + calls.set(calls.get() + 1); + Ok(42) + }, + |_error| true, + ) + .await; + + assert_eq!(result, Ok(42)); + assert_eq!(calls.get(), 1); + } + + #[tokio::test] + async fn succeeds_after_retriable_failures() { + let calls = Cell::new(0usize); + let result: Result = execute_with_retry( + 5, + TEST_MAX_BACKOFF, + async || { + let attempt = calls.get(); + calls.set(attempt + 1); + if attempt < 2 { Err(-1) } else { Ok(7) } + }, + |_error| true, + ) + .await; + + assert_eq!(result, Ok(7)); + assert_eq!(calls.get(), 3); + } + + #[tokio::test] + async fn non_retriable_error_returns_immediately() { + let calls = Cell::new(0usize); + let result: Result = execute_with_retry( + 3, + TEST_MAX_BACKOFF, + async || { + calls.set(calls.get() + 1); + Err(99) + }, + |error| *error != 99, + ) + .await; + + assert_eq!(result, Err(99)); + assert_eq!(calls.get(), 1); + } + + #[tokio::test] + async fn retries_are_exhausted() { + let max_retries = 4usize; + let calls = Cell::new(0usize); + let result: Result = execute_with_retry( + max_retries, + TEST_MAX_BACKOFF, + async || { + calls.set(calls.get() + 1); + Err(-7) + }, + |_error| true, + ) + .await; + + assert_eq!(result, Err(-7)); + assert_eq!(calls.get(), max_retries + 1); + } + + #[tokio::test] + async fn call_with_retry_retries_unavailable_then_succeeds() { + let config = RetryConfig { + max_retries: 5, + max_backoff: TEST_MAX_BACKOFF, + }; + let calls = Cell::new(0usize); + let result: Result = call_with_retry(config, async || { + let attempt = calls.get(); + calls.set(attempt + 1); + if attempt < 2 { + Err(Status::unavailable("connection lost")) + } else { + Ok(11) + } + }) + .await; + + assert_eq!( + result.expect("call_with_retry should succeed after retriable failures"), + 11 + ); + assert_eq!(calls.get(), 3); + } + + #[tokio::test] + async fn call_with_retry_returns_immediately_on_non_retriable_status() { + let config = RetryConfig { + max_retries: 5, + max_backoff: TEST_MAX_BACKOFF, + }; + let calls = Cell::new(0usize); + let result: Result = call_with_retry(config, async || { + calls.set(calls.get() + 1); + Err(Status::not_found("missing")) + }) + .await; + + assert_eq!( + result + .expect_err("call_with_retry should propagate a non-retriable status") + .code(), + Code::NotFound + ); + assert_eq!(calls.get(), 1); + } + + #[test] + fn backoff_stays_within_jitter_bounds() { + /// The expected initial backoff, mirroring the implementation's private constant. + const EXPECTED_INITIAL_BACKOFF: Duration = Duration::from_millis(100); + /// The expected maximum jitter, mirroring the implementation's private constant. + const EXPECTED_MAX_JITTER: Duration = Duration::from_millis(20); + + for (retry, max_backoff) in [ + (0usize, Duration::from_millis(1000)), + (3usize, Duration::from_millis(1000)), + (5usize, Duration::from_millis(1000)), + ] { + let capped_expected = (EXPECTED_INITIAL_BACKOFF * (1u32 << retry)).min(max_backoff); + for _ in 0..100 { + let actual = backoff(retry, max_backoff); + assert!( + actual >= capped_expected, + "backoff {actual:?} is below the expected floor {capped_expected:?}" + ); + assert!( + actual <= capped_expected + EXPECTED_MAX_JITTER, + "backoff {actual:?} exceeds the expected ceiling {:?}", + capped_expected + EXPECTED_MAX_JITTER + ); + } + } + } +} diff --git a/tests/huntsman/e2e/src/test_driver.rs b/tests/huntsman/e2e/src/test_driver.rs index ce8e87d8d..48dc08643 100644 --- a/tests/huntsman/e2e/src/test_driver.rs +++ b/tests/huntsman/e2e/src/test_driver.rs @@ -150,14 +150,17 @@ impl SpiderTestDriver { /// * The endpoint environment variable is unset. /// * The endpoint value is not a valid Spider endpoint. /// * Forwards [`read_concurrency`]'s return values on failure. - /// * Forwards [`SpiderClient::connect`]'s return values on failure. + /// * Forwards [`spider_client::SpiderClientBuilder::connect`]'s return values on failure. async fn init() -> anyhow::Result { const ENDPOINT_ENV_VAR: &str = "SPIDER_ENDPOINT"; let endpoint_string = std::env::var(ENDPOINT_ENV_VAR) .with_context(|| format!("{ENDPOINT_ENV_VAR} is not set"))?; let endpoint = Endpoint::from_shared(endpoint_string).context("invalid spider endpoint")?; let concurrency = read_concurrency()?; - let client = SpiderClient::connect(endpoint, concurrency).await?; + let client = SpiderClient::builder(endpoint) + .pool_size(concurrency) + .connect() + .await?; Ok(Self { client: RwLock::new(client), concurrency_limiter: Semaphore::new(concurrency.get()), From 7f8c4b3a1fa19688eca82c1310fd3d3213537afe Mon Sep 17 00:00:00 2001 From: LinZhihao-723 Date: Mon, 13 Jul 2026 22:21:23 -0400 Subject: [PATCH 2/3] Fix docstrings. --- components/spider-client/src/client.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/components/spider-client/src/client.rs b/components/spider-client/src/client.rs index 0f359f510..3879db990 100644 --- a/components/spider-client/src/client.rs +++ b/components/spider-client/src/client.rs @@ -266,7 +266,8 @@ impl SpiderClientBuilder { /// /// Returns an error if: /// - /// * [`ClientError::Transport`] if tonic cannot create or connect to the endpoint. + /// * Forwards [`JobOrchestrationClient::connect`]'s return values on failure. + /// * Forwards [`ResourceGroupManagementClient::connect`]'s return values on failure. pub async fn connect(self) -> Result { let (job_orchestration, resource_group) = tokio::try_join!( JobOrchestrationClient::connect( From fe1d8154297162b15156d1901aaa3179e43dceda Mon Sep 17 00:00:00 2001 From: LinZhihao-723 Date: Tue, 14 Jul 2026 20:00:05 -0400 Subject: [PATCH 3/3] Check Unknown status as well for retrying. --- components/spider-utils/src/grpc/retry.rs | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/components/spider-utils/src/grpc/retry.rs b/components/spider-utils/src/grpc/retry.rs index d3a49f1a8..b2fc0d4db 100644 --- a/components/spider-utils/src/grpc/retry.rs +++ b/components/spider-utils/src/grpc/retry.rs @@ -1,5 +1,6 @@ //! An async retry helper for transient gRPC call failures. +use std::error::Error; use std::time::Duration; use rand::Rng; @@ -119,9 +120,19 @@ pub async fn call_with_retry< /// /// # Returns /// -/// Whether `status`'s code is [`Code::Unavailable`], the only code treated as retriable. +/// Whether `status` is retriable, which holds when its code is: +/// +/// * [`Code::Unavailable`], the server's signal that the request may be retried; or +/// * [`Code::Unknown`] carrying a [`tonic::transport::Error`] source, which is how tonic surfaces a +/// client-side transport failure (such as a dropped connection) that is worth retrying. fn is_retriable_status(status: &Status) -> bool { - matches!(status.code(), Code::Unavailable) + match status.code() { + Code::Unavailable => true, + Code::Unknown => status + .source() + .is_some_and(|source| source.downcast_ref::().is_some()), + _ => false, + } } /// The default number of retries allowed after the initial attempt.