Skip to content

feat(spider-client)!: Retry gRPC connection failures with a configurable exponential-backoff policy. #396

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
LinZhihao-723 merged 4 commits into from
Jul 15, 2026
Merged

feat(spider-client)!: Retry gRPC connection failures with a configurable exponential-backoff policy. #396

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
4d49222
Implementation done.
LinZhihao-723 Jul 14, 2026
2112177
Merge branch 'oss-main' into retry
LinZhihao-723 Jul 14, 2026
7f8c4b3
Fix docstrings.
LinZhihao-723 Jul 14, 2026
fe1d815
Check Unknown status as well for retrying.
LinZhihao-723 Jul 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

104 changes: 86 additions & 18 deletions components/spider-client/src/client.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,13 +1,15 @@
//! [`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;
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;
Expand All @@ -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<Self, ClientError> {
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
Expand Down Expand Up @@ -222,3 +214,79 @@ 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:
///
/// * Forwards [`JobOrchestrationClient::connect`]'s return values on failure.
/// * Forwards [`ResourceGroupManagementClient::connect`]'s return values on failure.
pub async fn connect(self) -> Result<SpiderClient, ClientError> {
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();
104 changes: 60 additions & 44 deletions components/spider-client/src/grpc/job.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -27,6 +29,7 @@ use crate::error::to_transport_error;
#[derive(Debug, Clone)]
pub struct JobOrchestrationClient {
connection_pool: ConnectionPool<JobOrchestrationServiceClient<Channel>>,
retry_config: RetryConfig,
}

impl JobOrchestrationClient {
Expand All @@ -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<Self, ClientError> {
pub async fn connect(
endpoint: Endpoint,
pool_size: NonZeroUsize,
retry_config: RetryConfig,
) -> Result<Self, ClientError> {
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
Expand Down Expand Up @@ -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();
Comment thread
coderabbitai[bot] marked this conversation as resolved.

Ok(JobId::from(response.job_id))
}
Expand All @@ -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)
}
Expand All @@ -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)
}
Expand All @@ -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)
}
Expand All @@ -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()))
Expand All @@ -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)
}
Expand Down
Loading
Loading