This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs-object-store.git
The following commit(s) were added to refs/heads/main by this push:
new c37db72 feat: retry failed multipart part uploads (#849)
c37db72 is described below
commit c37db72b92b78608291c1b33feb79bb5bb6d2d1c
Author: Chris <[email protected]>
AuthorDate: Wed Sep 9 11:43:46 2026 -0700
feat: retry failed multipart part uploads (#849)
* Retry multipart part uploads with configurable policies
* Fix retry feature gate for HTTP-only builds
* Remove Copy from RetryContext
---------
Co-authored-by: Andrew Lamb <[email protected]>
---
src/aws/mod.rs | 108 +++++++++++++++++++---
src/azure/client.rs | 16 +++-
src/azure/mod.rs | 28 +++++-
src/gcp/mod.rs | 25 +++--
src/lib.rs | 5 +-
src/retry.rs | 261 ++++++++++++++++++++++++++++++++++++++++++++++++++++
6 files changed, 415 insertions(+), 28 deletions(-)
diff --git a/src/aws/mod.rs b/src/aws/mod.rs
index dd6c05e..8fcb440 100644
--- a/src/aws/mod.rs
+++ b/src/aws/mod.rs
@@ -44,6 +44,7 @@ use crate::client::CredentialProvider;
use crate::client::get::GetClientExt;
use crate::client::list::{ListClient, ListClientExt};
use crate::multipart::{MultipartStore, PartId};
+use crate::retry::{MultipartRetry, RetryPolicy};
use crate::signer::{SignedUrlOptions, Signer};
use crate::util::{STRICT_ENCODE_SET, validate_signed_url_extras};
use crate::{
@@ -313,6 +314,7 @@ impl ObjectStore for AmazonS3 {
location: &Path,
opts: PutMultipartOptions,
) -> Result<Box<dyn MultipartUpload>> {
+ let retry_policy = opts.retry_policy();
let upload_id = self.client.create_multipart(location, opts).await?;
Ok(Box::new(S3MultiPartUpload {
@@ -322,6 +324,7 @@ impl ObjectStore for AmazonS3 {
location: location.clone(),
upload_id: upload_id.clone(),
parts: Default::default(),
+ retry_policy,
}),
}))
}
@@ -498,6 +501,7 @@ struct UploadState {
location: Path,
upload_id: String,
client: Arc<S3Client>,
+ retry_policy: Option<Arc<dyn RetryPolicy>>,
}
#[async_trait]
@@ -507,17 +511,26 @@ impl MultipartUpload for S3MultiPartUpload {
self.part_idx += 1;
let state = Arc::clone(&self.state);
Box::pin(async move {
- let part = state
- .client
- .put_part(
- &state.location,
- &state.upload_id,
- idx,
- PutPartPayload::Part(data),
- )
- .await?;
- state.parts.put(idx, part);
- Ok(())
+ let mut retry = MultipartRetry::new(state.retry_policy.clone());
+ loop {
+ match state
+ .client
+ .put_part(
+ &state.location,
+ &state.upload_id,
+ idx,
+ PutPartPayload::Part(data.clone()),
+ )
+ .await
+ {
+ Ok(part) => {
+ state.parts.put(idx, part);
+ return Ok(());
+ }
+ Err(error) if retry.should_retry(&error).await => continue,
+ Err(error) => return Err(error),
+ }
+ }
})
}
@@ -727,6 +740,8 @@ mod tests {
#[cfg(feature = "reqwest")]
use crate::client::SpawnedReqwestConnector;
use crate::client::get::GetClient;
+ #[cfg(feature = "reqwest")]
+ use crate::client::mock_server::MockServer;
use crate::client::retry::RetryContext;
use crate::integration::*;
use crate::tests::*;
@@ -734,6 +749,12 @@ mod tests {
use base64::prelude::BASE64_STANDARD;
use http::HeaderMap;
use http::HeaderValue;
+ #[cfg(feature = "reqwest")]
+ use http::Response;
+ #[cfg(feature = "reqwest")]
+ use http::header::ETAG;
+ #[cfg(feature = "reqwest")]
+ use std::sync::atomic::{AtomicUsize, Ordering};
const NON_EXISTENT_NAME: &str = "nonexistentname";
@@ -779,6 +800,71 @@ mod tests {
store
}
+ #[cfg(feature = "reqwest")]
+ #[derive(Debug, Default)]
+ struct RetryOnce(AtomicUsize);
+
+ #[cfg(feature = "reqwest")]
+ #[async_trait]
+ impl RetryPolicy for RetryOnce {
+ async fn retry(&self, _: crate::retry::RetryContext) -> bool {
+ self.0.fetch_add(1, Ordering::SeqCst) == 0
+ }
+ }
+
+ #[cfg(feature = "reqwest")]
+ #[tokio::test]
+ async fn retries_multipart_part_operation() {
+ let mock = MockServer::new().await;
+ mock.push(
+ Response::builder()
+ .status(StatusCode::OK)
+ .body(
+
"<InitiateMultipartUploadResult><UploadId>upload-id</UploadId></InitiateMultipartUploadResult>"
+ .to_string(),
+ )
+ .unwrap(),
+ );
+ mock.push(
+ Response::builder()
+ .status(StatusCode::SERVICE_UNAVAILABLE)
+ .body(String::new())
+ .unwrap(),
+ );
+ mock.push(
+ Response::builder()
+ .status(StatusCode::OK)
+ .header(ETAG, "etag")
+ .body(String::new())
+ .unwrap(),
+ );
+
+ let store = AmazonS3Builder::new()
+ .with_endpoint(mock.url())
+ .with_bucket_name("test-bucket")
+ .with_region("us-east-1")
+ .with_allow_http(true)
+ .with_skip_signature(true)
+ .with_retry(crate::RetryConfig {
+ max_retries: 0,
+ ..Default::default()
+ })
+ .build()
+ .unwrap();
+ let policy = Arc::new(RetryOnce::default());
+ let opts = PutMultipartOptions::default()
+ .with_retry_policy(Arc::clone(&policy) as Arc<dyn RetryPolicy>);
+
+ let mut upload = store
+ .put_multipart_opts(&Path::from("multipart"), opts)
+ .await
+ .unwrap();
+ upload.put_part(PutPayload::from("data")).await.unwrap();
+
+ assert_eq!(policy.0.load(Ordering::SeqCst), 1);
+ mock.shutdown().await;
+ }
+
#[tokio::test]
async fn write_multipart_file_with_signature() {
maybe_skip_integration!();
diff --git a/src/azure/client.rs b/src/azure/client.rs
index c5b0a52..54da8c1 100644
--- a/src/azure/client.rs
+++ b/src/azure/client.rs
@@ -780,16 +780,20 @@ impl AzureClient {
)
}
+ /// Generate the identity of a block in a multipart upload
+ pub(crate) fn new_block_id() -> String {
+ let block_id = u128::from_be_bytes(rand::rng().random());
+ format!("{block_id:032x}")
+ }
+
/// PUT a block
<https://learn.microsoft.com/en-us/rest/api/storageservices/put-block>
pub(crate) async fn put_block(
&self,
path: &Path,
- _part_idx: usize,
+ content_id: &str,
payload: PutPayload,
) -> Result<PartId> {
- let part_idx = u128::from_be_bytes(rand::rng().random());
- let content_id = format!("{part_idx:032x}");
- let block_id = BASE64_STANDARD.encode(&content_id);
+ let block_id = BASE64_STANDARD.encode(content_id);
self.put_request(path, payload)
.query(&[("comp", "block"), ("blockid", &block_id)])
@@ -797,7 +801,9 @@ impl AzureClient {
.send()
.await?;
- Ok(PartId { content_id })
+ Ok(PartId {
+ content_id: content_id.into(),
+ })
}
/// PUT a block list
<https://learn.microsoft.com/en-us/rest/api/storageservices/put-block-list>
diff --git a/src/azure/mod.rs b/src/azure/mod.rs
index 91e4a96..e13fe0a 100644
--- a/src/azure/mod.rs
+++ b/src/azure/mod.rs
@@ -45,6 +45,7 @@ use url::Url;
use crate::client::get::GetClientExt;
use crate::client::list::{ListClient, ListClientExt};
use crate::client::{CredentialProvider, crypto_provider};
+use crate::retry::{MultipartRetry, RetryPolicy};
pub use credential::{AzureAccessKey, AzureAuthorizer, authority_hosts};
mod builder;
@@ -106,6 +107,7 @@ impl ObjectStore for MicrosoftAzure {
location: &Path,
opts: PutMultipartOptions,
) -> Result<Box<dyn MultipartUpload>> {
+ let retry_policy = opts.retry_policy();
Ok(Box::new(AzureMultiPartUpload {
part_idx: 0,
opts,
@@ -113,6 +115,7 @@ impl ObjectStore for MicrosoftAzure {
client: Arc::clone(&self.client),
location: location.clone(),
parts: Default::default(),
+ retry_policy,
}),
}))
}
@@ -279,6 +282,7 @@ struct UploadState {
location: Path,
parts: Parts,
client: Arc<AzureClient>,
+ retry_policy: Option<Arc<dyn RetryPolicy>>,
}
#[async_trait]
@@ -287,10 +291,23 @@ impl MultipartUpload for AzureMultiPartUpload {
let idx = self.part_idx;
self.part_idx += 1;
let state = Arc::clone(&self.state);
+ let content_id = AzureClient::new_block_id();
Box::pin(async move {
- let part = state.client.put_block(&state.location, idx,
data).await?;
- state.parts.put(idx, part);
- Ok(())
+ let mut retry = MultipartRetry::new(state.retry_policy.clone());
+ loop {
+ match state
+ .client
+ .put_block(&state.location, &content_id, data.clone())
+ .await
+ {
+ Ok(part) => {
+ state.parts.put(idx, part);
+ return Ok(());
+ }
+ Err(error) if retry.should_retry(&error).await => continue,
+ Err(error) => return Err(error),
+ }
+ }
})
}
@@ -392,10 +409,11 @@ impl MultipartStore for MicrosoftAzure {
&self,
path: &Path,
_: &MultipartId,
- part_idx: usize,
+ _: usize,
data: PutPayload,
) -> Result<PartId> {
- self.client.put_block(path, part_idx, data).await
+ let content_id = AzureClient::new_block_id();
+ self.client.put_block(path, &content_id, data).await
}
async fn complete_multipart(
diff --git a/src/gcp/mod.rs b/src/gcp/mod.rs
index 2dbb0c7..d33b7e8 100644
--- a/src/gcp/mod.rs
+++ b/src/gcp/mod.rs
@@ -43,6 +43,7 @@ use std::time::Duration;
use crate::CopyOptions;
use crate::client::{CredentialProvider, crypto_provider};
use crate::gcp::credential::GCSAuthorizer;
+use crate::retry::{MultipartRetry, RetryPolicy};
use crate::signer::{SignedUrlOptions, Signer};
use crate::util::validate_signed_url_extras;
use crate::{
@@ -117,6 +118,7 @@ struct UploadState {
path: Path,
multipart_id: MultipartId,
parts: Parts,
+ retry_policy: Option<Arc<dyn RetryPolicy>>,
}
#[async_trait]
@@ -126,12 +128,21 @@ impl MultipartUpload for GCSMultipartUpload {
self.part_idx += 1;
let state = Arc::clone(&self.state);
Box::pin(async move {
- let part = state
- .client
- .put_part(&state.path, &state.multipart_id, idx, payload)
- .await?;
- state.parts.put(idx, part);
- Ok(())
+ let mut retry = MultipartRetry::new(state.retry_policy.clone());
+ loop {
+ match state
+ .client
+ .put_part(&state.path, &state.multipart_id, idx,
payload.clone())
+ .await
+ {
+ Ok(part) => {
+ state.parts.put(idx, part);
+ return Ok(());
+ }
+ Err(error) if retry.should_retry(&error).await => continue,
+ Err(error) => return Err(error),
+ }
+ }
})
}
@@ -168,6 +179,7 @@ impl ObjectStore for GoogleCloudStorage {
location: &Path,
opts: PutMultipartOptions,
) -> Result<Box<dyn MultipartUpload>> {
+ let retry_policy = opts.retry_policy();
let upload_id = self.client.multipart_initiate(location, opts).await?;
Ok(Box::new(GCSMultipartUpload {
@@ -177,6 +189,7 @@ impl ObjectStore for GoogleCloudStorage {
path: location.clone(),
multipart_id: upload_id.clone(),
parts: Default::default(),
+ retry_policy,
}),
}))
}
diff --git a/src/lib.rs b/src/lib.rs
index 1dd749e..57713be 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -727,6 +727,8 @@ pub mod memory;
pub mod path;
pub mod prefix;
pub mod registry;
+#[cfg(any(feature = "aws-base", feature = "azure-base", feature = "gcp-base"))]
+pub mod retry;
#[cfg(feature = "cloud-base")]
pub mod signer;
#[cfg(feature = "tokio")]
@@ -2023,7 +2025,8 @@ pub struct PutMultipartOptions {
/// Implementation-specific extensions. Intended for use by
[`ObjectStore`] implementations
/// that need to pass context-specific information (like tracing spans)
via trait methods.
///
- /// These extensions are ignored entirely by backends offered through this
crate.
+ /// Cloud backends offered through this crate use extensions installed by
methods such as
+ /// `PutMultipartOptions::with_retry_policy`. Other extensions are ignored.
///
/// They are also excluded from [`PartialEq`] and [`Eq`].
pub extensions: Extensions,
diff --git a/src/retry.rs b/src/retry.rs
new file mode 100644
index 0000000..2aa46e6
--- /dev/null
+++ b/src/retry.rs
@@ -0,0 +1,261 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Retry policies for HTTP-backed multipart uploads
+//!
+//! Each multipart part first uses the backend's existing bounded request retry
+//! loop. If those retries are exhausted, a [`RetryPolicy`] can restart the
+//! entire part operation. Restarting at this boundary rebuilds the request,
+//! including fetching credentials and generating a new signature.
+//!
+//! Configure a policy with [`PutMultipartOptions::with_retry_policy`]. The
+//! policy applies independently to each part and does not change the retry
+//! behavior of multipart initiation, completion, or abort requests.
+
+use crate::client::retry::{RequestError, RetryError};
+use crate::client::{HttpError, HttpErrorKind};
+use crate::{Error, PutMultipartOptions};
+use async_trait::async_trait;
+use http::StatusCode;
+use std::error::Error as StdError;
+use std::fmt::Debug;
+use std::sync::Arc;
+use std::time::Duration;
+#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
+use std::time::Instant;
+#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
+use web_time::Instant;
+
+/// A normalized HTTP failure passed to a [`RetryPolicy`]
+///
+/// Backends wrap request failures in provider-specific [`Error`] variants.
+/// Before invoking a policy, `object_store` reduces a supported failure to an
+/// HTTP response status or transport error kind. Failures that cannot be
+/// represented by this enum are returned without invoking the policy.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[non_exhaustive]
+pub enum RetryFailure {
+ /// The HTTP status of a response that the backend classified as failed
+ ///
+ /// This is usually a non-success status. Some services, including S3, can
+ /// report an error in the body of a successful response. In that case this
+ /// contains the actual successful status, such as `200 OK`.
+ Status(StatusCode),
+
+ /// The request failed before producing a usable HTTP response
+ ///
+ /// The kind distinguishes failures such as connection errors, timeouts,
+ /// interrupted requests, and response decoding errors.
+ Transport(HttpErrorKind),
+}
+
+/// The failure and retry state for one operation attempt
+///
+/// For multipart uploads, an attempt is one complete call to upload a part,
+/// including its bounded request retries. Each part has an independent attempt
+/// count and elapsed time.
+//
+// Note: deliberately does not implement `Copy` so that non-`Copy` details,
+// such as the underlying error, can be added in the future without a
+// breaking change.
+#[derive(Debug, Clone, PartialEq, Eq)]
+#[non_exhaustive]
+pub struct RetryContext {
+ /// The normalized HTTP failure
+ pub failure: RetryFailure,
+
+ /// The failed operation attempt number, starting at one
+ pub attempt: usize,
+
+ /// Time since the first operation attempt started
+ ///
+ /// This includes time spent in requests and in earlier policy calls.
+ pub elapsed: Duration,
+}
+
+/// Controls whether and when a failed operation should be attempted again
+///
+/// The policy is shared by concurrently uploaded parts. Implementations must
+/// manage their own retry limit and backoff. `object_store` does not add a
delay
+/// or impose an outer retry limit after this method returns `true`.
+#[async_trait]
+pub trait RetryPolicy: Debug + Send + Sync + 'static {
+ /// Decide whether to attempt the failed operation again
+ ///
+ /// Because this method is asynchronous, implementations can wait using
+ /// their own backoff strategy, clock, or sleeper before returning `true`.
+ /// Returning `false` returns the current operation error to the caller.
+ async fn retry(&self, context: RetryContext) -> bool;
+}
+
+#[derive(Clone)]
+struct MultipartRetryPolicy(Arc<dyn RetryPolicy>);
+
+impl PutMultipartOptions {
+ /// Retry failed multipart part uploads according to `policy`
+ ///
+ /// The policy runs after the existing bounded request retries are
+ /// exhausted. A retry repeats the same logical part from the provider
+ /// operation boundary, allowing credentials and request signatures to be
+ /// refreshed. Initiating, completing, and aborting the multipart upload
+ /// retain their existing retry behavior.
+ #[must_use]
+ pub fn with_retry_policy(mut self, policy: Arc<dyn RetryPolicy>) -> Self {
+ self.extensions.insert(MultipartRetryPolicy(policy));
+ self
+ }
+
+ pub(crate) fn retry_policy(&self) -> Option<Arc<dyn RetryPolicy>> {
+ self.extensions
+ .get::<MultipartRetryPolicy>()
+ .map(|policy| Arc::clone(&policy.0))
+ }
+}
+
+pub(crate) struct MultipartRetry {
+ policy: Option<Arc<dyn RetryPolicy>>,
+ attempt: usize,
+ start: Instant,
+}
+
+impl MultipartRetry {
+ pub(crate) fn new(policy: Option<Arc<dyn RetryPolicy>>) -> Self {
+ Self {
+ policy,
+ attempt: 0,
+ start: Instant::now(),
+ }
+ }
+
+ pub(crate) async fn should_retry(&mut self, error: &Error) -> bool {
+ let Some(policy) = self.policy.as_ref() else {
+ return false;
+ };
+ let Some(failure) = classify_http_failure(error) else {
+ return false;
+ };
+
+ self.attempt += 1;
+ policy
+ .retry(RetryContext {
+ failure,
+ attempt: self.attempt,
+ elapsed: self.start.elapsed(),
+ })
+ .await
+ }
+}
+
+/// Find and normalize an HTTP failure in an [`Error`] source chain.
+///
+/// Provider errors may wrap a terminal [`RetryError`] or [`HttpError`] several
+/// levels deep. Returns `None` when the chain contains no failure
representable
+/// by [`RetryFailure`].
+fn classify_http_failure(error: &Error) -> Option<RetryFailure> {
+ let mut current: &(dyn StdError + 'static) = error;
+ loop {
+ if let Some(error) = current.downcast_ref::<RetryError>() {
+ return classify_request_error(error.inner());
+ }
+ if let Some(error) = current.downcast_ref::<HttpError>() {
+ return Some(RetryFailure::Transport(error.kind()));
+ }
+ current = current.source()?;
+ }
+}
+
+/// Normalize a terminal request error for a [`RetryPolicy`].
+///
+/// Error responses retain their actual status, including successful statuses
+/// whose bodies contain a provider error. A bare redirect has no status or
+/// transport kind to expose and is therefore not policy-controlled.
+fn classify_request_error(error: &RequestError) -> Option<RetryFailure> {
+ match error {
+ RequestError::Status { status, .. } | RequestError::Response { status,
.. } => {
+ Some(RetryFailure::Status(*status))
+ }
+ RequestError::Http(error) =>
Some(RetryFailure::Transport(error.kind())),
+ RequestError::BareRedirect => None,
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use parking_lot::Mutex;
+
+ #[derive(Debug, Default)]
+ struct RecordingPolicy(Mutex<Vec<RetryContext>>);
+
+ #[async_trait]
+ impl RetryPolicy for RecordingPolicy {
+ async fn retry(&self, context: RetryContext) -> bool {
+ self.0.lock().push(context);
+ true
+ }
+ }
+
+ #[derive(Debug, thiserror::Error)]
+ #[error("test error")]
+ struct TestError;
+
+ #[tokio::test]
+ async fn retries_http_errors() {
+ let policy = Arc::new(RecordingPolicy::default());
+ let mut retry = MultipartRetry::new(Some(Arc::clone(&policy) as
Arc<dyn RetryPolicy>));
+ let error = Error::Generic {
+ store: "test",
+ source: Box::new(HttpError::new(HttpErrorKind::Timeout,
TestError)),
+ };
+
+ assert!(retry.should_retry(&error).await);
+
+ let contexts = policy.0.lock();
+ assert_eq!(contexts.len(), 1);
+ assert_eq!(contexts[0].attempt, 1);
+ assert_eq!(
+ contexts[0].failure,
+ RetryFailure::Transport(HttpErrorKind::Timeout)
+ );
+ }
+
+ #[tokio::test]
+ async fn does_not_retry_non_http_errors() {
+ let policy = Arc::new(RecordingPolicy::default());
+ let mut retry = MultipartRetry::new(Some(Arc::clone(&policy) as
Arc<dyn RetryPolicy>));
+ let error = Error::Generic {
+ store: "test",
+ source: Box::new(TestError),
+ };
+
+ assert!(!retry.should_retry(&error).await);
+ assert!(policy.0.lock().is_empty());
+ }
+
+ #[test]
+ fn classifies_error_responses_with_success_status() {
+ let error = RequestError::Response {
+ status: StatusCode::OK,
+ body: "InternalError".into(),
+ };
+
+ assert_eq!(
+ classify_request_error(&error),
+ Some(RetryFailure::Status(StatusCode::OK))
+ );
+ }
+}
