From 96bf6197683562e75446af845cb21cd82f740cd1 Mon Sep 17 00:00:00 2001 From: lcian <17258265+lcian@users.noreply.github.com> Date: Wed, 21 Oct 2026 13:17:27 +0200 Subject: [PATCH 1/4] feat(rust): Add automatic resumable uploads Resume eligible buffer and file uploads from authoritative server progress, with configurable thresholds and temporary-file compression. Keep automatic uploads available without the manual protocol feature. --- Cargo.lock | 2 + clients/rust/Cargo.toml | 6 +- clients/rust/README.md | 18 +- clients/rust/src/client.rs | 33 ++ clients/rust/src/error.rs | 9 +- clients/rust/src/lib.rs | 1 - clients/rust/src/many.rs | 94 +++-- clients/rust/src/put.rs | 121 +++++- clients/rust/src/put/automatic.rs | 334 ++++++++++++++++ clients/rust/src/resumable.rs | 34 +- clients/rust/tests/automatic.rs | 619 ++++++++++++++++++++++++++++++ clients/rust/tests/e2e.rs | 268 +++++++------ 12 files changed, 1382 insertions(+), 157 deletions(-) create mode 100644 clients/rust/src/put/automatic.rs create mode 100644 clients/rust/tests/automatic.rs diff --git a/Cargo.lock b/Cargo.lock index 7574e8aa..9ac5cc93 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2822,6 +2822,7 @@ version = "0.2.10" dependencies = [ "anyhow", "async-compression", + "axum", "base64", "bytes", "futures-util", @@ -2831,6 +2832,7 @@ dependencies = [ "multer", "objectstore-test", "objectstore-types", + "rand 0.10.2", "reqwest 0.13.4", "sentry-core", "serde", diff --git a/clients/rust/Cargo.toml b/clients/rust/Cargo.toml index 7d5e2d3a..894e60c0 100644 --- a/clients/rust/Cargo.toml +++ b/clients/rust/Cargo.toml @@ -22,24 +22,26 @@ infer = { workspace = true } jsonwebtoken = { workspace = true } multer = { workspace = true } objectstore-types = { workspace = true } +rand = { workspace = true } # Pinned below workspace version for MSRV compatibility reqwest = { version = "0.13.1", default-features = false, features = ["charset", "http2", "system-proxy", "json", "stream", "multipart", "query"] } sentry-core = { version = ">=0.41", default-features = false, features = ["client"] } serde = { workspace = true } +tempfile = { workspace = true } thiserror = { workspace = true } -tokio = { version = "1.45.0", default-features = false, features = ["fs"] } +tokio = { version = "1.45.0", default-features = false, features = ["fs", "io-util", "rt", "sync", "time"] } tokio-util = { workspace = true, features = ["io"] } url = { workspace = true } zstd-safe = { workspace = true } [dev-dependencies] anyhow = { workspace = true } +axum = { workspace = true } futures-util = { workspace = true } http-body-util = { workspace = true } jsonwebtoken = { workspace = true } objectstore-test = { workspace = true } serde = { workspace = true } -tempfile = { workspace = true } tokio = { workspace = true, features = ["macros", "rt"] } zstd = { workspace = true } diff --git a/clients/rust/README.md b/clients/rust/README.md index c48651c1..317fa224 100644 --- a/clients/rust/README.md +++ b/clients/rust/README.md @@ -151,6 +151,17 @@ session.put("payload") .send().await?; ``` +### Automatic Resumable Uploads + +Large buffer and regular-file uploads automatically use resumable uploads, including uploads +queued through `many()`. The default threshold is 256 MiB before client compression. Compressed +payloads are prepared in temporary files; uncompressed and precompressed inputs are reused. +Streams and generic readers use direct uploads. No feature flag is required. + +See [`PutBuilder`] for eligibility, fallback, and recovery behavior, +[`Usecase::with_resumable_threshold_bytes`] and [`PutBuilder::resumable_threshold_bytes`] to +configure or disable it, and [`ClientBuilder::spool_dir`] to choose the temporary-file directory. + ### Resumable Upload API > **Feature flag required:** Enable `resumable-upload-api` to use this API. @@ -176,7 +187,8 @@ If a non-final chunk is larger than the granularity but its size is not a multip server may persist only its aligned prefix. ```rust,no_run -#![cfg(feature = "resumable-upload-api")] +# #[cfg(feature = "resumable-upload-api")] +# mod example { use bytes::Bytes; use anyhow::{Context as _, Result}; @@ -189,7 +201,7 @@ async fn upload_large_object(session: &Session, object: Bytes) -> Result<()> { .content_type("application/octet-stream") .send() .await? - // Fall back to a regular [`Session:;put`]. + // Fall back to a regular Session::put. .context("resumable upload declined for this object")?; let mut offset = 0; @@ -208,6 +220,8 @@ async fn upload_large_object(session: &Session, object: Bytes) -> Result<()> { }; } } +# } +# fn main() {} ``` ### Multipart Upload API diff --git a/clients/rust/src/client.rs b/clients/rust/src/client.rs index 64eb0cd1..3708cbb8 100644 --- a/clients/rust/src/client.rs +++ b/clients/rust/src/client.rs @@ -1,4 +1,5 @@ use std::io; +use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; @@ -21,6 +22,7 @@ struct ClientBuilderInner { propagate_traces: bool, reqwest_builder: reqwest::ClientBuilder, token: Option, + spool_dir: Option, } impl ClientBuilderInner { @@ -71,6 +73,7 @@ impl ClientBuilder { propagate_traces: false, reqwest_builder, token: None, + spool_dir: None, })) } @@ -124,6 +127,16 @@ impl ClientBuilder { Self(Ok(inner)) } + /// Sets the directory for temporary files used to prepare compressed resumable uploads. + /// + /// By default, the system temporary directory is used. The directory must already exist. + /// If a temporary file cannot be created, uploads fall back to direct streaming. + pub fn spool_dir(self, path: impl Into) -> Self { + let Ok(mut inner) = self.0 else { return self }; + inner.spool_dir = Some(path.into()); + Self(Ok(inner)) + } + /// Returns a [`Client`] that uses this [`ClientBuilder`] configuration. /// /// # Errors @@ -141,6 +154,7 @@ impl ClientBuilder { service_url: inner.service_url, propagate_traces: inner.propagate_traces, token: inner.token, + spool_dir: inner.spool_dir, }), }) } @@ -161,6 +175,7 @@ pub struct Usecase { name: Arc, compression: Option, expiration_policy: ExpirationPolicy, + resumable_threshold_bytes: Option, } impl Usecase { @@ -188,6 +203,7 @@ impl Usecase { name: Arc::new(UsecaseName::new(name)?), compression: Some(Compression::Zstd), expiration_policy: ExpirationPolicy::default(), + resumable_threshold_bytes: Some(256 * 1024 * 1024), }) } @@ -234,6 +250,22 @@ impl Usecase { } } + /// Returns the minimum remaining input size for automatic resumable uploads. + pub fn resumable_threshold_bytes(&self) -> Option { + self.resumable_threshold_bytes + } + + /// Sets the minimum remaining input size for automatic resumable uploads. + /// + /// Defaults to 256 MiB, measured before client compression. Pass [`None`] to disable + /// automatic resumables or zero to include every supported payload. This also applies to + /// uploads queued through [`Session::many`]. See [`crate::PutBuilder`] for eligibility. + /// Individual uploads can override this using [`crate::PutBuilder::resumable_threshold_bytes`]. + pub fn with_resumable_threshold_bytes(mut self, threshold: impl Into>) -> Self { + self.resumable_threshold_bytes = threshold.into(); + self + } + /// Creates a new custom [`Scope`]. /// /// Add parts to it using [`Scope::push`]. @@ -325,6 +357,7 @@ impl Scope { #[derive(Debug)] pub(crate) struct ClientInner { + pub(crate) spool_dir: Option, reqwest: reqwest::Client, service_url: Url, propagate_traces: bool, diff --git a/clients/rust/src/error.rs b/clients/rust/src/error.rs index 1ccf46d9..8dfcbe37 100644 --- a/clients/rust/src/error.rs +++ b/clients/rust/src/error.rs @@ -69,19 +69,22 @@ pub enum Error { /// The error message. message: String, }, + /// An automatic resumable upload failed after session creation. + /// + /// The source preserves the failure that caused the upload to stop. Cancellation errors + /// are ignored so they cannot replace that failure. + #[error("upload failed: {0}")] + UploadFailed(#[source] Box), /// A resumable upload session expired, was canceled, or could otherwise not be found. - #[cfg(feature = "resumable-upload-api")] #[error("resumable upload session is not available")] ResumableUploadUnavailable, /// A resumable chunk starts at a different offset than the server currently holds. - #[cfg(feature = "resumable-upload-api")] #[error("upload offset mismatch (server holds {offset} bytes)")] UploadOffsetMismatch { /// The server's authoritative upload offset in bytes. offset: u64, }, /// A non-final chunk is shorter than the upload granularity. - #[cfg(feature = "resumable-upload-api")] #[error( "non-final chunk {chunk_length} is smaller than upload granularity {upload_granularity}" )] diff --git a/clients/rust/src/lib.rs b/clients/rust/src/lib.rs index 32649572..83db3369 100644 --- a/clients/rust/src/lib.rs +++ b/clients/rust/src/lib.rs @@ -15,7 +15,6 @@ mod many; mod multipart; mod put; mod response; -#[cfg(feature = "resumable-upload-api")] mod resumable; pub mod utils; diff --git a/clients/rust/src/many.rs b/clients/rust/src/many.rs index 0f9b12a7..d1b6d3f9 100644 --- a/clients/rust/src/many.rs +++ b/clients/rust/src/many.rs @@ -84,6 +84,7 @@ enum BatchOperation { metadata: Metadata, compression: Option, body: PutBody, + resumable_threshold_bytes: Option, }, Delete { key: ObjectKey, @@ -116,6 +117,7 @@ impl From for BatchOperation { metadata, compression, body, + resumable_threshold_bytes, session: _session, } = value; BatchOperation::Insert { @@ -123,6 +125,7 @@ impl From for BatchOperation { metadata, compression, body, + resumable_threshold_bytes, } } } @@ -159,6 +162,7 @@ impl BatchOperation { mut metadata, compression, body, + resumable_threshold_bytes: _, } => { let mut headers = operation_headers("insert", key.as_deref()); metadata.compression = compression.map(CompressionMode::compression); @@ -531,17 +535,9 @@ async fn send_batch( Ok(results) } -fn classify_fail(key: Option, error: Error) -> Classified { - Classified::Failed(OperationResult::Put( - key.unwrap_or_else(|| "".to_owned()), - Err(error), - )) -} - /// Classifies a single operation for batch processing. /// -/// Insert operations whose body exceeds [`MAX_BATCH_PART_SIZE`] are marked as -/// [`Classified::Individual`]. Everything else is [`Classified::Batchable`]. +/// Oversized, unknown-size, and resumable-eligible uploads are sent individually. async fn classify(op: BatchOperation) -> Classified { match op { BatchOperation::Insert { @@ -549,20 +545,22 @@ async fn classify(op: BatchOperation) -> Classified { metadata, compression, body, + resumable_threshold_bytes, } => { - let size = match &body { - PutBody::Buffer(bytes) => Some(bytes.len() as u64), - PutBody::File(file) => match file.metadata().await { - Ok(meta) => Some(meta.len()), - Err(err) => return classify_fail(key, err.into()), - }, - PutBody::Path(path) => match tokio::fs::metadata(path).await { - Ok(meta) => Some(meta.len()), - Err(err) => return classify_fail(key, err.into()), - }, - // Streams have unknown size and must not go through the batch endpoint. - PutBody::Stream(_) => None, + let mut body = body; + let size = match body.remaining_size().await { + Ok(size) => size, + Err(error) if matches!(body, PutBody::Path(_)) => { + return Classified::Failed(error_result( + OperationContext::Insert { key }, + error.into(), + )); + } + Err(_) => None, }; + let resumable = size + .zip(resumable_threshold_bytes) + .is_some_and(|(size, threshold)| size >= threshold); // Client-side compression happens while streaming the body, so the on-wire size is // only known up to its worst case. Precompressed bodies are sent as-is. @@ -581,8 +579,12 @@ async fn classify(op: BatchOperation) -> Classified { metadata, compression, body, + resumable_threshold_bytes, }; + if resumable { + return Classified::Individual(op); + } match size { Some(s) if s <= MAX_BATCH_PART_SIZE as usize => Classified::Batchable(op, s as u64), _ => Classified::Individual(op), @@ -637,6 +639,7 @@ async fn execute_individual(op: BatchOperation, session: &Session) -> OperationR metadata, compression, body, + resumable_threshold_bytes, } => { let error_key = key.clone().unwrap_or_else(|| "".to_owned()); let put = PutBuilder { @@ -645,6 +648,7 @@ async fn execute_individual(op: BatchOperation, session: &Session) -> OperationR compression, key, body, + resumable_threshold_bytes, }; match put.send().await { Ok(response) => OperationResult::Put(response.key.clone(), Ok(response)), @@ -760,7 +764,8 @@ impl ManyBuilder { /// Sets the maximum number of concurrent individual (non-batch) requests. /// - /// Operations that exceed the per-part size limit are sent as individual requests. + /// Operations that exceed the per-part size limit or qualify for resumable uploads + /// are sent as individual requests. /// This controls how many such requests can be in-flight simultaneously. /// Defaults to 5 if not set. pub fn max_individual_concurrency(mut self, concurrency: usize) -> Self { @@ -822,6 +827,7 @@ mod tests { metadata: Metadata::default(), compression: Some(compression), body: PutBody::Buffer(vec![0; size].into()), + resumable_threshold_bytes: None, } } @@ -867,6 +873,50 @@ mod tests { ); } + #[tokio::test] + async fn threshold_overrides_control_batch_classification() { + let session = crate::Client::builder("http://localhost") + .build() + .unwrap() + .session( + crate::Usecase::new("testing") + .with_resumable_threshold_bytes(4) + .for_organization(123), + ) + .unwrap(); + assert!(matches!( + classify(session.put("data").into()).await, + Classified::Individual(_) + )); + assert!(matches!( + classify(session.put("data").resumable_threshold_bytes(None).into()).await, + Classified::Batchable(_, _) + )); + assert!(matches!( + classify(session.put("data").resumable_threshold_bytes(5).into()).await, + Classified::Batchable(_, _) + )); + let file = tempfile::NamedTempFile::new().unwrap(); + std::fs::write(file.path(), b"prefix-data").unwrap(); + let mut input = tokio::fs::File::open(file.path()).await.unwrap(); + tokio::io::AsyncSeekExt::seek(&mut input, std::io::SeekFrom::Start(8)) + .await + .unwrap(); + assert!(matches!( + classify(session.put_file(input).into()).await, + Classified::Batchable(_, _) + )); + let Classified::Individual(BatchOperation::Insert { body, .. }) = + classify(session.put_path(file.path()).into()).await + else { + panic!("eligible path should be sent individually"); + }; + assert!( + matches!(body, PutBody::Path(_)), + "classification must not open paths" + ); + } + #[test] fn iter_batches_empty() { assert!(batches(vec![]).is_empty()); diff --git a/clients/rust/src/put.rs b/clients/rust/src/put.rs index fdfbcd3b..04221eca 100644 --- a/clients/rust/src/put.rs +++ b/clients/rust/src/put.rs @@ -10,7 +10,7 @@ use objectstore_types::metadata::Metadata; use reqwest::Body; use serde::Deserialize; use tokio::fs::File; -use tokio::io::{AsyncRead, BufReader}; +use tokio::io::{AsyncRead, AsyncSeekExt, AsyncWriteExt, BufReader}; use tokio_util::io::{ReaderStream, StreamReader}; pub use objectstore_types::metadata::{Compression, ExpirationPolicy}; @@ -18,6 +18,8 @@ pub use objectstore_types::metadata::{Compression, ExpirationPolicy}; use crate::response::ResponseExt as _; use crate::{ClientStream, ObjectKey, Session}; +mod automatic; + /// The response returned from the service after uploading an object. #[derive(Debug, Deserialize)] pub struct PutResponse { @@ -32,6 +34,30 @@ pub(crate) enum PutBody { Path(PathBuf), } +impl PutBody { + /// Measures remaining bytes without opening queued paths or consuming a stream. + pub(crate) async fn remaining_size(&mut self) -> io::Result> { + match self { + Self::Buffer(bytes) => Ok(Some(bytes.len() as u64)), + Self::File(file) => { + file.flush().await?; + let metadata = file.metadata().await?; + if !metadata.is_file() { + return Ok(None); + } + Ok(Some( + metadata.len().saturating_sub(file.stream_position().await?), + )) + } + Self::Path(path) => { + let metadata = tokio::fs::metadata(path).await?; + Ok(metadata.is_file().then_some(metadata.len())) + } + Self::Stream(_) => Ok(None), + } + } +} + impl fmt::Debug for PutBody { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_tuple("PutBody").finish_non_exhaustive() @@ -75,6 +101,7 @@ impl Session { .compression() .map(CompressionMode::Compress), key: None, + resumable_threshold_bytes: self.scope.usecase().resumable_threshold_bytes(), body, } } @@ -85,11 +112,15 @@ impl Session { } /// Creates or replaces an object using a streaming payload. + /// + /// Streams use direct uploads, regardless of the resumable threshold. pub fn put_stream(&self, body: ClientStream) -> PutBuilder { self.put_body(PutBody::Stream(body)) } /// Creates or replaces an object using an [`AsyncRead`] payload. + /// + /// Readers use direct uploads. Use [`Self::put_file`] for automatic resumable file uploads. pub fn put_read(&self, body: R) -> PutBuilder where R: AsyncRead + Send + Sync + 'static, @@ -125,6 +156,30 @@ impl Session { } /// A [`put`](Session::put) request builder. +/// +/// # Automatic resumable uploads +/// +/// Buffers and regular files (including paths) use resumable uploads when their remaining +/// input size meets the configured threshold, by default 256 MiB before client compression. +/// Streams, generic readers, and files whose size cannot be determined use direct uploads. +/// Files must remain unchanged until the upload finishes. File uploads start at the current +/// position; path uploads open the file only when executed. +/// +/// Eligible uploads reuse uncompressed or precompressed input. Client compression first spools +/// the complete compressed payload to a temporary file, using [`crate::ClientBuilder::spool_dir`] +/// when configured. If file creation fails, the client falls back to direct streaming. Later +/// preparation errors propagate. Unsupported or failed session creation also falls back to a +/// direct upload of the prepared bytes, without compressing them again. +/// +/// Transient failures and writes making no progress are retried twice, with exponential delays +/// of 2 and 4 seconds plus up to 1 second of jitter. Recovery queries the server's persisted +/// offset before replaying data. The retry budget resets whenever that offset advances. +/// Terminal failures after session creation attempt cancellation and return +/// [`crate::Error::UploadFailed`] preserving the cause. Dropping the future releases local +/// resources; the remote session is left to expire. +/// +/// Automatic uploads require no feature flag. The `resumable-upload-api` feature exposes the +/// manual protocol API. See [`Self::resumable_threshold_bytes`] to override eligibility. #[derive(Debug)] pub struct PutBuilder { pub(crate) session: Session, @@ -132,9 +187,28 @@ pub struct PutBuilder { pub(crate) compression: Option, pub(crate) key: Option, pub(crate) body: PutBody, + pub(crate) resumable_threshold_bytes: Option, } impl PutBuilder { + /// Sets the minimum remaining input size for automatic resumable uploads. + /// + /// Inherits [`crate::Usecase::resumable_threshold_bytes`] unless overridden. Pass [`None`] + /// to disable or zero to include every supported payload. Eligible uploads queued through + /// [`Session::many`] are sent individually rather than batched. See [`PutBuilder`] for + /// preparation and recovery behavior. + /// + /// ```no_run + /// # async fn example(session: objectstore_client::Session) -> objectstore_client::Result<()> { + /// session.put("payload").resumable_threshold_bytes(None).send().await?; + /// # Ok(()) + /// # } + /// ``` + pub fn resumable_threshold_bytes(mut self, threshold: impl Into>) -> Self { + self.resumable_threshold_bytes = threshold.into(); + self + } + /// Sets an explicit object key. /// /// If a key is specified, the object will be stored under that key. Otherwise, the Objectstore @@ -146,7 +220,7 @@ impl PutBuilder { /// Sets an explicit compression algorithm to be used for this payload. /// - /// The client compresses the payload while uploading it and records the algorithm in the + /// The client compresses the payload and records the algorithm in the /// object's metadata. [`None`] should be used if no compression should be performed by the /// client, either because the payload is uncompressible (such as a media format), or if the /// compression should not be recorded for this object. @@ -329,15 +403,44 @@ impl PutBuilder { None => reqwest::Method::POST, }; - let mut builder = self - .session - .request(method, self.key.as_deref().unwrap_or_default())?; - self.metadata.compression = self.compression.map(CompressionMode::compression); - let body = encode_body(self.body, self.compression).await?; - - builder = builder.headers(self.metadata.to_headers("")?); + let headers = self.metadata.to_headers("")?; + let prepared = automatic::prepare( + self.body, + self.compression, + self.resumable_threshold_bytes, + self.session.client.spool_dir.clone(), + ) + .await?; + let mut content_length = None; + // Keep the file reader alive until a direct fallback request finishes. + let mut replay = None; + let body = match prepared { + automatic::PreparedBody::Direct(body) => body, + automatic::PreparedBody::Replay(body) => { + let create = crate::resumable::CreateResumableUploadBuilder { + session: self.session.clone(), + total_length: body.len(), + key: self.key.clone(), + metadata: self.metadata, + }; + if let Ok(Some(upload)) = create.send().await { + return automatic::upload(body, upload).await; + } + content_length = Some(body.len()); + Body::wrap_stream(replay.insert(body).stream(0).await?) + } + }; + // Prepare authentication only when ready to send: spooling can take long enough for + // a token minted before preparation to expire. + let mut builder = self + .session + .request(method, self.key.as_deref().unwrap_or_default())? + .headers(headers); + if let Some(length) = content_length { + builder = builder.header(reqwest::header::CONTENT_LENGTH, length); + } let response = builder.body(body).send().await?; Ok(response.error_for_status_and_drain().await?.json().await?) } diff --git a/clients/rust/src/put/automatic.rs b/clients/rust/src/put/automatic.rs new file mode 100644 index 00000000..84f5f2cb --- /dev/null +++ b/clients/rust/src/put/automatic.rs @@ -0,0 +1,334 @@ +//! Preparation and recovery for automatic resumable uploads. + +use std::io::{self, SeekFrom}; +use std::path::PathBuf; +use std::time::Duration; + +use futures_util::{StreamExt, stream}; +use tokio::fs::File; +use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt}; +use tokio::sync::{mpsc, oneshot}; + +use super::{CompressionMode, PutBody, PutResponse, encode_body}; +use crate::resumable::{ResumableUpload, UploadProgress}; +use crate::{ClientStream, Error}; + +pub(super) enum PreparedBody { + Direct(reqwest::Body), + Replay(ReplayBody), +} + +pub(super) enum ReplayBody { + Buffer(bytes::Bytes), + File { + file: Option, + start: u64, + length: u64, + }, +} + +/// One reader owns the file for an attempt and returns it after stopping. +pub(super) enum ReplayFile { + Ready(File), + Reading { + stop: oneshot::Sender<()>, + task: tokio::task::JoinHandle, + }, +} + +/// Marks local read failures even when reqwest wraps them in a transport error. +#[derive(Debug, thiserror::Error)] +#[error("could not read upload payload: {0}")] +struct PayloadReadError(#[source] io::Error); + +#[derive(Debug, thiserror::Error)] +#[error("upload made no progress")] +struct NoUploadProgress; + +pub(super) async fn prepare( + mut body: PutBody, + compression: Option, + threshold: Option, + spool_dir: Option, +) -> crate::Result { + if let Some(threshold) = threshold { + if let PutBody::Path(path) = body { + body = PutBody::File(tokio::fs::File::open(path).await?); + } + if let Some(size) = body.remaining_size().await.ok().flatten() + && size >= threshold + { + let mut length = size; + if let Some(CompressionMode::Compress(crate::Compression::Zstd)) = compression { + let temporary = tokio::task::spawn_blocking(move || match spool_dir { + Some(dir) => tempfile::tempfile_in(dir), + None => tempfile::tempfile(), + }) + .await + .map_err(io::Error::other)?; + let mut temporary = match temporary { + Ok(file) => tokio::fs::File::from_std(file), + // No input has been consumed, so direct streaming is still possible. + Err(_) => { + return Ok(PreparedBody::Direct(encode_body(body, compression).await?)); + } + }; + let reader: Box = match body { + PutBody::Buffer(bytes) => Box::new(io::Cursor::new(bytes)), + PutBody::File(file) => Box::new(file), + _ => unreachable!("only buffers and opened regular files have a known size"), + }; + let mut encoder = super::ZstdEncoder::new(super::BufReader::new(reader)); + length = tokio::io::copy(&mut encoder, &mut temporary).await?; + temporary.flush().await?; + temporary.rewind().await?; + body = PutBody::File(temporary); + } + let replay = match body { + PutBody::Buffer(bytes) => ReplayBody::Buffer(bytes), + PutBody::File(mut file) => { + let start = file.stream_position().await?; + ReplayBody::File { + file: Some(ReplayFile::Ready(file)), + start, + length, + } + } + _ => unreachable!("only buffers and opened regular files have a known size"), + }; + return Ok(PreparedBody::Replay(replay)); + } + } + Ok(PreparedBody::Direct(encode_body(body, compression).await?)) +} + +impl ReplayBody { + pub(super) fn len(&self) -> u64 { + match self { + Self::Buffer(bytes) => bytes.len() as u64, + Self::File { length, .. } => *length, + } + } + + /// Creates a non-replayable HTTP stream, waiting for the previous file reader to finish. + pub(super) async fn stream(&mut self, offset: u64) -> io::Result { + match self { + Self::Buffer(bytes) => { + let bytes = bytes.slice(offset as usize..); + Ok(stream::once(async move { Ok(bytes) }).boxed()) + } + Self::File { + file, + start, + length, + } => { + let mut owned = match file.take().expect("file is present between attempts") { + ReplayFile::Ready(file) => file, + ReplayFile::Reading { stop, task } => { + // The HTTP layer may retain the old body after returning a response. + // Stop its reader explicitly, then wait for exclusive file ownership. + drop(stop); + task.await.map_err(io::Error::other)? + } + }; + owned.seek(SeekFrom::Start(*start + offset)).await?; + let (send, receive) = mpsc::channel(1); + let (stop, stopped) = oneshot::channel(); + let length = *length - offset; + let task = tokio::spawn(async move { + { + let read = read_file(&mut owned, length, send); + futures_util::pin_mut!(read, stopped); + let _ = futures_util::future::select(stopped, read).await; + } + // Canceling an async read can leave disk IO running. Wait for it + // before returning the file to the next attempt, without reopening it. + File::from_std(owned.into_std().await) + }); + *file = Some(ReplayFile::Reading { stop, task }); + Ok(stream::unfold(receive, |mut receive| async { + receive.recv().await.map(|bytes| (bytes, receive)) + }) + .boxed()) + } + } + } +} + +/// Reads only while the request accepts data, with at most one buffered chunk. +async fn read_file( + file: &mut File, + mut remaining: u64, + send: mpsc::Sender>, +) { + while remaining > 0 { + let Ok(permit) = send.reserve().await else { + break; + }; + let mut bytes = vec![0; remaining.min(64 * 1024) as usize]; + let result = file.read(&mut bytes).await.and_then(|count| { + if count == 0 { + Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "upload file became shorter", + )) + } else { + remaining -= count as u64; + bytes.truncate(count); + Ok(bytes::Bytes::from(bytes)) + } + }); + let failed = result.is_err(); + permit.send(result.map_err(|error| io::Error::other(PayloadReadError(error)))); + if failed { + break; + } + } +} + +pub(super) async fn upload( + mut body: ReplayBody, + handle: ResumableUpload, +) -> crate::Result { + let result = upload_inner(&mut body, &handle).await; + if result.is_err() { + let _ = handle.cancel().send().await; + } + result.map_err(Error::UploadFailed) +} + +async fn upload_inner( + body: &mut ReplayBody, + handle: &ResumableUpload, +) -> Result> { + let mut offset = 0; + let mut retries = 0; + let mut probe = false; + loop { + let result = if probe { + handle.progress().send().await + } else { + handle + .put_stream(offset, body.len() - offset, body.stream(offset).await?) + .send() + .await + }; + let progress = match result { + Err(Error::UploadOffsetMismatch { offset }) => { + Ok(UploadProgress::Incomplete { offset }) + } + other => other, + }; + let failure: Box = match progress { + Ok(UploadProgress::Complete) => { + return Ok(PutResponse { + key: handle.key().to_owned(), + }); + } + Ok(UploadProgress::Incomplete { offset: next }) => { + if next < offset || next > body.len() { + return Err(Error::MalformedResponse("invalid upload offset".into()).into()); + } + if next == offset && !probe { + Box::new(NoUploadProgress) + } else { + if next > offset { + retries = 0; + } + offset = next; + probe = false; + continue; + } + } + Err(error) => { + if !is_transient(&error) { + return Err(error.into()); + } + error.into() + } + }; + if retries >= 2 { + return Err(failure); + } + backoff(retries).await; + retries += 1; + probe = true; + } +} + +async fn backoff(retries: u32) { + let delay = Duration::from_secs(2 << retries) + Duration::from_secs_f64(rand::random::()); + tokio::time::sleep(delay).await; +} + +fn is_transient(error: &Error) -> bool { + let mut source: Option<&(dyn std::error::Error + 'static)> = Some(error); + while let Some(error) = source { + if error.is::() + || error + .downcast_ref::() + .and_then(io::Error::get_ref) + .is_some_and(|inner| inner.is::()) + { + return false; + } + source = error.source(); + } + match error { + Error::Reqwest(error) => match error.status() { + Some(status) => matches!(status.as_u16(), 408 | 429 | 502 | 503 | 504), + None => { + error.is_timeout() || error.is_connect() || error.is_request() || error.is_body() + } + }, + _ => false, + } +} + +#[cfg(test)] +mod tests { + use futures_util::TryStreamExt; + + use super::*; + + #[test] + fn retry_stops_previous_reader_and_waits_for_pending_read() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .max_blocking_threads(1) + .build() + .unwrap(); + runtime.block_on(async { + let data: Vec = (0..200_000).map(|n| (n % 251) as u8).collect(); + let mut file = tempfile::tempfile().unwrap(); + std::io::Write::write_all(&mut file, &data).unwrap(); + let mut body = ReplayBody::File { + file: Some(ReplayFile::Ready(File::from_std(file))), + start: 7, + length: 150_000, + }; + let mut old = body.stream(0).await.unwrap(); + + // Occupy the sole blocking worker so the old read remains pending. + let (release, blocked) = std::sync::mpsc::channel(); + let (started, running) = oneshot::channel(); + let worker = tokio::task::spawn_blocking(move || { + started.send(()).unwrap(); + blocked.recv().unwrap(); + }); + running.await.unwrap(); + tokio::task::yield_now().await; + assert!(futures_util::poll!(old.next()).is_pending()); + // Retain the old HTTP body: retry must stop its reader explicitly. + let retry = body.stream(100); + tokio::pin!(retry); + assert!(futures_util::poll!(&mut retry).is_pending()); + release.send(()).unwrap(); + worker.await.unwrap(); + let retry = retry.await.unwrap(); + let remaining: Vec<_> = retry.try_collect().await.unwrap(); + assert_eq!(remaining.concat(), data[107..150_007]); + drop(old); + }); + } +} diff --git a/clients/rust/src/resumable.rs b/clients/rust/src/resumable.rs index b4dc58bb..7cd9c315 100644 --- a/clients/rust/src/resumable.rs +++ b/clients/rust/src/resumable.rs @@ -3,14 +3,15 @@ //! Enable the `resumable-upload-api` feature to use this API. //! It maps directly to the server operations and is intended for advanced use cases that need //! low-level protocol calls. -//! A future release of `objectstore-client` will use this API under the hood for eligible -//! `Session::put` calls, so users can benefit from this feature without the need to interact with -//! this API directly. +//! Eligible uploads use this protocol automatically, without requiring the feature flag. +#[cfg(feature = "resumable-upload-api")] use std::borrow::Cow; +#[cfg(feature = "resumable-upload-api")] use std::collections::BTreeMap; use std::fmt; +#[cfg(feature = "resumable-upload-api")] use bytes::Bytes; use objectstore_types::metadata::Metadata; use objectstore_types::resumable::{ @@ -19,13 +20,17 @@ use objectstore_types::resumable::{ }; use reqwest::{Body, Method, Response, StatusCode}; use serde::Serialize; +#[cfg(feature = "resumable-upload-api")] use tokio::io::AsyncRead; +#[cfg(feature = "resumable-upload-api")] use tokio_util::io::ReaderStream; pub use objectstore_types::resumable::{SessionToken, UploadProgress}; use crate::response::ResponseExt as _; -use crate::{ClientStream, Compression, Error, ExpirationPolicy, ObjectKey, Session}; +use crate::{ClientStream, Error, ObjectKey, Session}; +#[cfg(feature = "resumable-upload-api")] +use crate::{Compression, ExpirationPolicy}; #[derive(Serialize)] #[serde(rename_all = "snake_case")] @@ -56,6 +61,7 @@ pub struct ResumableUpload { granularity: Option, } +#[cfg(feature = "resumable-upload-api")] impl Session { /// Starts building a resumable upload for an object with `object_length` bytes. /// @@ -102,6 +108,7 @@ impl ResumableUpload { } /// Returns the session token bound to this upload. + #[cfg(feature = "resumable-upload-api")] pub fn token(&self) -> &SessionToken { &self.token } @@ -110,6 +117,7 @@ impl ResumableUpload { /// /// The granularity is the persistence unit for non-final chunks: chunks shorter than one /// unit are rejected, and larger chunks may persist only an aligned prefix. + #[cfg(feature = "resumable-upload-api")] pub fn granularity(&self) -> Option { self.granularity } @@ -129,6 +137,7 @@ impl ResumableUpload { /// /// The returned progress contains the authoritative server offset, which should be used for /// subsequent requests. + #[cfg(feature = "resumable-upload-api")] pub fn put(&self, offset: u64, chunk: impl Into) -> PutChunkBuilder { let chunk = chunk.into(); self.put_body(offset, chunk.len() as u64, chunk.into()) @@ -146,6 +155,7 @@ impl ResumableUpload { /// /// The reader must produce exactly `length` bytes. If this upload records compression, /// those bytes must come from the object after it has been compressed in full. + #[cfg(feature = "resumable-upload-api")] pub fn put_read(&self, offset: u64, length: u64, reader: R) -> PutChunkBuilder where R: AsyncRead + Send + Sync + 'static, @@ -201,15 +211,16 @@ impl ResumableUpload { } } -/// A builder for [`Session::create_upload`]. +/// A builder for creating a resumable upload. #[derive(Debug)] pub struct CreateResumableUploadBuilder { - session: Session, - total_length: u64, - key: Option, - metadata: Metadata, + pub(crate) session: Session, + pub(crate) total_length: u64, + pub(crate) key: Option, + pub(crate) metadata: Metadata, } +#[cfg(feature = "resumable-upload-api")] impl CreateResumableUploadBuilder { /// Sets an explicit object key instead of asking the server to generate one. pub fn key(mut self, key: impl Into) -> Self { @@ -264,7 +275,9 @@ impl CreateResumableUploadBuilder { self.metadata.custom.insert(key.into(), value.into()); self } +} +impl CreateResumableUploadBuilder { /// Creates the resumable upload and returns a handle bound to it. /// /// Returns `None` when the server declines resumable uploads for this object. @@ -336,8 +349,7 @@ impl UploadProgressBuilder { } } -/// A builder for [`ResumableUpload::put`], [`ResumableUpload::put_stream`], or -/// [`ResumableUpload::put_read`]. +/// A builder for writing a resumable upload chunk. pub struct PutChunkBuilder { upload: ResumableUpload, offset: u64, diff --git a/clients/rust/tests/automatic.rs b/clients/rust/tests/automatic.rs new file mode 100644 index 00000000..d9b60cff --- /dev/null +++ b/clients/rust/tests/automatic.rs @@ -0,0 +1,619 @@ +//! Automatic upload tests deliberately run without `resumable-upload-api` too. + +use std::collections::VecDeque; +use std::io::{self, SeekFrom, Write}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use axum::body::{Body, Bytes, to_bytes}; +use axum::extract::{Request, State}; +use axum::http::{HeaderMap, Method, Response, Uri}; +use futures_util::{StreamExt, stream}; +use objectstore_client::{Client, Compression, Error, ExpirationPolicy, Session, Usecase}; +use tokio::io::AsyncSeekExt; + +#[derive(Clone, Debug)] +struct Received { + method: Method, + uri: Uri, + headers: HeaderMap, + body: Bytes, +} + +enum Reply { + Create, + Complete, + Offset(u16, u64), + EarlyOffset(u64), + Status(u16), + LostCompletion, + Response(Response), +} + +impl Reply { + fn response(self) -> Response { + let mut response = Response::builder(); + let body = match self { + Self::Create => Body::from(r#"{"key":"key","session":"dG9rZW4","granularity":4}"#), + Self::Complete => { + response = response.status(201); + Body::from(r#"{"key":"key"}"#) + } + Self::Offset(status, offset) => { + response = response.status(status).header("upload-offset", offset); + Body::empty() + } + Self::EarlyOffset(offset) => return Self::Offset(409, offset).response(), + Self::Status(status) => { + response = response.status(status); + Body::from(r#"{"key":"key"}"#) + } + Self::LostCompletion => { + response = response.status(201).header("content-length", "100"); + Body::from_stream(stream::once(async { + Err::(io::Error::other("connection lost after commit")) + })) + } + Self::Response(response) => return response, + }; + response.body(body).unwrap() + } +} + +#[derive(Default)] +struct Script { + replies: VecDeque, + received: Vec, +} + +struct Server { + url: String, + script: Arc>, + task: tokio::task::JoinHandle<()>, +} + +impl Server { + async fn new(replies: impl IntoIterator) -> Self { + async fn handle( + State(script): State>>, + request: Request, + ) -> Response { + let (parts, body) = request.into_parts(); + // Model a server that responds before consuming the request body. + { + let mut script = script.lock().unwrap(); + if matches!(script.replies.front(), Some(Reply::EarlyOffset(_))) { + script.received.push(Received { + method: parts.method, + uri: parts.uri, + headers: parts.headers, + body: Bytes::new(), + }); + // Drain in the background so HTTP/1 can deliver the early response + // without resetting the connection due to unread request data. + tokio::spawn(async move { + let _ = to_bytes(body, 16 * 1024 * 1024).await; + }); + return script.replies.pop_front().unwrap().response(); + } + } + // A truncated request body is expected in the local-file-failure test. + let Ok(body) = to_bytes(body, 16 * 1024 * 1024).await else { + return Reply::Status(400).response(); + }; + let reply = { + let mut script = script.lock().unwrap(); + script.received.push(Received { + method: parts.method, + uri: parts.uri, + headers: parts.headers, + body, + }); + script.replies.pop_front().expect("unexpected HTTP request") + }; + reply.response() + } + let script = Arc::new(Mutex::new(Script { + replies: replies.into_iter().collect(), + ..Default::default() + })); + let app = axum::Router::new() + .fallback(handle) + .with_state(script.clone()); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}", listener.local_addr().unwrap()); + let task = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + Self { url, script, task } + } + + fn client(&self) -> objectstore_client::ClientBuilder { + Client::builder(&self.url) + .configure_reqwest(|builder| builder.no_proxy().connect_timeout(Duration::from_secs(5))) + } + + fn session(&self, threshold: Option) -> Session { + self.client() + .build() + .unwrap() + .session( + Usecase::new("testing") + .with_compression(None) + .with_resumable_threshold_bytes(threshold) + .for_organization(123), + ) + .unwrap() + } + + fn finish(&self) -> Vec { + let script = self.script.lock().unwrap(); + assert!( + script.replies.is_empty(), + "{} unused replies", + script.replies.len() + ); + script.received.clone() + } +} + +impl Drop for Server { + fn drop(&mut self) { + self.task.abort(); + } +} + +#[tokio::test] +async fn threshold_inheritance_overrides_and_empty_payloads() { + assert_eq!( + Usecase::new("testing").resumable_threshold_bytes(), + Some(256 * 1024 * 1024) + ); + for (threshold, override_threshold, data, resumable) in [ + (Some(4), None, "abc", false), + (Some(4), None, "abcd", true), + (Some(4), Some(None), "abcd", false), + (None, None, "abcd", false), + (None, Some(Some(4)), "abcd", true), + (Some(0), None, "", true), + ] { + let replies = if resumable { + vec![Reply::Create, Reply::Complete] + } else { + vec![Reply::Status(200)] + }; + let server = Server::new(replies).await; + let mut put = server.session(threshold).put(data); + if let Some(threshold) = override_threshold { + put = put.resumable_threshold_bytes(threshold); + } + assert_eq!(put.send().await.unwrap().key, "key"); + let requests = server.finish(); + assert_eq!(requests.last().unwrap().body.as_ref(), data.as_bytes()); + assert_eq!(requests[0].headers.contains_key("upload-length"), resumable); + } +} + +#[tokio::test] +async fn streams_and_readers_remain_direct() { + let server = Server::new([Reply::Status(200), Reply::Status(200)]).await; + let session = server.session(Some(0)); + session + .put_stream( + stream::iter([Ok(Bytes::from_static(b"a")), Ok(Bytes::from_static(b"bc"))]).boxed(), + ) + .send() + .await + .unwrap(); + session + .put_read(io::Cursor::new(b"def")) + .send() + .await + .unwrap(); + let requests = server.finish(); + assert_eq!(requests[0].body, "abc"); + assert_eq!(requests[1].body, "def"); + assert!( + requests + .iter() + .all(|r| !r.headers.contains_key("upload-length")) + ); +} + +#[tokio::test] +async fn files_use_remaining_length_and_paths_open_lazily() { + let server = Server::new([ + Reply::Create, + Reply::Offset(204, 2), + Reply::Complete, + Reply::Status(200), + Reply::Create, + Reply::Complete, + ]) + .await; + let session = server.session(Some(4)); + let mut temporary = tempfile::NamedTempFile::new().unwrap(); + temporary.write_all(b"prefix-data").unwrap(); + let mut file = tokio::fs::File::open(temporary.path()).await.unwrap(); + file.seek(SeekFrom::Start(7)).await.unwrap(); + session.put_file(file).send().await.unwrap(); + let mut file = tokio::fs::File::open(temporary.path()).await.unwrap(); + file.seek(SeekFrom::Start(8)).await.unwrap(); + session.put_file(file).send().await.unwrap(); + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("created-after-builder"); + let put = session.put_path(&path); + tokio::fs::write(path, b"path").await.unwrap(); + put.send().await.unwrap(); + let requests = server.finish(); + assert_eq!(requests[0].headers["upload-length"], "4"); + assert_eq!(requests[1].body, "data"); + assert_eq!(requests[2].body, "ta"); + assert_eq!(requests[3].body, "ata"); + assert_eq!(requests[5].body, "path"); +} + +#[tokio::test] +async fn file_retry_after_response_before_body_is_consumed() { + let server = Server::new([Reply::Create, Reply::EarlyOffset(4), Reply::Complete]).await; + let data: Vec = (0..8 * 1024 * 1024).map(|n| (n % 251) as u8).collect(); + let mut file = tempfile::tempfile().unwrap(); + file.write_all(&data).unwrap(); + std::io::Seek::rewind(&mut file).unwrap(); + let put = server + .session(Some(0)) + .put_file(tokio::fs::File::from_std(file)); + tokio::time::timeout(Duration::from_secs(10), put.send()) + .await + .expect("retry must recover ownership of the previous request's file") + .unwrap(); + let requests = server.finish(); + assert_eq!(requests[2].method, Method::PUT); + assert_eq!(requests[2].body.len(), data.len() - 4); + assert!(requests[2].body.as_ref() == &data[4..]); +} + +#[tokio::test] +async fn compression_spooling_metadata_and_creation_fallback() { + for creation in [Reply::Create, Reply::Status(501), Reply::Status(503)] { + let fallback = !matches!(creation, Reply::Create); + let server = Server::new([ + creation, + if fallback { + Reply::Status(200) + } else { + Reply::Complete + }, + ]) + .await; + let dir = tempfile::tempdir().unwrap(); + let client = server.client().spool_dir(dir.path()).build().unwrap(); + let session = client + .session( + Usecase::new("testing") + .with_resumable_threshold_bytes(1024) + .for_organization(123), + ) + .unwrap(); + session + .put(vec![b'x'; 1024]) + .key("key") + .content_type("text/plain") + .filename("file.txt") + .origin("127.0.0.1") + .append_metadata("test", "value") + .expiration_policy(ExpirationPolicy::TimeToLive(Duration::from_secs(60))) + .send() + .await + .unwrap(); + let requests = server.finish(); + let encoded = &requests[1].body; + assert_eq!( + zstd::decode_all(encoded.as_ref()).unwrap(), + vec![b'x'; 1024] + ); + assert!(encoded.len() < 1024); // Eligibility was decided before compression. + assert_eq!( + requests[0].headers["upload-length"], + encoded.len().to_string() + ); + assert_eq!( + requests[1].headers["content-length"], + encoded.len().to_string() + ); + assert_eq!(requests[0].method, Method::PUT); + let metadata = + objectstore_types::metadata::Metadata::from_headers(&requests[0].headers, "").unwrap(); + assert_eq!(metadata.compression, Some(Compression::Zstd)); + assert_eq!(metadata.content_type, "text/plain"); + assert_eq!(metadata.filename.as_deref(), Some("file.txt")); + assert_eq!(metadata.origin.as_deref(), Some("127.0.0.1")); + assert_eq!(metadata.custom["test"], "value"); + assert_eq!( + metadata.expiration_policy, + ExpirationPolicy::TimeToLive(Duration::from_secs(60)) + ); + if fallback { + assert_eq!(requests[1].headers["content-encoding"], "zstd"); + } + assert_eq!(std::fs::read_dir(dir.path()).unwrap().count(), 0); + } +} + +#[tokio::test] +async fn missing_spool_directory_falls_back_and_precompressed_input_needs_no_spool() { + let server = Server::new([Reply::Status(200), Reply::Create, Reply::Complete]).await; + let dir = tempfile::tempdir().unwrap(); + let client = server + .client() + .spool_dir(dir.path().join("missing")) + .build() + .unwrap(); + let session = client + .session( + Usecase::new("testing") + .with_resumable_threshold_bytes(0) + .for_organization(123), + ) + .unwrap(); + session.put("original").send().await.unwrap(); + let compressed = zstd::encode_all(&b"original"[..], 0).unwrap(); + session + .put(compressed.clone()) + .precompressed(Compression::Zstd) + .send() + .await + .unwrap(); + let requests = server.finish(); + assert!(!requests[0].headers.contains_key("upload-length")); + assert_eq!( + zstd::decode_all(requests[0].body.as_ref()).unwrap(), + b"original" + ); + assert_eq!(requests[2].body.as_ref(), compressed); +} + +#[tokio::test] +async fn many_honors_threshold_overrides() { + let server = Server::new([Reply::Create, Reply::Complete]).await; + let session = server.session(None); + session + .many() + .push(session.put("data").resumable_threshold_bytes(4)) + .send() + .await + .error_for_failures() + .await + .map_err(|errors| errors.collect::>()) + .unwrap(); + let requests = server.finish(); + assert_eq!(requests[0].uri.query(), Some("upload_type=resumable")); + assert_eq!(requests[1].body, "data"); +} + +#[tokio::test] +async fn partial_progress_and_conflicts_replay_only_the_suffix() { + let server = Server::new([ + Reply::Create, + Reply::Offset(204, 3), + Reply::Offset(409, 5), + Reply::Complete, + ]) + .await; + server + .session(Some(0)) + .put("abcdefgh") + .send() + .await + .unwrap(); + let requests = server.finish(); + for (request, offset, body) in [ + (&requests[1], "0", "abcdefgh"), + (&requests[2], "3", "defgh"), + (&requests[3], "5", "fgh"), + ] { + assert_eq!(request.headers["upload-offset"], offset); + assert_eq!(request.body, body); + } +} + +#[tokio::test] +async fn recovery_probes_and_progress_resets_the_budget() { + let server = Server::new([ + Reply::Create, + Reply::Status(503), + Reply::Status(502), + Reply::Offset(204, 2), + Reply::Status(429), + Reply::Status(504), + Reply::Offset(204, 4), + Reply::Complete, + ]) + .await; + // Even with middleware retries enabled, a failed write must be followed by a probe. + let client = server + .client() + .configure_reqwest(|b| { + b.retry( + reqwest::retry::for_host("127.0.0.1") + .no_budget() + .classify_fn(|r| { + if r.status() == Some(reqwest::StatusCode::SERVICE_UNAVAILABLE) { + r.retryable() + } else { + r.success() + } + }), + ) + }) + .build() + .unwrap(); + let session = client + .session( + Usecase::new("testing") + .with_compression(None) + .with_resumable_threshold_bytes(0) + .for_organization(123), + ) + .unwrap(); + session.put("abcdef").send().await.unwrap(); + let requests = server.finish(); + let offsets: Vec<_> = requests[1..] + .iter() + .map(|r| r.headers["upload-offset"].to_str().unwrap()) + .collect(); + assert_eq!(offsets, ["0", "*", "*", "2", "*", "*", "4"]); + assert_eq!(requests[4].body, "cdef"); + assert_eq!(requests[7].body, "ef"); +} + +#[tokio::test] +async fn no_progress_and_probe_failures_share_the_retry_budget() { + let server = Server::new([ + Reply::Create, + Reply::Offset(204, 0), + Reply::Status(503), + Reply::Offset(204, 0), + Reply::Offset(409, 0), + Reply::Status(500), + ]) + .await; + let error = server + .session(Some(0)) + .put("data") + .send() + .await + .unwrap_err(); + assert!(matches!(&error, Error::UploadFailed(_))); + assert!(error.to_string().contains("no progress")); + let requests = server.finish(); + assert_eq!(requests[2].headers["upload-offset"], "*"); + assert_eq!(requests[3].headers["upload-offset"], "*"); + assert_eq!(requests.last().unwrap().method, Method::DELETE); +} + +#[tokio::test] +async fn permanent_errors_and_invalid_offsets_cancel_without_restarting() { + for replies in [ + vec![Reply::Create, Reply::Status(403), Reply::Status(500)], + vec![Reply::Create, Reply::Offset(204, 10), Reply::Status(204)], + vec![ + Reply::Create, + Reply::Offset(204, 2), + Reply::Offset(204, 1), + Reply::Status(204), + ], + vec![Reply::Create, Reply::Status(410), Reply::Status(204)], + ] { + let server = Server::new(replies).await; + let error = server + .session(Some(0)) + .put("data") + .send() + .await + .unwrap_err(); + assert!(matches!(error, Error::UploadFailed(_))); + assert!(std::error::Error::source(&error).is_some()); + assert_eq!(server.finish().last().unwrap().method, Method::DELETE); + } +} + +#[tokio::test] +async fn lost_completion_is_recovered_without_reuploading() { + let server = Server::new([Reply::Create, Reply::LostCompletion, Reply::Complete]).await; + server.session(Some(0)).put("data").send().await.unwrap(); + let requests = server.finish(); + assert_eq!(requests[2].headers["upload-offset"], "*"); + assert!(requests[2].body.is_empty()); +} + +#[tokio::test] +async fn local_file_errors_are_not_retried() { + let mut file = tempfile::NamedTempFile::new().unwrap(); + file.write_all(b"data").unwrap(); + let path = file.path().to_owned(); + let (parts, body) = Reply::Create.response().into_parts(); + let body = Body::from_stream(stream::once(async move { + // The client has measured the file, but cannot upload until it receives this body. + tokio::fs::write(path, []).await?; + to_bytes(body, 1024).await.map_err(io::Error::other) + })); + let server = Server::new([ + Reply::Response(Response::from_parts(parts, body)), + Reply::Status(204), + ]) + .await; + let error = server + .session(Some(0)) + .put_path(file.path()) + .send() + .await + .unwrap_err(); + assert!(matches!(error, Error::UploadFailed(_))); + assert!(error.to_string().contains("upload failed")); + let requests = server.finish(); + assert_eq!(requests.len(), 2); + assert_eq!(requests[1].method, Method::DELETE); +} + +#[tokio::test] +async fn compressed_file_replays_the_prepared_bytes() { + let server = Server::new([Reply::Create, Reply::Offset(204, 4), Reply::Complete]).await; + let mut file = tempfile::NamedTempFile::new().unwrap(); + file.write_all(b"prefix-payload").unwrap(); + let mut input = tokio::fs::File::open(file.path()).await.unwrap(); + input.seek(SeekFrom::Start(7)).await.unwrap(); + server + .session(Some(0)) + .put_file(input) + .compress(Compression::Zstd) + .send() + .await + .unwrap(); + let requests = server.finish(); + let encoded = &requests[1].body; + assert_eq!(requests[2].body, encoded.slice(4..)); + assert_eq!(zstd::decode_all(encoded.as_ref()).unwrap(), b"payload"); +} + +#[tokio::test] +async fn preparation_read_failure_does_not_fall_back() { + let server = Server::new([]).await; + let mut file = tempfile::NamedTempFile::new().unwrap(); + file.write_all(b"payload").unwrap(); + let write_only = std::fs::OpenOptions::new() + .write(true) + .open(file.path()) + .unwrap(); + let error = server + .session(Some(0)) + .put_file(tokio::fs::File::from_std(write_only)) + .compress(Compression::Zstd) + .send() + .await + .unwrap_err(); + assert!(matches!(error, Error::Io(_))); + assert!(server.finish().is_empty()); +} + +#[cfg(unix)] +#[tokio::test] +async fn non_regular_files_use_direct_uploads() { + let server = Server::new([Reply::Status(200)]).await; + let file = tokio::fs::File::open("/dev/null").await.unwrap(); + server.session(Some(0)).put_file(file).send().await.unwrap(); + let requests = server.finish(); + assert!(!requests[0].headers.contains_key("upload-length")); + assert!(requests[0].body.is_empty()); +} + +#[test] +fn upload_futures_are_send() { + fn assert_send(_: impl Send) {} + let session = Client::builder("http://localhost") + .build() + .unwrap() + .session(Usecase::new("testing").for_organization(123)) + .unwrap(); + assert_send(session.put("data").send()); + assert_send(session.put_path("data").send()); + assert_send(session.many().push(session.put("data")).send()); +} diff --git a/clients/rust/tests/e2e.rs b/clients/rust/tests/e2e.rs index efee1071..4b27c91a 100644 --- a/clients/rust/tests/e2e.rs +++ b/clients/rust/tests/e2e.rs @@ -71,14 +71,23 @@ async fn stores_uncompressed() { let body = "oh hai!"; - let stored_id = session.put(body).compress(None).send().await.unwrap().key; + for threshold in [None, Some(0)] { + let stored_id = session + .put(body) + .resumable_threshold_bytes(threshold) + .compress(None) + .send() + .await + .unwrap() + .key; - let response = session.get(&stored_id).send().await.unwrap().unwrap(); - assert_eq!(response.metadata.compression, None); - assert!(response.metadata.time_created.is_some()); + let response = session.get(&stored_id).send().await.unwrap().unwrap(); + assert_eq!(response.metadata.compression, None); + assert!(response.metadata.time_created.is_some()); - let received = response.payload().await.unwrap(); - assert_eq!(received, "oh hai!"); + let received = response.payload().await.unwrap(); + assert_eq!(received, "oh hai!"); + } } #[tokio::test] @@ -93,23 +102,31 @@ async fn uses_zstd_by_default() { let session = client.session(usecase.for_organization(12345)).unwrap(); let body = "oh hai!"; - let stored_id = session.put(body).send().await.unwrap().key; + for threshold in [None, Some(0)] { + let stored_id = session + .put(body) + .resumable_threshold_bytes(threshold) + .send() + .await + .unwrap() + .key; - // when the user indicates that it can deal with zstd, it gets zstd - let request = session.get(&stored_id).decompress(false); - let response = request.send().await.unwrap().unwrap(); - assert_eq!(response.metadata.compression, Some(Compression::Zstd)); + // when the user indicates that it can deal with zstd, it gets zstd + let request = session.get(&stored_id).decompress(false); + let response = request.send().await.unwrap().unwrap(); + assert_eq!(response.metadata.compression, Some(Compression::Zstd)); - let received_compressed = response.payload().await.unwrap(); - let decompressed = zstd::bulk::decompress(&received_compressed, 1024).unwrap(); - assert_eq!(decompressed, b"oh hai!"); + let received_compressed = response.payload().await.unwrap(); + let decompressed = zstd::bulk::decompress(&received_compressed, 1024).unwrap(); + assert_eq!(decompressed, b"oh hai!"); - // otherwise, the client does the decompression - let response = session.get(&stored_id).send().await.unwrap().unwrap(); - assert_eq!(response.metadata.compression, None); + // otherwise, the client does the decompression + let response = session.get(&stored_id).send().await.unwrap().unwrap(); + assert_eq!(response.metadata.compression, None); - let received = response.payload().await.unwrap(); - assert_eq!(received, "oh hai!"); + let received = response.payload().await.unwrap(); + assert_eq!(received, "oh hai!"); + } } #[tokio::test] @@ -124,29 +141,32 @@ async fn stores_precompressed_payload_verbatim() { let session = client.session(usecase.for_organization(12345)).unwrap(); let compressed = zstd::bulk::compress(b"oh hai!", 0).unwrap(); - let stored_id = session - .put(compressed.clone()) - .precompressed(Compression::Zstd) - .send() - .await - .unwrap() - .key; - - // the payload is stored as-is, without another compression pass - let response = session - .get(&stored_id) - .accept_encoding([Compression::Zstd]) - .send() - .await - .unwrap() - .unwrap(); - assert_eq!(response.metadata.compression, Some(Compression::Zstd)); - assert_eq!(response.payload().await.unwrap(), compressed); + for threshold in [None, Some(0)] { + let stored_id = session + .put(compressed.clone()) + .resumable_threshold_bytes(threshold) + .precompressed(Compression::Zstd) + .send() + .await + .unwrap() + .key; - // and the recorded encoding makes downloads decompress transparently - let response = session.get(&stored_id).send().await.unwrap().unwrap(); - assert_eq!(response.metadata.compression, None); - assert_eq!(response.payload().await.unwrap(), "oh hai!"); + // the payload is stored as-is, without another compression pass + let response = session + .get(&stored_id) + .accept_encoding([Compression::Zstd]) + .send() + .await + .unwrap() + .unwrap(); + assert_eq!(response.metadata.compression, Some(Compression::Zstd)); + assert_eq!(response.payload().await.unwrap(), compressed); + + // and the recorded encoding makes downloads decompress transparently + let response = session.get(&stored_id).send().await.unwrap().unwrap(); + assert_eq!(response.metadata.compression, None); + assert_eq!(response.payload().await.unwrap(), "oh hai!"); + } } #[tokio::test] @@ -484,8 +504,8 @@ async fn batch_operations() { let usecase = Usecase::new("usecase"); let session = client.session(usecase.for_project(12345, 1337)).unwrap(); - // First batch: 4 PUT operations - // key-2 uses default compression (zstd), others are uncompressed + // First batch: batched PUTs alongside one automatic resumable upload. + // key-2 uses default compression; key-5 is precompressed. let results: Vec<_> = session .many() .push( @@ -498,12 +518,19 @@ async fn batch_operations() { .push(session.put("second object").key("key-2")) .push(session.put("third object").compress(None).key("key-3")) .push(session.put("fourth object").compress(None).key("key-4")) + .push( + session + .put(zstd::encode_all(&b"fifth object"[..], 0).unwrap()) + .precompressed(Compression::Zstd) + .resumable_threshold_bytes(0) + .key("key-5"), + ) .send() .await .collect() .await; - assert_eq!(results.len(), 4); + assert_eq!(results.len(), 5); let mut keys: Vec = results .iter() .map(|r| match r { @@ -512,7 +539,7 @@ async fn batch_operations() { }) .collect(); keys.sort(); - assert_eq!(keys, vec!["key-1", "key-2", "key-3", "key-4"]); + assert_eq!(keys, vec!["key-1", "key-2", "key-3", "key-4", "key-5"]); // Second batch: GET key-1, GET key-2 (automatic decompression), DELETE key-3, PUT key-4 (override) let results: Vec<_> = session @@ -588,6 +615,10 @@ async fn batch_operations() { // Verify the deleted object is gone let response = session.get("key-3").send().await.unwrap(); assert!(response.is_none()); + + // The precompressed automatic upload remains readable through normal downloads. + let response = session.get("key-5").send().await.unwrap().unwrap(); + assert_eq!(response.payload().await.unwrap(), "fifth object"); } #[tokio::test] @@ -732,48 +763,62 @@ async fn batch_put_files() { // This should be routed to a single PUT request instead of being batched. let large_body: Vec = vec![0xAB; 2 * 1024 * 1024]; // 2 MB - let mut many = session.many(); - - for (name, content) in &small_bodies { - let path = dir.path().join(name); - let mut f = std::fs::File::create(&path).unwrap(); - f.write_all(content).unwrap(); - let file = tokio::fs::File::open(&path).await.unwrap(); - many = many.push(session.put_file(file).compress(None).key(*name)); - } - - let large_path = dir.path().join("large"); - std::fs::File::create(&large_path) - .unwrap() - .write_all(&large_body) - .unwrap(); - let large_file = tokio::fs::File::open(&large_path).await.unwrap(); - many = many.push(session.put_file(large_file).compress(None).key("large")); - - let results: Vec<_> = many.send().await.collect().await; - - assert_eq!(results.len(), 4); + for threshold in [None, Some(0)] { + let mut many = session.many(); + + for (name, content) in &small_bodies { + let path = dir.path().join(name); + let mut f = std::fs::File::create(&path).unwrap(); + f.write_all(content).unwrap(); + let file = tokio::fs::File::open(&path).await.unwrap(); + many = many.push( + session + .put_file(file) + .resumable_threshold_bytes(threshold) + .compress(None) + .key(*name), + ); + } - let mut keys: Vec = results - .iter() - .map(|r| match r { - OperationResult::Put(key, Ok(_)) => key.clone(), - other => panic!("Expected successful Put result, got: {other:?}"), - }) - .collect(); - keys.sort(); - assert_eq!(keys, vec!["large", "small-1", "small-2", "small-3"]); + let large_path = dir.path().join("large"); + std::fs::File::create(&large_path) + .unwrap() + .write_all(&large_body) + .unwrap(); + let large_file = tokio::fs::File::open(&large_path).await.unwrap(); + many = many.push( + session + .put_file(large_file) + .resumable_threshold_bytes(threshold) + .compress(None) + .key("large"), + ); + + let results: Vec<_> = many.send().await.collect().await; + + assert_eq!(results.len(), 4); + + let mut keys: Vec = results + .iter() + .map(|r| match r { + OperationResult::Put(key, Ok(_)) => key.clone(), + other => panic!("Expected successful Put result, got: {other:?}"), + }) + .collect(); + keys.sort(); + assert_eq!(keys, vec!["large", "small-1", "small-2", "small-3"]); + + // Verify all objects are retrievable with the correct content. + for (name, expected) in &small_bodies { + let response = session.get(name).send().await.unwrap().unwrap(); + let payload = response.payload().await.unwrap(); + assert_eq!(payload.as_ref(), expected.as_slice(), "mismatch for {name}"); + } - // Verify all objects are retrievable with the correct content. - for (name, expected) in &small_bodies { - let response = session.get(name).send().await.unwrap().unwrap(); + let response = session.get("large").send().await.unwrap().unwrap(); let payload = response.payload().await.unwrap(); - assert_eq!(payload.as_ref(), expected.as_slice(), "mismatch for {name}"); + assert_eq!(payload.as_ref(), large_body.as_slice()); } - - let response = session.get("large").send().await.unwrap().unwrap(); - let payload = response.payload().await.unwrap(); - assert_eq!(payload.as_ref(), large_body.as_slice()); } #[tokio::test] @@ -819,30 +864,39 @@ async fn put_file_with_compression() { let body = "hello compressed file!"; std::fs::write(&path, body).unwrap(); - // Default compression is zstd - let file = tokio::fs::File::open(&path).await.unwrap(); - let stored_id = session.put_file(file).send().await.unwrap().key; - - // When requesting raw, the stored object should be zstd-compressed - let response = session - .get(&stored_id) - .decompress(false) - .send() - .await - .unwrap() - .unwrap(); - assert_eq!(response.metadata.compression, Some(Compression::Zstd)); - - let compressed = response.payload().await.unwrap(); - let decompressed = zstd::bulk::decompress(&compressed, 1024).unwrap(); - assert_eq!(decompressed, body.as_bytes()); - - // Default get decompresses transparently - let response = session.get(&stored_id).send().await.unwrap().unwrap(); - assert_eq!(response.metadata.compression, None); - - let received = response.payload().await.unwrap(); - assert_eq!(received, body); + for threshold in [None, Some(0)] { + let file = tokio::fs::File::open(&path).await.unwrap(); + for put in [session.put_file(file), session.put_path(&path)] { + // Default compression is zstd. + let stored_id = put + .resumable_threshold_bytes(threshold) + .send() + .await + .unwrap() + .key; + + // When requesting raw, the stored object should be zstd-compressed + let response = session + .get(&stored_id) + .decompress(false) + .send() + .await + .unwrap() + .unwrap(); + assert_eq!(response.metadata.compression, Some(Compression::Zstd)); + + let compressed = response.payload().await.unwrap(); + let decompressed = zstd::bulk::decompress(&compressed, 1024).unwrap(); + assert_eq!(decompressed, body.as_bytes()); + + // Default get decompresses transparently + let response = session.get(&stored_id).send().await.unwrap().unwrap(); + assert_eq!(response.metadata.compression, None); + + let received = response.payload().await.unwrap(); + assert_eq!(received, body); + } + } } #[tokio::test] From 64aae596a98b72b607f13e9019df50dbb815fc6b Mon Sep 17 00:00:00 2001 From: lcian <17258265+lcian@users.noreply.github.com> Date: Wed, 21 Oct 2026 14:48:52 +0200 Subject: [PATCH 2/4] ref(rust-client): Simplify automatic upload recovery Separate payload preparation and reader lifetime from protocol recovery. Stop and join each request reader before returning its result, report local IO errors directly, and keep progress probes in a dedicated retry loop. --- clients/rust/src/put.rs | 47 ++-- clients/rust/src/put/automatic.rs | 353 +++++------------------------- clients/rust/src/put/payload.rs | 227 +++++++++++++++++++ 3 files changed, 309 insertions(+), 318 deletions(-) create mode 100644 clients/rust/src/put/payload.rs diff --git a/clients/rust/src/put.rs b/clients/rust/src/put.rs index 04221eca..acb72810 100644 --- a/clients/rust/src/put.rs +++ b/clients/rust/src/put.rs @@ -19,6 +19,7 @@ use crate::response::ResponseExt as _; use crate::{ClientStream, ObjectKey, Session}; mod automatic; +mod payload; /// The response returned from the service after uploading an object. #[derive(Debug, Deserialize)] @@ -405,44 +406,46 @@ impl PutBuilder { self.metadata.compression = self.compression.map(CompressionMode::compression); let headers = self.metadata.to_headers("")?; - let prepared = automatic::prepare( + let prepared = payload::prepare( self.body, self.compression, self.resumable_threshold_bytes, self.session.client.spool_dir.clone(), ) .await?; - let mut content_length = None; - // Keep the file reader alive until a direct fallback request finishes. - let mut replay = None; - let body = match prepared { - automatic::PreparedBody::Direct(body) => body, - automatic::PreparedBody::Replay(body) => { + let session = &self.session; + let key = self.key.as_deref().unwrap_or_default(); + let send_direct = |body: Body, content_length: Option| async move { + // Prepare authentication only when ready to send: spooling can take long enough for + // a token minted before preparation to expire. + let mut builder = session.request(method, key)?.headers(headers); + if let Some(length) = content_length { + builder = builder.header(reqwest::header::CONTENT_LENGTH, length); + } + let response = builder.body(body).send().await?; + Ok(response.error_for_status_and_drain().await?.json().await?) + }; + match prepared { + payload::PreparedBody::Direct(body) => { + send_direct(encode_body(body, self.compression).await?, None).await + } + payload::PreparedBody::Replay(mut body) => { + let length = body.len(); let create = crate::resumable::CreateResumableUploadBuilder { session: self.session.clone(), - total_length: body.len(), + total_length: length, key: self.key.clone(), metadata: self.metadata, }; if let Ok(Some(upload)) = create.send().await { return automatic::upload(body, upload).await; } - content_length = Some(body.len()); - Body::wrap_stream(replay.insert(body).stream(0).await?) + body.with_stream(0, |stream| async { + send_direct(Body::wrap_stream(stream), Some(length)).await + }) + .await } - }; - - // Prepare authentication only when ready to send: spooling can take long enough for - // a token minted before preparation to expire. - let mut builder = self - .session - .request(method, self.key.as_deref().unwrap_or_default())? - .headers(headers); - if let Some(length) = content_length { - builder = builder.header(reqwest::header::CONTENT_LENGTH, length); } - let response = builder.body(body).send().await?; - Ok(response.error_for_status_and_drain().await?.json().await?) } } diff --git a/clients/rust/src/put/automatic.rs b/clients/rust/src/put/automatic.rs index 84f5f2cb..168e3770 100644 --- a/clients/rust/src/put/automatic.rs +++ b/clients/rust/src/put/automatic.rs @@ -1,279 +1,88 @@ -//! Preparation and recovery for automatic resumable uploads. +//! Write prepared payloads, recovering from authoritative server progress. -use std::io::{self, SeekFrom}; -use std::path::PathBuf; +use std::io; use std::time::Duration; -use futures_util::{StreamExt, stream}; -use tokio::fs::File; -use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt}; -use tokio::sync::{mpsc, oneshot}; - -use super::{CompressionMode, PutBody, PutResponse, encode_body}; +use super::PutResponse; +use super::payload::ReplayBody; +use crate::Error; use crate::resumable::{ResumableUpload, UploadProgress}; -use crate::{ClientStream, Error}; - -pub(super) enum PreparedBody { - Direct(reqwest::Body), - Replay(ReplayBody), -} - -pub(super) enum ReplayBody { - Buffer(bytes::Bytes), - File { - file: Option, - start: u64, - length: u64, - }, -} - -/// One reader owns the file for an attempt and returns it after stopping. -pub(super) enum ReplayFile { - Ready(File), - Reading { - stop: oneshot::Sender<()>, - task: tokio::task::JoinHandle, - }, -} - -/// Marks local read failures even when reqwest wraps them in a transport error. -#[derive(Debug, thiserror::Error)] -#[error("could not read upload payload: {0}")] -struct PayloadReadError(#[source] io::Error); - -#[derive(Debug, thiserror::Error)] -#[error("upload made no progress")] -struct NoUploadProgress; - -pub(super) async fn prepare( - mut body: PutBody, - compression: Option, - threshold: Option, - spool_dir: Option, -) -> crate::Result { - if let Some(threshold) = threshold { - if let PutBody::Path(path) = body { - body = PutBody::File(tokio::fs::File::open(path).await?); - } - if let Some(size) = body.remaining_size().await.ok().flatten() - && size >= threshold - { - let mut length = size; - if let Some(CompressionMode::Compress(crate::Compression::Zstd)) = compression { - let temporary = tokio::task::spawn_blocking(move || match spool_dir { - Some(dir) => tempfile::tempfile_in(dir), - None => tempfile::tempfile(), - }) - .await - .map_err(io::Error::other)?; - let mut temporary = match temporary { - Ok(file) => tokio::fs::File::from_std(file), - // No input has been consumed, so direct streaming is still possible. - Err(_) => { - return Ok(PreparedBody::Direct(encode_body(body, compression).await?)); - } - }; - let reader: Box = match body { - PutBody::Buffer(bytes) => Box::new(io::Cursor::new(bytes)), - PutBody::File(file) => Box::new(file), - _ => unreachable!("only buffers and opened regular files have a known size"), - }; - let mut encoder = super::ZstdEncoder::new(super::BufReader::new(reader)); - length = tokio::io::copy(&mut encoder, &mut temporary).await?; - temporary.flush().await?; - temporary.rewind().await?; - body = PutBody::File(temporary); - } - let replay = match body { - PutBody::Buffer(bytes) => ReplayBody::Buffer(bytes), - PutBody::File(mut file) => { - let start = file.stream_position().await?; - ReplayBody::File { - file: Some(ReplayFile::Ready(file)), - start, - length, - } - } - _ => unreachable!("only buffers and opened regular files have a known size"), - }; - return Ok(PreparedBody::Replay(replay)); - } - } - Ok(PreparedBody::Direct(encode_body(body, compression).await?)) -} - -impl ReplayBody { - pub(super) fn len(&self) -> u64 { - match self { - Self::Buffer(bytes) => bytes.len() as u64, - Self::File { length, .. } => *length, - } - } - - /// Creates a non-replayable HTTP stream, waiting for the previous file reader to finish. - pub(super) async fn stream(&mut self, offset: u64) -> io::Result { - match self { - Self::Buffer(bytes) => { - let bytes = bytes.slice(offset as usize..); - Ok(stream::once(async move { Ok(bytes) }).boxed()) - } - Self::File { - file, - start, - length, - } => { - let mut owned = match file.take().expect("file is present between attempts") { - ReplayFile::Ready(file) => file, - ReplayFile::Reading { stop, task } => { - // The HTTP layer may retain the old body after returning a response. - // Stop its reader explicitly, then wait for exclusive file ownership. - drop(stop); - task.await.map_err(io::Error::other)? - } - }; - owned.seek(SeekFrom::Start(*start + offset)).await?; - let (send, receive) = mpsc::channel(1); - let (stop, stopped) = oneshot::channel(); - let length = *length - offset; - let task = tokio::spawn(async move { - { - let read = read_file(&mut owned, length, send); - futures_util::pin_mut!(read, stopped); - let _ = futures_util::future::select(stopped, read).await; - } - // Canceling an async read can leave disk IO running. Wait for it - // before returning the file to the next attempt, without reopening it. - File::from_std(owned.into_std().await) - }); - *file = Some(ReplayFile::Reading { stop, task }); - Ok(stream::unfold(receive, |mut receive| async { - receive.recv().await.map(|bytes| (bytes, receive)) - }) - .boxed()) - } - } - } -} - -/// Reads only while the request accepts data, with at most one buffered chunk. -async fn read_file( - file: &mut File, - mut remaining: u64, - send: mpsc::Sender>, -) { - while remaining > 0 { - let Ok(permit) = send.reserve().await else { - break; - }; - let mut bytes = vec![0; remaining.min(64 * 1024) as usize]; - let result = file.read(&mut bytes).await.and_then(|count| { - if count == 0 { - Err(io::Error::new( - io::ErrorKind::UnexpectedEof, - "upload file became shorter", - )) - } else { - remaining -= count as u64; - bytes.truncate(count); - Ok(bytes::Bytes::from(bytes)) - } - }); - let failed = result.is_err(); - permit.send(result.map_err(|error| io::Error::other(PayloadReadError(error)))); - if failed { - break; - } - } -} pub(super) async fn upload( mut body: ReplayBody, handle: ResumableUpload, ) -> crate::Result { - let result = upload_inner(&mut body, &handle).await; - if result.is_err() { - let _ = handle.cancel().send().await; + match upload_inner(&mut body, &handle).await { + Ok(()) => Ok(PutResponse { + key: handle.key().to_owned(), + }), + Err(error) => { + let _ = handle.cancel().send().await; + Err(Error::UploadFailed(Box::new(error))) + } } - result.map_err(Error::UploadFailed) } -async fn upload_inner( - body: &mut ReplayBody, - handle: &ResumableUpload, -) -> Result> { +async fn upload_inner(body: &mut ReplayBody, handle: &ResumableUpload) -> crate::Result<()> { let mut offset = 0; let mut retries = 0; - let mut probe = false; loop { - let result = if probe { - handle.progress().send().await - } else { - handle - .put_stream(offset, body.len() - offset, body.stream(offset).await?) - .send() - .await - }; - let progress = match result { - Err(Error::UploadOffsetMismatch { offset }) => { - Ok(UploadProgress::Incomplete { offset }) + let remaining = body.len() - offset; + let write = body + .with_stream(offset, |stream| async { + handle.put_stream(offset, remaining, stream).send().await + }) + .await; + let progress = match authoritative_progress(write) { + Ok(UploadProgress::Incomplete { offset: next }) if next == offset => { + let error = io::Error::other("upload made no progress").into(); + recover(handle, &mut retries, error).await? } - other => other, + Err(error) if is_transient(&error) => recover(handle, &mut retries, error).await?, + other => other?, }; - let failure: Box = match progress { - Ok(UploadProgress::Complete) => { - return Ok(PutResponse { - key: handle.key().to_owned(), - }); - } - Ok(UploadProgress::Incomplete { offset: next }) => { + match progress { + UploadProgress::Complete => return Ok(()), + UploadProgress::Incomplete { offset: next } => { if next < offset || next > body.len() { - return Err(Error::MalformedResponse("invalid upload offset".into()).into()); - } - if next == offset && !probe { - Box::new(NoUploadProgress) - } else { - if next > offset { - retries = 0; - } - offset = next; - probe = false; - continue; + return Err(Error::MalformedResponse("invalid upload offset".into())); } - } - Err(error) => { - if !is_transient(&error) { - return Err(error.into()); + if next > offset { + retries = 0; } - error.into() + offset = next; } - }; - if retries >= 2 { - return Err(failure); } - backoff(retries).await; - retries += 1; - probe = true; } } -async fn backoff(retries: u32) { - let delay = Duration::from_secs(2 << retries) + Duration::from_secs_f64(rand::random::()); - tokio::time::sleep(delay).await; +/// Writes and failed probes share this budget; only acknowledged progress resets it. +async fn recover( + handle: &ResumableUpload, + retries: &mut u32, + mut failure: Error, +) -> crate::Result { + while *retries < 2 { + let delay = + Duration::from_secs(2 << *retries) + Duration::from_secs_f64(rand::random::()); + tokio::time::sleep(delay).await; + *retries += 1; + match authoritative_progress(handle.progress().send().await) { + Err(error) if is_transient(&error) => failure = error, + result => return result, + } + } + Err(failure) +} + +fn authoritative_progress(result: crate::Result) -> crate::Result { + match result { + Err(Error::UploadOffsetMismatch { offset }) => Ok(UploadProgress::Incomplete { offset }), + other => other, + } } fn is_transient(error: &Error) -> bool { - let mut source: Option<&(dyn std::error::Error + 'static)> = Some(error); - while let Some(error) = source { - if error.is::() - || error - .downcast_ref::() - .and_then(io::Error::get_ref) - .is_some_and(|inner| inner.is::()) - { - return false; - } - source = error.source(); - } match error { Error::Reqwest(error) => match error.status() { Some(status) => matches!(status.as_u16(), 408 | 429 | 502 | 503 | 504), @@ -284,51 +93,3 @@ fn is_transient(error: &Error) -> bool { _ => false, } } - -#[cfg(test)] -mod tests { - use futures_util::TryStreamExt; - - use super::*; - - #[test] - fn retry_stops_previous_reader_and_waits_for_pending_read() { - let runtime = tokio::runtime::Builder::new_current_thread() - .enable_all() - .max_blocking_threads(1) - .build() - .unwrap(); - runtime.block_on(async { - let data: Vec = (0..200_000).map(|n| (n % 251) as u8).collect(); - let mut file = tempfile::tempfile().unwrap(); - std::io::Write::write_all(&mut file, &data).unwrap(); - let mut body = ReplayBody::File { - file: Some(ReplayFile::Ready(File::from_std(file))), - start: 7, - length: 150_000, - }; - let mut old = body.stream(0).await.unwrap(); - - // Occupy the sole blocking worker so the old read remains pending. - let (release, blocked) = std::sync::mpsc::channel(); - let (started, running) = oneshot::channel(); - let worker = tokio::task::spawn_blocking(move || { - started.send(()).unwrap(); - blocked.recv().unwrap(); - }); - running.await.unwrap(); - tokio::task::yield_now().await; - assert!(futures_util::poll!(old.next()).is_pending()); - // Retain the old HTTP body: retry must stop its reader explicitly. - let retry = body.stream(100); - tokio::pin!(retry); - assert!(futures_util::poll!(&mut retry).is_pending()); - release.send(()).unwrap(); - worker.await.unwrap(); - let retry = retry.await.unwrap(); - let remaining: Vec<_> = retry.try_collect().await.unwrap(); - assert_eq!(remaining.concat(), data[107..150_007]); - drop(old); - }); - } -} diff --git a/clients/rust/src/put/payload.rs b/clients/rust/src/put/payload.rs new file mode 100644 index 00000000..d9eb9d45 --- /dev/null +++ b/clients/rust/src/put/payload.rs @@ -0,0 +1,227 @@ +//! Prepare replayable wire bytes and give each file request exclusive reader ownership. + +use std::io::{self, SeekFrom}; +use std::path::PathBuf; + +use bytes::Bytes; +use futures_util::{StreamExt, stream}; +use tokio::fs::File; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeekExt, AsyncWriteExt}; +use tokio::sync::{mpsc, oneshot}; + +use super::{CompressionMode, PutBody}; +use crate::ClientStream; + +pub(super) enum PreparedBody { + Direct(PutBody), + Replay(ReplayBody), +} + +pub(super) enum ReplayBody { + Buffer(Bytes), + File { file: File, start: u64, length: u64 }, +} + +pub(super) async fn prepare( + mut body: PutBody, + compression: Option, + threshold: Option, + spool_dir: Option, +) -> crate::Result { + let Some(threshold) = threshold else { + return Ok(PreparedBody::Direct(body)); + }; + if let PutBody::Path(path) = body { + body = PutBody::File(File::open(path).await?); + } + let Some(mut length) = body.remaining_size().await.ok().flatten() else { + return Ok(PreparedBody::Direct(body)); + }; + if length < threshold { + return Ok(PreparedBody::Direct(body)); + } + + if let Some(CompressionMode::Compress(crate::Compression::Zstd)) = compression { + let temporary = tokio::task::spawn_blocking(move || match spool_dir { + Some(dir) => tempfile::tempfile_in(dir), + None => tempfile::tempfile(), + }) + .await + .map_err(io::Error::other)?; + let Ok(temporary) = temporary else { + // Creation has not consumed input, so direct streaming is still possible. + return Ok(PreparedBody::Direct(body)); + }; + let mut temporary = File::from_std(temporary); + let reader: Box = match body { + PutBody::Buffer(bytes) => Box::new(io::Cursor::new(bytes)), + PutBody::File(file) => Box::new(file), + _ => unreachable!("only buffers and opened regular files have a known size"), + }; + let mut encoder = super::ZstdEncoder::new(super::BufReader::new(reader)); + length = tokio::io::copy(&mut encoder, &mut temporary).await?; + temporary.flush().await?; + temporary.rewind().await?; + body = PutBody::File(temporary); + } + + Ok(PreparedBody::Replay(match body { + PutBody::Buffer(bytes) => ReplayBody::Buffer(bytes), + PutBody::File(mut file) => ReplayBody::File { + start: file.stream_position().await?, + file, + length, + }, + _ => unreachable!("only buffers and opened regular files have a known size"), + })) +} + +impl ReplayBody { + pub(super) fn len(&self) -> u64 { + match self { + Self::Buffer(bytes) => bytes.len() as u64, + Self::File { length, .. } => *length, + } + } + + /// Runs a request with a suffix, then stops its reader before returning the result. + /// + /// The request body may outlive the response. Keeping reader ownership here makes local + /// read errors explicit and prevents an old request from reading after a retry seeks. + pub(super) async fn with_stream( + &mut self, + offset: u64, + request: impl FnOnce(ClientStream) -> F, + ) -> crate::Result + where + F: std::future::Future>, + { + match self { + Self::Buffer(bytes) => { + let bytes = bytes.slice(offset as usize..); + // Streaming bodies prevent independent replays by HTTP middleware. + request(stream::once(async move { Ok(bytes) }).boxed()).await + } + Self::File { + file, + start, + length, + } => { + let (stream, reader) = + FileReader::start(file, *start + offset, *length - offset).await?; + let response = request(stream).await; + // Read failures are returned directly, independently of reqwest's error wrapping. + reader.finish().await?; + response + } + } + } +} + +/// Explicitly stoppable because HTTP may retain the request body after returning a response. +/// Dropping this handle signals cancellation; finishing also waits for pending disk IO. +struct FileReader { + stop: oneshot::Sender<()>, + task: tokio::task::JoinHandle>, +} + +impl FileReader { + async fn start(file: &File, offset: u64, length: u64) -> io::Result<(ClientStream, Self)> { + // Duplicate the open handle, never reopen a path. Clones share the OS cursor, so + // with_stream must finish this reader before starting another request. + let mut file = file.try_clone().await?; + file.seek(SeekFrom::Start(offset)).await?; + let (send, receive) = mpsc::channel(1); + let (stop, stopped) = oneshot::channel(); + let task = tokio::spawn(async move { + let result = { + let read = forward_file(&mut file, length, send); + futures_util::pin_mut!(read, stopped); + match futures_util::future::select(stopped, read).await { + futures_util::future::Either::Left(_) => Ok(()), + futures_util::future::Either::Right((result, _)) => result, + } + }; + // An async read may leave blocking IO running after cancellation. + drop(file.into_std().await); + result + }); + let stream = stream::unfold(receive, |mut receive| async { + receive.recv().await.map(|bytes| (bytes, receive)) + }) + .boxed(); + Ok((stream, Self { stop, task })) + } + + async fn finish(self) -> io::Result<()> { + drop(self.stop); + self.task.await.map_err(io::Error::other)? + } +} + +async fn forward_file( + file: &mut File, + length: u64, + send: mpsc::Sender>, +) -> io::Result<()> { + let mut remaining = length; + while remaining > 0 { + let Ok(permit) = send.reserve().await else { + break; + }; + let mut bytes = vec![0; remaining.min(64 * 1024) as usize]; + if let Err(error) = file.read_exact(&mut bytes).await { + // Fail the HTTP body too, but keep the original error for the caller. + permit.send(Err(io::Error::other(error.to_string()))); + return Err(error); + } + remaining -= bytes.len() as u64; + permit.send(Ok(Bytes::from(bytes))); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use futures_util::TryStreamExt; + + use super::*; + + #[test] + fn stopping_reader_waits_for_pending_read_even_if_http_keeps_body() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .max_blocking_threads(1) + .build() + .unwrap(); + runtime.block_on(async { + let data: Vec = (0..200_000).map(|n| (n % 251) as u8).collect(); + let mut file = tempfile::tempfile().unwrap(); + std::io::Write::write_all(&mut file, &data).unwrap(); + let file = File::from_std(file); + let (mut old, reader) = FileReader::start(&file, 7, 150_000).await.unwrap(); + + // Occupy the sole blocking worker so the old read remains pending. + let (release, blocked) = std::sync::mpsc::channel(); + let (started, running) = oneshot::channel(); + let worker = tokio::task::spawn_blocking(move || { + started.send(()).unwrap(); + blocked.recv().unwrap(); + }); + running.await.unwrap(); + tokio::task::yield_now().await; + assert!(futures_util::poll!(old.next()).is_pending()); + let finish = reader.finish(); + tokio::pin!(finish); + assert!(futures_util::poll!(&mut finish).is_pending()); + release.send(()).unwrap(); + worker.await.unwrap(); + finish.await.unwrap(); + let (stream, reader) = FileReader::start(&file, 107, 149_900).await.unwrap(); + let remaining: Vec<_> = stream.try_collect().await.unwrap(); + reader.finish().await.unwrap(); + assert_eq!(remaining.concat(), data[107..150_007]); + drop(old); + }); + } +} From e0e411384f29eaa9c4a0ee286d8838da4c263f54 Mon Sep 17 00:00:00 2001 From: lcian <17258265+lcian@users.noreply.github.com> Date: Wed, 21 Oct 2026 15:48:03 +0200 Subject: [PATCH 3/4] ref(rust-client): Address automatic upload API feedback Unify direct fallbacks, make sizing optional, and limit resumables to standalone uploads. --- clients/rust/README.md | 11 +-- clients/rust/src/client.rs | 4 +- clients/rust/src/many.rs | 97 +++++++------------------ clients/rust/src/put.rs | 117 +++++++++++++----------------- clients/rust/src/put/automatic.rs | 4 +- clients/rust/src/put/payload.rs | 25 ++++++- clients/rust/tests/automatic.rs | 14 ++-- clients/rust/tests/e2e.rs | 109 +++++++++++----------------- 8 files changed, 160 insertions(+), 221 deletions(-) diff --git a/clients/rust/README.md b/clients/rust/README.md index 317fa224..1d5fb86f 100644 --- a/clients/rust/README.md +++ b/clients/rust/README.md @@ -153,12 +153,13 @@ session.put("payload") ### Automatic Resumable Uploads -Large buffer and regular-file uploads automatically use resumable uploads, including uploads -queued through `many()`. The default threshold is 256 MiB before client compression. Compressed -payloads are prepared in temporary files; uncompressed and precompressed inputs are reused. -Streams and generic readers use direct uploads. No feature flag is required. +Large buffer and regular-file uploads automatically use resumable uploads. The default +threshold is 256 MiB before client compression. Compressed payloads are prepared in temporary +files; uncompressed and precompressed inputs are reused. +Streams, generic readers, and uploads queued through `many()` do not use resumables. +No feature flag is required. -See [`PutBuilder`] for eligibility, fallback, and recovery behavior, +See [`PutBuilder`] for eligibility and recovery behavior, [`Usecase::with_resumable_threshold_bytes`] and [`PutBuilder::resumable_threshold_bytes`] to configure or disable it, and [`ClientBuilder::spool_dir`] to choose the temporary-file directory. diff --git a/clients/rust/src/client.rs b/clients/rust/src/client.rs index 3708cbb8..566e3ba3 100644 --- a/clients/rust/src/client.rs +++ b/clients/rust/src/client.rs @@ -258,8 +258,8 @@ impl Usecase { /// Sets the minimum remaining input size for automatic resumable uploads. /// /// Defaults to 256 MiB, measured before client compression. Pass [`None`] to disable - /// automatic resumables or zero to include every supported payload. This also applies to - /// uploads queued through [`Session::many`]. See [`crate::PutBuilder`] for eligibility. + /// automatic resumables or zero to include every supported payload. Uploads queued through + /// [`Session::many`] ignore this setting. See [`crate::PutBuilder`] for eligibility. /// Individual uploads can override this using [`crate::PutBuilder::resumable_threshold_bytes`]. pub fn with_resumable_threshold_bytes(mut self, threshold: impl Into>) -> Self { self.resumable_threshold_bytes = threshold.into(); diff --git a/clients/rust/src/many.rs b/clients/rust/src/many.rs index d1b6d3f9..bdc2a590 100644 --- a/clients/rust/src/many.rs +++ b/clients/rust/src/many.rs @@ -84,7 +84,6 @@ enum BatchOperation { metadata: Metadata, compression: Option, body: PutBody, - resumable_threshold_bytes: Option, }, Delete { key: ObjectKey, @@ -117,15 +116,14 @@ impl From for BatchOperation { metadata, compression, body, - resumable_threshold_bytes, session: _session, + resumable_threshold_bytes: _, } = value; BatchOperation::Insert { key, metadata, compression, body, - resumable_threshold_bytes, } } } @@ -162,7 +160,6 @@ impl BatchOperation { mut metadata, compression, body, - resumable_threshold_bytes: _, } => { let mut headers = operation_headers("insert", key.as_deref()); metadata.compression = compression.map(CompressionMode::compression); @@ -535,9 +532,17 @@ async fn send_batch( Ok(results) } +fn classify_fail(key: Option, error: Error) -> Classified { + Classified::Failed(OperationResult::Put( + key.unwrap_or_else(|| "".to_owned()), + Err(error), + )) +} + /// Classifies a single operation for batch processing. /// -/// Oversized, unknown-size, and resumable-eligible uploads are sent individually. +/// Insert operations whose body exceeds [`MAX_BATCH_PART_SIZE`] are marked as +/// [`Classified::Individual`]. Everything else is [`Classified::Batchable`]. async fn classify(op: BatchOperation) -> Classified { match op { BatchOperation::Insert { @@ -545,22 +550,20 @@ async fn classify(op: BatchOperation) -> Classified { metadata, compression, body, - resumable_threshold_bytes, } => { - let mut body = body; - let size = match body.remaining_size().await { - Ok(size) => size, - Err(error) if matches!(body, PutBody::Path(_)) => { - return Classified::Failed(error_result( - OperationContext::Insert { key }, - error.into(), - )); - } - Err(_) => None, + let size = match &body { + PutBody::Buffer(bytes) => Some(bytes.len() as u64), + PutBody::File(file) => match file.metadata().await { + Ok(meta) => Some(meta.len()), + Err(err) => return classify_fail(key, err.into()), + }, + PutBody::Path(path) => match tokio::fs::metadata(path).await { + Ok(meta) => Some(meta.len()), + Err(err) => return classify_fail(key, err.into()), + }, + // Streams have unknown size and must not go through the batch endpoint. + PutBody::Stream(_) => None, }; - let resumable = size - .zip(resumable_threshold_bytes) - .is_some_and(|(size, threshold)| size >= threshold); // Client-side compression happens while streaming the body, so the on-wire size is // only known up to its worst case. Precompressed bodies are sent as-is. @@ -579,12 +582,8 @@ async fn classify(op: BatchOperation) -> Classified { metadata, compression, body, - resumable_threshold_bytes, }; - if resumable { - return Classified::Individual(op); - } match size { Some(s) if s <= MAX_BATCH_PART_SIZE as usize => Classified::Batchable(op, s as u64), _ => Classified::Individual(op), @@ -639,16 +638,16 @@ async fn execute_individual(op: BatchOperation, session: &Session) -> OperationR metadata, compression, body, - resumable_threshold_bytes, } => { let error_key = key.clone().unwrap_or_else(|| "".to_owned()); let put = PutBuilder { + // Automatic resumables are currently limited to standalone uploads. + resumable_threshold_bytes: None, session: session.clone(), metadata, compression, key, body, - resumable_threshold_bytes, }; match put.send().await { Ok(response) => OperationResult::Put(response.key.clone(), Ok(response)), @@ -764,8 +763,7 @@ impl ManyBuilder { /// Sets the maximum number of concurrent individual (non-batch) requests. /// - /// Operations that exceed the per-part size limit or qualify for resumable uploads - /// are sent as individual requests. + /// Operations that exceed the per-part size limit are sent as individual requests. /// This controls how many such requests can be in-flight simultaneously. /// Defaults to 5 if not set. pub fn max_individual_concurrency(mut self, concurrency: usize) -> Self { @@ -827,7 +825,6 @@ mod tests { metadata: Metadata::default(), compression: Some(compression), body: PutBody::Buffer(vec![0; size].into()), - resumable_threshold_bytes: None, } } @@ -873,50 +870,6 @@ mod tests { ); } - #[tokio::test] - async fn threshold_overrides_control_batch_classification() { - let session = crate::Client::builder("http://localhost") - .build() - .unwrap() - .session( - crate::Usecase::new("testing") - .with_resumable_threshold_bytes(4) - .for_organization(123), - ) - .unwrap(); - assert!(matches!( - classify(session.put("data").into()).await, - Classified::Individual(_) - )); - assert!(matches!( - classify(session.put("data").resumable_threshold_bytes(None).into()).await, - Classified::Batchable(_, _) - )); - assert!(matches!( - classify(session.put("data").resumable_threshold_bytes(5).into()).await, - Classified::Batchable(_, _) - )); - let file = tempfile::NamedTempFile::new().unwrap(); - std::fs::write(file.path(), b"prefix-data").unwrap(); - let mut input = tokio::fs::File::open(file.path()).await.unwrap(); - tokio::io::AsyncSeekExt::seek(&mut input, std::io::SeekFrom::Start(8)) - .await - .unwrap(); - assert!(matches!( - classify(session.put_file(input).into()).await, - Classified::Batchable(_, _) - )); - let Classified::Individual(BatchOperation::Insert { body, .. }) = - classify(session.put_path(file.path()).into()).await - else { - panic!("eligible path should be sent individually"); - }; - assert!( - matches!(body, PutBody::Path(_)), - "classification must not open paths" - ); - } - #[test] fn iter_batches_empty() { assert!(batches(vec![]).is_empty()); diff --git a/clients/rust/src/put.rs b/clients/rust/src/put.rs index acb72810..db492f0d 100644 --- a/clients/rust/src/put.rs +++ b/clients/rust/src/put.rs @@ -36,25 +36,27 @@ pub(crate) enum PutBody { } impl PutBody { - /// Measures remaining bytes without opening queued paths or consuming a stream. - pub(crate) async fn remaining_size(&mut self) -> io::Result> { + /// Returns the remaining byte length, or `None` when it cannot be determined. + pub(crate) async fn len(&mut self) -> Option { match self { - Self::Buffer(bytes) => Ok(Some(bytes.len() as u64)), + Self::Buffer(bytes) => Some(bytes.len() as u64), Self::File(file) => { - file.flush().await?; - let metadata = file.metadata().await?; + file.flush().await.ok()?; + let metadata = file.metadata().await.ok()?; if !metadata.is_file() { - return Ok(None); + return None; } - Ok(Some( - metadata.len().saturating_sub(file.stream_position().await?), - )) + Some( + metadata + .len() + .saturating_sub(file.stream_position().await.ok()?), + ) } Self::Path(path) => { - let metadata = tokio::fs::metadata(path).await?; - Ok(metadata.is_file().then_some(metadata.len())) + let metadata = tokio::fs::metadata(path).await.ok()?; + metadata.is_file().then_some(metadata.len()) } - Self::Stream(_) => Ok(None), + Self::Stream(_) => None, } } } @@ -160,27 +162,17 @@ impl Session { /// /// # Automatic resumable uploads /// -/// Buffers and regular files (including paths) use resumable uploads when their remaining -/// input size meets the configured threshold, by default 256 MiB before client compression. -/// Streams, generic readers, and files whose size cannot be determined use direct uploads. -/// Files must remain unchanged until the upload finishes. File uploads start at the current -/// position; path uploads open the file only when executed. +/// Buffers and regular files (including paths) use resumable uploads at 256 MiB by default, +/// measured from the current position before client compression. Override or disable this with +/// [`Self::resumable_threshold_bytes`]. Streams, generic readers, and files of unknown size use +/// direct uploads. Files must remain unchanged during upload. /// -/// Eligible uploads reuse uncompressed or precompressed input. Client compression first spools -/// the complete compressed payload to a temporary file, using [`crate::ClientBuilder::spool_dir`] -/// when configured. If file creation fails, the client falls back to direct streaming. Later -/// preparation errors propagate. Unsupported or failed session creation also falls back to a -/// direct upload of the prepared bytes, without compressing them again. +/// Input is reused unless compression is needed, in which case the compressed payload is spooled +/// to a temporary file. Use [`crate::ClientBuilder::spool_dir`] to choose its directory. +/// Transient failures and writes making no progress are retried up to twice. Acknowledged +/// progress resets the retry count. /// -/// Transient failures and writes making no progress are retried twice, with exponential delays -/// of 2 and 4 seconds plus up to 1 second of jitter. Recovery queries the server's persisted -/// offset before replaying data. The retry budget resets whenever that offset advances. -/// Terminal failures after session creation attempt cancellation and return -/// [`crate::Error::UploadFailed`] preserving the cause. Dropping the future releases local -/// resources; the remote session is left to expire. -/// -/// Automatic uploads require no feature flag. The `resumable-upload-api` feature exposes the -/// manual protocol API. See [`Self::resumable_threshold_bytes`] to override eligibility. +/// No feature flag is required. Uploads queued through [`Session::many`] do not use resumables. #[derive(Debug)] pub struct PutBuilder { pub(crate) session: Session, @@ -195,9 +187,8 @@ impl PutBuilder { /// Sets the minimum remaining input size for automatic resumable uploads. /// /// Inherits [`crate::Usecase::resumable_threshold_bytes`] unless overridden. Pass [`None`] - /// to disable or zero to include every supported payload. Eligible uploads queued through - /// [`Session::many`] are sent individually rather than batched. See [`PutBuilder`] for - /// preparation and recovery behavior. + /// to disable or zero to include every supported payload. Uploads queued through + /// [`Session::many`] ignore this setting. See [`PutBuilder`] for preparation and recovery. /// /// ```no_run /// # async fn example(session: objectstore_client::Session) -> objectstore_client::Result<()> { @@ -406,46 +397,40 @@ impl PutBuilder { self.metadata.compression = self.compression.map(CompressionMode::compression); let headers = self.metadata.to_headers("")?; - let prepared = payload::prepare( + let mut prepared = payload::prepare( self.body, self.compression, self.resumable_threshold_bytes, self.session.client.spool_dir.clone(), ) .await?; - let session = &self.session; - let key = self.key.as_deref().unwrap_or_default(); - let send_direct = |body: Body, content_length: Option| async move { - // Prepare authentication only when ready to send: spooling can take long enough for - // a token minted before preparation to expire. - let mut builder = session.request(method, key)?.headers(headers); - if let Some(length) = content_length { - builder = builder.header(reqwest::header::CONTENT_LENGTH, length); - } - let response = builder.body(body).send().await?; - Ok(response.error_for_status_and_drain().await?.json().await?) - }; - match prepared { - payload::PreparedBody::Direct(body) => { - send_direct(encode_body(body, self.compression).await?, None).await - } - payload::PreparedBody::Replay(mut body) => { - let length = body.len(); - let create = crate::resumable::CreateResumableUploadBuilder { - session: self.session.clone(), - total_length: length, - key: self.key.clone(), - metadata: self.metadata, - }; - if let Ok(Some(upload)) = create.send().await { - return automatic::upload(body, upload).await; - } - body.with_stream(0, |stream| async { - send_direct(Body::wrap_stream(stream), Some(length)).await - }) - .await + if let payload::PreparedBody::Replay(body) = &mut prepared { + let create = crate::resumable::CreateResumableUploadBuilder { + session: self.session.clone(), + total_length: body.len(), + key: self.key.clone(), + metadata: self.metadata, + }; + if let Ok(Some(upload)) = create.send().await { + return automatic::upload(body, upload).await; } } + + // All direct uploads, including preparation and session-creation fallbacks, converge here. + prepared + .with_body(self.compression, |body, content_length| async move { + // Mint authentication after preparation, which may outlast a token's validity. + let mut builder = self + .session + .request(method, self.key.as_deref().unwrap_or_default())? + .headers(headers); + if let Some(length) = content_length { + builder = builder.header(reqwest::header::CONTENT_LENGTH, length); + } + let response = builder.body(body).send().await?; + Ok(response.error_for_status_and_drain().await?.json().await?) + }) + .await } } diff --git a/clients/rust/src/put/automatic.rs b/clients/rust/src/put/automatic.rs index 168e3770..55d71f6c 100644 --- a/clients/rust/src/put/automatic.rs +++ b/clients/rust/src/put/automatic.rs @@ -9,10 +9,10 @@ use crate::Error; use crate::resumable::{ResumableUpload, UploadProgress}; pub(super) async fn upload( - mut body: ReplayBody, + body: &mut ReplayBody, handle: ResumableUpload, ) -> crate::Result { - match upload_inner(&mut body, &handle).await { + match upload_inner(body, &handle).await { Ok(()) => Ok(PutResponse { key: handle.key().to_owned(), }), diff --git a/clients/rust/src/put/payload.rs b/clients/rust/src/put/payload.rs index d9eb9d45..d6465c3a 100644 --- a/clients/rust/src/put/payload.rs +++ b/clients/rust/src/put/payload.rs @@ -17,6 +17,29 @@ pub(super) enum PreparedBody { Replay(ReplayBody), } +impl PreparedBody { + /// Supplies a direct request with either the original input or prepared wire bytes. + pub(super) async fn with_body( + self, + compression: Option, + request: impl FnOnce(reqwest::Body, Option) -> F, + ) -> crate::Result + where + F: std::future::Future>, + { + match self { + Self::Direct(body) => request(super::encode_body(body, compression).await?, None).await, + Self::Replay(mut body) => { + let length = body.len(); + body.with_stream(0, |stream| { + request(reqwest::Body::wrap_stream(stream), Some(length)) + }) + .await + } + } + } +} + pub(super) enum ReplayBody { Buffer(Bytes), File { file: File, start: u64, length: u64 }, @@ -34,7 +57,7 @@ pub(super) async fn prepare( if let PutBody::Path(path) = body { body = PutBody::File(File::open(path).await?); } - let Some(mut length) = body.remaining_size().await.ok().flatten() else { + let Some(mut length) = body.len().await else { return Ok(PreparedBody::Direct(body)); }; if length < threshold { diff --git a/clients/rust/tests/automatic.rs b/clients/rust/tests/automatic.rs index d9b60cff..99701f30 100644 --- a/clients/rust/tests/automatic.rs +++ b/clients/rust/tests/automatic.rs @@ -372,12 +372,13 @@ async fn missing_spool_directory_falls_back_and_precompressed_input_needs_no_spo } #[tokio::test] -async fn many_honors_threshold_overrides() { - let server = Server::new([Reply::Create, Reply::Complete]).await; - let session = server.session(None); +async fn many_keeps_oversized_uploads_direct() { + let server = Server::new([Reply::Status(200)]).await; + let session = server.session(Some(0)); + let data = Bytes::from(vec![b'x'; 1024 * 1024 + 1]); session .many() - .push(session.put("data").resumable_threshold_bytes(4)) + .push(session.put(data.clone()).resumable_threshold_bytes(0)) .send() .await .error_for_failures() @@ -385,8 +386,9 @@ async fn many_honors_threshold_overrides() { .map_err(|errors| errors.collect::>()) .unwrap(); let requests = server.finish(); - assert_eq!(requests[0].uri.query(), Some("upload_type=resumable")); - assert_eq!(requests[1].body, "data"); + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].uri.query(), None); + assert_eq!(requests[0].body, data); } #[tokio::test] diff --git a/clients/rust/tests/e2e.rs b/clients/rust/tests/e2e.rs index 4b27c91a..52ea57a2 100644 --- a/clients/rust/tests/e2e.rs +++ b/clients/rust/tests/e2e.rs @@ -504,8 +504,8 @@ async fn batch_operations() { let usecase = Usecase::new("usecase"); let session = client.session(usecase.for_project(12345, 1337)).unwrap(); - // First batch: batched PUTs alongside one automatic resumable upload. - // key-2 uses default compression; key-5 is precompressed. + // First batch: 4 PUT operations + // key-2 uses default compression (zstd), others are uncompressed let results: Vec<_> = session .many() .push( @@ -518,19 +518,12 @@ async fn batch_operations() { .push(session.put("second object").key("key-2")) .push(session.put("third object").compress(None).key("key-3")) .push(session.put("fourth object").compress(None).key("key-4")) - .push( - session - .put(zstd::encode_all(&b"fifth object"[..], 0).unwrap()) - .precompressed(Compression::Zstd) - .resumable_threshold_bytes(0) - .key("key-5"), - ) .send() .await .collect() .await; - assert_eq!(results.len(), 5); + assert_eq!(results.len(), 4); let mut keys: Vec = results .iter() .map(|r| match r { @@ -539,7 +532,7 @@ async fn batch_operations() { }) .collect(); keys.sort(); - assert_eq!(keys, vec!["key-1", "key-2", "key-3", "key-4", "key-5"]); + assert_eq!(keys, vec!["key-1", "key-2", "key-3", "key-4"]); // Second batch: GET key-1, GET key-2 (automatic decompression), DELETE key-3, PUT key-4 (override) let results: Vec<_> = session @@ -615,10 +608,6 @@ async fn batch_operations() { // Verify the deleted object is gone let response = session.get("key-3").send().await.unwrap(); assert!(response.is_none()); - - // The precompressed automatic upload remains readable through normal downloads. - let response = session.get("key-5").send().await.unwrap().unwrap(); - assert_eq!(response.payload().await.unwrap(), "fifth object"); } #[tokio::test] @@ -763,62 +752,48 @@ async fn batch_put_files() { // This should be routed to a single PUT request instead of being batched. let large_body: Vec = vec![0xAB; 2 * 1024 * 1024]; // 2 MB - for threshold in [None, Some(0)] { - let mut many = session.many(); - - for (name, content) in &small_bodies { - let path = dir.path().join(name); - let mut f = std::fs::File::create(&path).unwrap(); - f.write_all(content).unwrap(); - let file = tokio::fs::File::open(&path).await.unwrap(); - many = many.push( - session - .put_file(file) - .resumable_threshold_bytes(threshold) - .compress(None) - .key(*name), - ); - } + let mut many = session.many(); - let large_path = dir.path().join("large"); - std::fs::File::create(&large_path) - .unwrap() - .write_all(&large_body) - .unwrap(); - let large_file = tokio::fs::File::open(&large_path).await.unwrap(); - many = many.push( - session - .put_file(large_file) - .resumable_threshold_bytes(threshold) - .compress(None) - .key("large"), - ); - - let results: Vec<_> = many.send().await.collect().await; - - assert_eq!(results.len(), 4); - - let mut keys: Vec = results - .iter() - .map(|r| match r { - OperationResult::Put(key, Ok(_)) => key.clone(), - other => panic!("Expected successful Put result, got: {other:?}"), - }) - .collect(); - keys.sort(); - assert_eq!(keys, vec!["large", "small-1", "small-2", "small-3"]); - - // Verify all objects are retrievable with the correct content. - for (name, expected) in &small_bodies { - let response = session.get(name).send().await.unwrap().unwrap(); - let payload = response.payload().await.unwrap(); - assert_eq!(payload.as_ref(), expected.as_slice(), "mismatch for {name}"); - } + for (name, content) in &small_bodies { + let path = dir.path().join(name); + let mut f = std::fs::File::create(&path).unwrap(); + f.write_all(content).unwrap(); + let file = tokio::fs::File::open(&path).await.unwrap(); + many = many.push(session.put_file(file).compress(None).key(*name)); + } - let response = session.get("large").send().await.unwrap().unwrap(); + let large_path = dir.path().join("large"); + std::fs::File::create(&large_path) + .unwrap() + .write_all(&large_body) + .unwrap(); + let large_file = tokio::fs::File::open(&large_path).await.unwrap(); + many = many.push(session.put_file(large_file).compress(None).key("large")); + + let results: Vec<_> = many.send().await.collect().await; + + assert_eq!(results.len(), 4); + + let mut keys: Vec = results + .iter() + .map(|r| match r { + OperationResult::Put(key, Ok(_)) => key.clone(), + other => panic!("Expected successful Put result, got: {other:?}"), + }) + .collect(); + keys.sort(); + assert_eq!(keys, vec!["large", "small-1", "small-2", "small-3"]); + + // Verify all objects are retrievable with the correct content. + for (name, expected) in &small_bodies { + let response = session.get(name).send().await.unwrap().unwrap(); let payload = response.payload().await.unwrap(); - assert_eq!(payload.as_ref(), large_body.as_slice()); + assert_eq!(payload.as_ref(), expected.as_slice(), "mismatch for {name}"); } + + let response = session.get("large").send().await.unwrap().unwrap(); + let payload = response.payload().await.unwrap(); + assert_eq!(payload.as_ref(), large_body.as_slice()); } #[tokio::test] From b06ec0ba812579cc770d14484ba2e1aa1fac5d30 Mon Sep 17 00:00:00 2001 From: "sentry-junior[bot]" <264270552+sentry-junior[bot]@users.noreply.github.com> Date: Fri, 23 Oct 2026 11:30:25 +0000 Subject: [PATCH 4/4] ref(rs-client): Simplify automatic upload preparation and trim duplicated docs Co-Authored-By: Lorenzo Cian <17258265+lcian@users.noreply.github.com> --- clients/rust/README.md | 13 ++---- clients/rust/src/client.rs | 7 ++-- clients/rust/src/put.rs | 31 +------------- clients/rust/src/put/payload.rs | 72 ++++++++++++++++++++++----------- clients/rust/tests/automatic.rs | 2 +- 5 files changed, 58 insertions(+), 67 deletions(-) diff --git a/clients/rust/README.md b/clients/rust/README.md index e6ffda98..b3f2c57d 100644 --- a/clients/rust/README.md +++ b/clients/rust/README.md @@ -153,15 +153,10 @@ session.put("payload") ### Automatic Resumable Uploads -Large buffer and regular-file uploads automatically use resumable uploads. The default -threshold is 256 MiB before client compression. Compressed payloads are prepared in temporary -files; uncompressed and precompressed inputs are reused. -Streams, generic readers, and uploads queued through `many()` do not use resumables. -No feature flag is required. - -See [`PutBuilder`] for eligibility and recovery behavior, -[`Usecase::with_resumable_threshold_bytes`] and [`PutBuilder::resumable_threshold_bytes`] to -configure or disable it, and [`ClientBuilder::spool_dir`] to choose the temporary-file directory. +Large uploads automatically use resumable uploads and resend only what the server hasn't +stored after a failed write. No feature flag is required. + +See [`PutBuilder`] for eligibility, configuration, and recovery behavior. ### Resumable Upload API diff --git a/clients/rust/src/client.rs b/clients/rust/src/client.rs index bc20a185..6f0fbb26 100644 --- a/clients/rust/src/client.rs +++ b/clients/rust/src/client.rs @@ -257,10 +257,9 @@ impl Usecase { /// Sets the minimum remaining input size for automatic resumable uploads. /// - /// Defaults to 256 MiB, measured before client compression. Pass [`None`] to disable - /// automatic resumables or zero to include every supported payload. Uploads queued through - /// [`Session::many`] ignore this setting. See [`crate::PutBuilder`] for eligibility. - /// Individual uploads can override this using [`crate::PutBuilder::resumable_threshold_bytes`]. + /// Defaults to 256 MiB. Pass [`None`] to disable automatic resumables or zero to include + /// every supported payload. Individual uploads can override this using + /// [`crate::PutBuilder::resumable_threshold_bytes`]. See [`crate::PutBuilder`] for eligibility. pub fn with_resumable_threshold_bytes(mut self, threshold: impl Into>) -> Self { self.resumable_threshold_bytes = threshold.into(); self diff --git a/clients/rust/src/put.rs b/clients/rust/src/put.rs index 91c1b233..24953663 100644 --- a/clients/rust/src/put.rs +++ b/clients/rust/src/put.rs @@ -10,7 +10,7 @@ use objectstore_types::metadata::Metadata; use reqwest::Body; use serde::Deserialize; use tokio::fs::File; -use tokio::io::{AsyncRead, AsyncSeekExt, AsyncWriteExt, BufReader}; +use tokio::io::{AsyncRead, BufReader}; use tokio_util::io::{ReaderStream, StreamReader}; pub use objectstore_types::metadata::{Compression, ExpirationPolicy}; @@ -35,32 +35,6 @@ pub(crate) enum PutBody { Path(PathBuf), } -impl PutBody { - /// Returns the remaining byte length, or `None` when it cannot be determined. - pub(crate) async fn len(&mut self) -> Option { - match self { - Self::Buffer(bytes) => Some(bytes.len() as u64), - Self::File(file) => { - file.flush().await.ok()?; - let metadata = file.metadata().await.ok()?; - if !metadata.is_file() { - return None; - } - Some( - metadata - .len() - .saturating_sub(file.stream_position().await.ok()?), - ) - } - Self::Path(path) => { - let metadata = tokio::fs::metadata(path).await.ok()?; - metadata.is_file().then_some(metadata.len()) - } - Self::Stream(_) => None, - } - } -} - impl fmt::Debug for PutBody { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_tuple("PutBody").finish_non_exhaustive() @@ -187,8 +161,7 @@ impl PutBuilder { /// Sets the minimum remaining input size for automatic resumable uploads. /// /// Inherits [`crate::Usecase::resumable_threshold_bytes`] unless overridden. Pass [`None`] - /// to disable or zero to include every supported payload. Uploads queued through - /// [`Session::many`] ignore this setting. See [`PutBuilder`] for preparation and recovery. + /// to disable or zero to include every supported payload. See [`PutBuilder`] for eligibility. /// /// ```no_run /// # async fn example(session: objectstore_client::Session) -> objectstore_client::Result<()> { diff --git a/clients/rust/src/put/payload.rs b/clients/rust/src/put/payload.rs index d6465c3a..d997324e 100644 --- a/clients/rust/src/put/payload.rs +++ b/clients/rust/src/put/payload.rs @@ -46,7 +46,7 @@ pub(super) enum ReplayBody { } pub(super) async fn prepare( - mut body: PutBody, + body: PutBody, compression: Option, threshold: Option, spool_dir: Option, @@ -54,14 +54,24 @@ pub(super) async fn prepare( let Some(threshold) = threshold else { return Ok(PreparedBody::Direct(body)); }; - if let PutBody::Path(path) = body { - body = PutBody::File(File::open(path).await?); - } - let Some(mut length) = body.len().await else { - return Ok(PreparedBody::Direct(body)); + let body = match body { + PutBody::Path(path) => PutBody::File(File::open(path).await?), + body => body, }; - if length < threshold { - return Ok(PreparedBody::Direct(body)); + let mut replay = match body { + PutBody::Buffer(bytes) => ReplayBody::Buffer(bytes), + PutBody::File(mut file) => match regular_file_span(&mut file).await { + Some((start, length)) => ReplayBody::File { + file, + start, + length, + }, + None => return Ok(PreparedBody::Direct(PutBody::File(file))), + }, + body => return Ok(PreparedBody::Direct(body)), + }; + if replay.len() < threshold { + return Ok(PreparedBody::Direct(replay.into_put_body())); } if let Some(CompressionMode::Compress(crate::Compression::Zstd)) = compression { @@ -73,33 +83,47 @@ pub(super) async fn prepare( .map_err(io::Error::other)?; let Ok(temporary) = temporary else { // Creation has not consumed input, so direct streaming is still possible. - return Ok(PreparedBody::Direct(body)); + return Ok(PreparedBody::Direct(replay.into_put_body())); }; let mut temporary = File::from_std(temporary); - let reader: Box = match body { - PutBody::Buffer(bytes) => Box::new(io::Cursor::new(bytes)), - PutBody::File(file) => Box::new(file), - _ => unreachable!("only buffers and opened regular files have a known size"), + let reader: Box = match replay { + ReplayBody::Buffer(bytes) => Box::new(io::Cursor::new(bytes)), + ReplayBody::File { file, .. } => Box::new(file), }; let mut encoder = super::ZstdEncoder::new(super::BufReader::new(reader)); - length = tokio::io::copy(&mut encoder, &mut temporary).await?; + let length = tokio::io::copy(&mut encoder, &mut temporary).await?; temporary.flush().await?; temporary.rewind().await?; - body = PutBody::File(temporary); + replay = ReplayBody::File { + file: temporary, + start: 0, + length, + }; } - Ok(PreparedBody::Replay(match body { - PutBody::Buffer(bytes) => ReplayBody::Buffer(bytes), - PutBody::File(mut file) => ReplayBody::File { - start: file.stream_position().await?, - file, - length, - }, - _ => unreachable!("only buffers and opened regular files have a known size"), - })) + Ok(PreparedBody::Replay(replay)) +} + +/// Returns the current position and remaining length of a regular file. +async fn regular_file_span(file: &mut File) -> Option<(u64, u64)> { + file.flush().await.ok()?; + let metadata = file.metadata().await.ok()?; + if !metadata.is_file() { + return None; + } + let start = file.stream_position().await.ok()?; + Some((start, metadata.len().saturating_sub(start))) } impl ReplayBody { + /// Returns the unprepared input for a direct upload, positioned where it started. + fn into_put_body(self) -> PutBody { + match self { + Self::Buffer(bytes) => PutBody::Buffer(bytes), + Self::File { file, .. } => PutBody::File(file), + } + } + pub(super) fn len(&self) -> u64 { match self { Self::Buffer(bytes) => bytes.len() as u64, diff --git a/clients/rust/tests/automatic.rs b/clients/rust/tests/automatic.rs index 31b6e674..2ca528e6 100644 --- a/clients/rust/tests/automatic.rs +++ b/clients/rust/tests/automatic.rs @@ -1,4 +1,4 @@ -//! Automatic upload tests deliberately run without `resumable-upload-api` too. +//! Automatic upload tests do not require the `resumable-upload-api` feature. use std::collections::VecDeque; use std::io::{self, SeekFrom, Write};