Skip to main content

hotdata/
uploads.rs

1//! Ergonomic, hand-written direct-to-storage (presigned) file uploads.
2//!
3//! This module is regeneration-immune: it is protected by `.openapi-generator-ignore`
4//! and is never emitted by the OpenAPI generator. It orchestrates the
5//! presigned-upload flow that the generated [`apis::uploads_api`](crate::apis::uploads_api)
6//! ops expose as raw building blocks:
7//!
8//! 1. `POST /v1/uploads` ([`create_upload_session_handler`]) opens a session.
9//!    A small file declares its size and gets a single `url` (`mode == "single"`)
10//!    or, for a known-size multipart upload, a full set of `part_urls` plus a
11//!    `part_size` (`mode == "multipart"`). A large file omits its declared size
12//!    to open a **streaming** session: `mode == "multipart"` with a `part_size`
13//!    but NO `part_urls` — the client mints each part URL on demand from
14//!    `POST /v1/uploads/{id}/parts` ([`mint_upload_parts_handler`]) just before
15//!    uploading that part, so a URL can't expire mid-transfer on a slow upload.
16//!    Every session also carries a one-time `finalize_token`.
17//! 2. The client `PUT`s the bytes **directly to object storage** — never back
18//!    through the API. Single uploads stream the whole file to `url`; multipart
19//!    uploads slice the file into `part_size`-byte chunks and `PUT` each chunk to
20//!    its part URL (pre-issued for known-size, minted on demand for streaming),
21//!    collecting the storage `ETag` per part.
22//! 3. `POST /v1/uploads/{upload_id}/finalize` ([`finalize_upload_handler`])
23//!    confirms the upload with the finalize token in the `X-Upload-Finalize-Token`
24//!    header (empty body for single; the ascending `{part_number, e_tag}` list
25//!    for multipart) and returns a [`models::FinalizeUploadResponse`].
26//!
27//! # Storage PUT header isolation
28//!
29//! A presigned storage URL already carries its authorization in the query string
30//! (or in the server-provided `headers` map). Object stores (S3 and compatible)
31//! reject a `PUT` with `403 SignatureDoesNotMatch` if it carries extra
32//! signed-ish headers, so [`put_to_storage`] sends a *bare* request: NONE of the
33//! SDK's bearer / workspace / session headers, only an explicit `Content-Length`
34//! and whatever the server placed in `headers` (currently always empty). It also
35//! refuses to let reqwest auto-append a charset to a `Content-Type` — a type is
36//! sent only when the server's `headers` map asks for one.
37//!
38//! No S3/AWS SDK is involved: storage `PUT`s are plain `reqwest`.
39
40use std::collections::HashMap;
41use std::path::Path;
42use std::sync::atomic::{AtomicU64, Ordering};
43use std::sync::Arc;
44use std::time::Duration;
45
46use crate::apis::configuration::Configuration;
47use crate::apis::{self, Error};
48use crate::models;
49
50/// One mebibyte, the unit the storage part-size range is expressed in.
51const MIB: u64 = 1024 * 1024;
52
53/// Default cap on concurrent part `PUT`s when the caller doesn't set
54/// [`UploadOptions::max_concurrency`]. Matches the boto3 / AWS CLI default of 10.
55/// The effective in-flight count is the MIN of this and a memory budget (see
56/// [`effective_in_flight`]).
57pub const DEFAULT_MAX_CONCURRENCY: usize = 10;
58
59/// Default part-size hint, in bytes (8 MiB), sent when the caller doesn't set
60/// [`UploadOptions::part_size`]. The server clamps the hint to its own range and
61/// returns the actual size. See [`auto_part_size_hint`].
62pub const DEFAULT_PART_SIZE: u64 = 8 * MIB;
63
64/// Target ceiling on part count when auto-scaling the part-size hint for very
65/// large files, with headroom under S3's hard 10,000-part limit. See
66/// [`auto_part_size_hint`].
67pub const TARGET_MAX_PARTS: u64 = 9000;
68
69/// Minimum part size storage accepts (5 MiB). The hint is clamped to at least
70/// this; the server enforces it too.
71pub const MIN_PART_SIZE: u64 = 5 * MIB;
72
73/// Maximum part size storage accepts (5 GiB). The hint is clamped to at most
74/// this.
75pub const MAX_PART_SIZE: u64 = 5 * 1024 * MIB;
76
77/// File-size boundary between the two upload strategies. A file at or below this
78/// size takes the known-size path — a single quick `PUT` (or a short eager
79/// multipart) that completes well within a presigned URL's TTL, so there is no
80/// expiry risk. A larger file uses the streaming just-in-time path, minting each
81/// part URL only moments before it is uploaded. Set to [`DEFAULT_PART_SIZE`],
82/// the server's default single-vs-multipart boundary, so small uploads keep the
83/// single-`PUT` fast path unchanged.
84pub const STREAMING_THRESHOLD: u64 = DEFAULT_PART_SIZE;
85
86/// Target peak-memory budget for in-flight part buffers (256 MiB). Each
87/// in-flight part buffers up to `part_size` bytes, so [`effective_in_flight`]
88/// derives the in-flight count as `budget / part_size`.
89///
90/// This is a TARGET, not a hard ceiling: it holds while `part_size` is small
91/// relative to the budget (the normal case — 8 MiB parts stay well under it). It
92/// cannot bound memory below one in-flight part, so when the server returns a
93/// very large `part_size` (e.g. a 5 GiB part on a huge file), a single in-flight
94/// part already exceeds this budget and peak memory is `1 * part_size`. In other
95/// words the budget caps *concurrency*, not the size of one part.
96pub const UPLOAD_MEMORY_BUDGET: u64 = 256 * MIB;
97
98/// Compute the part-size HINT to send to the server in
99/// `CreateUploadRequest.part_size` when the caller did not specify one.
100///
101/// Starts from [`DEFAULT_PART_SIZE`] (8 MiB) and grows only for files large
102/// enough that 8 MiB parts would exceed [`TARGET_MAX_PARTS`] — so the common
103/// case is unchanged and only very large files (beyond ~72 GiB) get a larger
104/// hint to keep the part count bounded. The result is rounded UP to a whole MiB
105/// and clamped to `[MIN_PART_SIZE, MAX_PART_SIZE]`. The server still has the
106/// final say and clamps to its own range.
107///
108/// Pure and total: `declared_size == 0` yields [`DEFAULT_PART_SIZE`].
109pub fn auto_part_size_hint(declared_size: u64) -> u64 {
110    // Smallest part size that keeps the count at or under the target.
111    let by_count = declared_size.div_ceil(TARGET_MAX_PARTS);
112    let raw = DEFAULT_PART_SIZE.max(by_count);
113    // Round up to a whole MiB so the hint is a clean multiple.
114    let rounded = raw.div_ceil(MIB) * MIB;
115    rounded.clamp(MIN_PART_SIZE, MAX_PART_SIZE)
116}
117
118/// Compute how many part `PUT`s to keep in flight, given the caller's
119/// `max_concurrency` (already defaulted to [`DEFAULT_MAX_CONCURRENCY`]) and the
120/// SERVER's actual returned `part_size`.
121///
122/// Peak buffered memory is `in_flight * part_size`, so we cap in-flight at
123/// `UPLOAD_MEMORY_BUDGET / part_size`, then at `max_concurrency`. Normal 8 MiB
124/// parts give `256/8 = 32`, capped to `max_concurrency`; a 64 MiB part gives `4`.
125///
126/// `max_concurrency` is honored as an explicit floor: a caller asking for `1`
127/// (or `0`) gets serial uploads (`1`), so the budget never *raises* concurrency
128/// above what was requested. The budget-derived count itself has a floor of 1
129/// (you must keep at least one part in flight to make progress), so the overall
130/// result is always `>= 1`.
131///
132/// Pure and total: a zero `part_size` is treated as 1 to avoid division by zero.
133pub fn effective_in_flight(max_concurrency: usize, part_size: u64) -> usize {
134    // Honor an explicit low request down to serial (1); never below 1.
135    let cap = max_concurrency.max(1);
136    let by_budget = (UPLOAD_MEMORY_BUDGET / part_size.max(1)).max(1) as usize;
137    by_budget.min(cap)
138}
139
140/// Progress callback: invoked as bytes flow with `(bytes_done_total, total)`,
141/// where `total` is the full declared file size. `bytes_done_total` is
142/// monotonically non-decreasing and reaches exactly `total` when the transfer
143/// completes. Shared (`Arc`) so it can be cloned across concurrent part tasks;
144/// it must therefore be `Send + Sync`.
145pub type UploadProgress = Arc<dyn Fn(u64, u64) + Send + Sync>;
146
147/// Options for [`Client::upload_file`](crate::Client::upload_file).
148///
149/// All fields are optional. `content_type` / `content_encoding` / `filename`
150/// are recorded with the upload (advisory metadata; they do not change where the
151/// bytes are stored). `part_size` is a hint the server clamps to its allowed
152/// range and ignores for single-`PUT` uploads. `progress`, when set, is invoked
153/// as bytes flow.
154#[derive(Default, Clone)]
155pub struct UploadOptions {
156    /// Content type to record for the uploaded file (e.g. a Parquet/CSV/JSON
157    /// MIME type). Advisory.
158    pub content_type: Option<String>,
159    /// Content encoding to record for the uploaded file (e.g. `gzip`). Advisory.
160    pub content_encoding: Option<String>,
161    /// Original file name, recorded for bookkeeping. Advisory. Defaults to the
162    /// source path's file name when not set.
163    pub filename: Option<String>,
164    /// Preferred part size, in bytes, for a large (multipart) upload. A hint;
165    /// the server clamps it and ignores it for single-`PUT` uploads. When unset,
166    /// the SDK auto-scales a hint via [`auto_part_size_hint`] (8 MiB for normal
167    /// files, larger only for very large ones to bound the part count).
168    pub part_size: Option<u64>,
169    /// Maximum number of part `PUT`s to keep in flight for a multipart upload.
170    /// `None` uses [`DEFAULT_MAX_CONCURRENCY`]. The effective in-flight count is
171    /// the MIN of this and a peak-memory budget derived from the server's actual
172    /// part size (see [`effective_in_flight`]), so memory stays bounded.
173    pub max_concurrency: Option<usize>,
174    /// Optional progress callback invoked with `(bytes_done_total, total)`.
175    pub progress: Option<UploadProgress>,
176}
177
178impl std::fmt::Debug for UploadOptions {
179    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
180        f.debug_struct("UploadOptions")
181            .field("content_type", &self.content_type)
182            .field("content_encoding", &self.content_encoding)
183            .field("filename", &self.filename)
184            .field("part_size", &self.part_size)
185            .field("max_concurrency", &self.max_concurrency)
186            .field("progress", &self.progress.as_ref().map(|_| "<callback>"))
187            .finish()
188    }
189}
190
191/// Error returned by [`Client::upload_file`](crate::Client::upload_file).
192///
193/// Marked `#[non_exhaustive]`: new variants may be added without a breaking
194/// change, so downstream `match`es should carry a wildcard arm.
195#[derive(Debug)]
196#[non_exhaustive]
197pub enum UploadError {
198    /// Opening or reading the local source file failed.
199    Io(std::io::Error),
200    /// Opening the upload session (`POST /v1/uploads`) failed. A `501`
201    /// `PRESIGN_UNSUPPORTED` lands here too — the presigned path is a hard
202    /// requirement and never falls back to the legacy `POST /v1/files` proxy.
203    CreateSession(Error<apis::uploads_api::CreateUploadSessionHandlerError>),
204    /// A storage `PUT` (or the request building / transport around it) failed.
205    Storage(reqwest::Error),
206    /// A storage `PUT` returned a non-2xx status. Carries the status and the
207    /// response body for diagnosis.
208    StorageStatus {
209        /// The HTTP status the storage endpoint returned.
210        status: reqwest::StatusCode,
211        /// The 1-based part number for a multipart `PUT`, or `None` for the
212        /// single-`PUT` path.
213        part_number: Option<i32>,
214        /// The storage response body (often XML for S3-style errors).
215        body: String,
216    },
217    /// Storage accepted a part `PUT` but returned no `ETag` header, so the part
218    /// cannot be finalized.
219    MissingETag {
220        /// The 1-based part number whose `PUT` response lacked an `ETag`.
221        part_number: i32,
222    },
223    /// The create-session response was internally inconsistent for its declared
224    /// `mode` (e.g. `single` without a `url`, or `multipart` without
225    /// `part_urls` / `part_size`).
226    MalformedSession(String),
227    /// A size (the file's declared size, or the part-size hint) did not fit the
228    /// wire's signed 64-bit field. Only reachable for pathological sizes beyond
229    /// `i64::MAX` bytes (~8 EiB).
230    SizeOverflow {
231        /// What overflowed (e.g. `"declared_size_bytes"`).
232        what: &'static str,
233        /// The offending value.
234        value: u64,
235    },
236    /// Finalizing the upload (`POST /v1/uploads/{id}/finalize`) failed.
237    Finalize(Error<apis::uploads_api::FinalizeUploadHandlerError>),
238    /// Minting part URLs on demand (`POST /v1/uploads/{id}/parts`) failed during
239    /// a streaming upload — either the initial batch or an on-403 re-mint.
240    MintParts(Error<apis::uploads_api::MintUploadPartsHandlerError>),
241}
242
243impl std::fmt::Display for UploadError {
244    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
245        match self {
246            UploadError::Io(e) => write!(f, "reading the source file failed: {e}"),
247            UploadError::CreateSession(e) => write!(f, "opening the upload session failed: {e}"),
248            UploadError::Storage(e) => write!(f, "uploading to storage failed: {e}"),
249            UploadError::StorageStatus {
250                status,
251                part_number,
252                body,
253            } => match part_number {
254                Some(n) => write!(f, "storage rejected part {n} with status {status}: {body}"),
255                None => write!(
256                    f,
257                    "storage rejected the upload with status {status}: {body}"
258                ),
259            },
260            UploadError::MissingETag { part_number } => write!(
261                f,
262                "storage returned no ETag for part {part_number}; cannot finalize"
263            ),
264            UploadError::SizeOverflow { what, value } => {
265                write!(
266                    f,
267                    "{what} ({value} bytes) exceeds the maximum supported size"
268                )
269            }
270            UploadError::MalformedSession(msg) => {
271                write!(f, "malformed upload session response: {msg}")
272            }
273            UploadError::Finalize(e) => write!(f, "finalizing the upload failed: {e}"),
274            UploadError::MintParts(e) => write!(f, "minting upload part URLs failed: {e}"),
275        }
276    }
277}
278
279impl std::error::Error for UploadError {
280    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
281        match self {
282            UploadError::Io(e) => Some(e),
283            UploadError::CreateSession(e) => Some(e),
284            UploadError::Storage(e) => Some(e),
285            UploadError::Finalize(e) => Some(e),
286            UploadError::MintParts(e) => Some(e),
287            _ => None,
288        }
289    }
290}
291
292impl From<std::io::Error> for UploadError {
293    fn from(e: std::io::Error) -> Self {
294        UploadError::Io(e)
295    }
296}
297
298/// Upload a local file directly to object storage and finalize it.
299///
300/// This is the orchestration behind [`Client::upload_file`](crate::Client::upload_file);
301/// see that method for the public contract. It stats `path` for the declared
302/// size, opens a session, drives the single-`PUT` or multipart path, and
303/// finalizes — never touching the legacy `POST /v1/files` proxy.
304pub(crate) async fn upload_file(
305    configuration: &Configuration,
306    path: &Path,
307    opts: UploadOptions,
308) -> Result<models::FinalizeUploadResponse, UploadError> {
309    let metadata = tokio::fs::metadata(path).await?;
310    let total = metadata.len();
311
312    let filename = opts
313        .filename
314        .clone()
315        .or_else(|| path.file_name().map(|n| n.to_string_lossy().into_owned()));
316
317    // Part-size hint: honor an explicit caller value, else auto-scale from the
318    // declared size so the common case stays at 8 MiB and only very large files
319    // grow the hint (bounding the part count). The server clamps it regardless.
320    let part_size_hint = opts.part_size.unwrap_or_else(|| auto_part_size_hint(total));
321
322    // The wire models the part-size hint as a signed i64; reject (rather than
323    // silently wrap) a pathological hint beyond i64::MAX.
324    let part_size_hint_i64 =
325        i64::try_from(part_size_hint).map_err(|_| UploadError::SizeOverflow {
326            what: "part_size",
327            value: part_size_hint,
328        })?;
329
330    // Default a large file to a JUST-IN-TIME (streaming) session: omit
331    // `declared_size_bytes` so the server mints NO part URLs up front. The client
332    // then mints each part URL moments before it uploads that part (see
333    // `upload_multipart_streaming`), so a URL cannot expire mid-transfer no
334    // matter how long a slow upload runs — the failure mode of the eager
335    // known-size path, whose URLs share a ~30-minute TTL. A small file is a
336    // single quick `PUT` with no expiry risk, so it keeps the known-size path
337    // (and the server's single-`PUT` fast path) by declaring its size.
338    //
339    // `declared_size_bytes` is sent (and so range-checked against the wire's
340    // i64) ONLY on the known-size path; a streaming upload omits it entirely, so
341    // a size beyond i64::MAX is never an obstacle to a streamed file.
342    let declared_size_bytes = if total > STREAMING_THRESHOLD {
343        None
344    } else {
345        let size = i64::try_from(total).map_err(|_| UploadError::SizeOverflow {
346            what: "declared_size_bytes",
347            value: total,
348        })?;
349        Some(Some(size))
350    };
351    let create = models::CreateUploadRequest {
352        content_type: opts.content_type.clone().map(Some),
353        content_encoding: opts.content_encoding.clone().map(Some),
354        filename: filename.map(Some),
355        part_size: Some(Some(part_size_hint_i64)),
356        declared_size_bytes,
357        ..models::CreateUploadRequest::new()
358    };
359    let session = apis::uploads_api::create_upload_session_handler(configuration, create)
360        .await
361        .map_err(UploadError::CreateSession)?;
362
363    // Report initial progress so a 0-byte file (or an instant single PUT) still
364    // emits a terminal (0/0 or total/total) tick.
365    if let Some(ref progress) = opts.progress {
366        progress(0, total);
367    }
368
369    let parts = match session.mode.as_str() {
370        "single" => {
371            upload_single(&session, path, total, opts.progress.as_ref()).await?;
372            None
373        }
374        "multipart" => {
375            let max_concurrency = opts.max_concurrency.unwrap_or(DEFAULT_MAX_CONCURRENCY);
376            // A streaming (unknown-size) session returns NO part URLs up front
377            // (the `part_urls` key is absent or null) — mint them on demand. A
378            // known-size session returns the full `part_urls` list to PUT to
379            // directly. An explicitly present (even empty) list is a known-size
380            // response; `upload_multipart` validates it and rejects an empty one.
381            let parts = if matches!(session.part_urls, Some(Some(_))) {
382                upload_multipart(
383                    configuration,
384                    &session,
385                    path,
386                    total,
387                    max_concurrency,
388                    opts.progress.as_ref(),
389                )
390                .await?
391            } else {
392                upload_multipart_streaming(
393                    configuration,
394                    &session,
395                    path,
396                    total,
397                    max_concurrency,
398                    opts.progress.as_ref(),
399                )
400                .await?
401            };
402            Some(parts)
403        }
404        other => {
405            return Err(UploadError::MalformedSession(format!(
406                "unknown upload mode `{other}`"
407            )))
408        }
409    };
410
411    // Finalize: single sends an empty object `{}`; multipart sends
412    // `{"parts": [...]}` with the ascending, non-duplicate parts list. The token
413    // rides the X-Upload-Finalize-Token header (handled by the generated op).
414    //
415    // The body MUST be a JSON object, never `null`: the server rejects a `null`
416    // finalize body ("invalid type: null, expected struct FinalizeUploadRequest")
417    // even though the field is logically optional for single uploads. So we wrap
418    // in `Some(..)` for both modes — the generated op then serializes a struct,
419    // and `parts` (skip_serializing_if = Option::is_none) drops out for single,
420    // yielding `{}`.
421    let finalize_body = Some(
422        parts
423            .map(|parts| models::FinalizeUploadRequest {
424                parts: Some(Some(parts)),
425            })
426            .unwrap_or_default(),
427    );
428
429    // Finalize is exactly-once on the server: a second finalize of the same
430    // upload is rejected. The generated op routes through `execute_retrying`,
431    // which would retry an ambiguous failure (a lost response, or a 429 the
432    // server actually processed) — turning a finalize that SUCCEEDED into a
433    // spurious "already finalized" error on the retry. So we call it with retries
434    // disabled (a single attempt). Part PUTs stay retryable (idempotent: storage
435    // overwrites a part by number); only finalize is single-shot.
436    let mut finalize_config = configuration.clone();
437    finalize_config.retry.max_retries = 0;
438
439    apis::uploads_api::finalize_upload_handler(
440        &finalize_config,
441        &session.upload_id,
442        &session.finalize_token,
443        finalize_body,
444    )
445    .await
446    .map_err(UploadError::Finalize)
447}
448
449/// Single-`PUT` path: stream the whole file to `session.url`, invoking the
450/// progress callback incrementally as chunks are sent to storage.
451///
452/// The body is a [`progress_stream`] wrapping the file reader, so progress is
453/// byte-granular (a multi-GB upload reports smooth `done/total` ticks rather
454/// than jumping 0% -> 100%). A streaming body is not clonable, so this single
455/// `PUT` is sent once with no 429/reset retry — an intentional trade for smooth
456/// progress on the large, common single-`PUT` path; a presigned storage `PUT`
457/// is not expected to be admission-shed.
458async fn upload_single(
459    session: &models::UploadSessionResponse,
460    path: &Path,
461    total: u64,
462    progress: Option<&UploadProgress>,
463) -> Result<(), UploadError> {
464    let url =
465        session.url.clone().flatten().ok_or_else(|| {
466            UploadError::MalformedSession("single upload missing `url`".to_owned())
467        })?;
468
469    let file = tokio::fs::File::open(path).await?;
470    let body = progress_stream(file, total, progress.cloned());
471
472    put_stream_to_storage(&url, &session.headers, body, total).await?;
473
474    // Guarantee a terminal tick at exactly `total`, even if the stream's last
475    // chunk boundary or an empty file left the counter short. Monotonic: the
476    // streamed ticks never exceed `total`.
477    if let Some(progress) = progress {
478        progress(total, total);
479    }
480    Ok(())
481}
482
483/// Wrap a file reader in a byte-counting stream of `Bytes` chunks. Each chunk
484/// advances a running total and invokes `progress(done, total)` as it is yielded
485/// to the request body, so progress reflects bytes actually handed to the
486/// transport. Monotonic non-decreasing; the running total never exceeds `total`.
487fn progress_stream(
488    file: tokio::fs::File,
489    total: u64,
490    progress: Option<UploadProgress>,
491) -> ProgressStream {
492    use tokio_util::codec::{BytesCodec, FramedRead};
493
494    ProgressStream {
495        inner: FramedRead::new(file, BytesCodec::new()),
496        done: 0,
497        total,
498        progress,
499    }
500}
501
502/// A [`Stream`](futures_core::Stream) of `Bytes` chunks read from a file that
503/// reports cumulative byte progress as each chunk is yielded. Hand-rolled over
504/// `futures_core` (the crate's only direct futures dep) rather than pulling in
505/// `futures_util`, mirroring how [`Client::upload_stream`](crate::Client::upload_stream)
506/// stays on `futures_core::Stream`.
507struct ProgressStream {
508    inner: tokio_util::codec::FramedRead<tokio::fs::File, tokio_util::codec::BytesCodec>,
509    done: u64,
510    total: u64,
511    progress: Option<UploadProgress>,
512}
513
514impl futures_core::Stream for ProgressStream {
515    type Item = std::io::Result<bytes::Bytes>;
516
517    fn poll_next(
518        self: std::pin::Pin<&mut Self>,
519        cx: &mut std::task::Context<'_>,
520    ) -> std::task::Poll<Option<Self::Item>> {
521        use std::task::Poll;
522        // `inner` (FramedRead) is Unpin, and our other fields are too, so a
523        // mutable projection through `get_mut` is sound without pin-project.
524        let this = self.get_mut();
525        match std::pin::Pin::new(&mut this.inner).poll_next(cx) {
526            Poll::Ready(Some(Ok(chunk))) => {
527                let chunk = chunk.freeze();
528                this.done = (this.done + chunk.len() as u64).min(this.total);
529                if let Some(ref progress) = this.progress {
530                    progress(this.done, this.total);
531                }
532                Poll::Ready(Some(Ok(chunk)))
533            }
534            Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))),
535            Poll::Ready(None) => Poll::Ready(None),
536            Poll::Pending => Poll::Pending,
537        }
538    }
539}
540
541/// A single part's upload work, independent of how the bytes actually get to
542/// storage. The known-size path builds one of these per part and hands them to
543/// [`upload_parts_resilient`]; tests substitute a fake uploader.
544#[derive(Clone, Debug, PartialEq, Eq)]
545struct PartPlan {
546    /// 0-based position — indexes both the results vector and the URL list.
547    index: usize,
548    /// 1-based S3 part number.
549    part_number: i32,
550    /// Byte offset of this part within the source file.
551    offset: u64,
552    /// Byte length of this part (the last part is the remainder).
553    len: u64,
554}
555
556/// Outer-loop retry policy for whole-upload resilience, layered ON TOP of the
557/// per-part transport retries in [`crate::http::execute_retrying`].
558///
559/// The inner per-part retries all burn within a few seconds, so they only ride
560/// out a momentary blip *during one part*. A longer network interruption (or a
561/// flaky uplink that resets connections under load) needs a genuinely later
562/// attempt. Each extra round re-sweeps ONLY the parts still failing, after a
563/// backoff, and at reduced concurrency — fewer in-flight connections reset less
564/// on a saturated link. Crucially, parts that already succeeded keep their
565/// ETags across rounds, so a single bad part never discards the whole transfer.
566#[derive(Clone, Copy, Debug)]
567struct RetryRounds {
568    /// Re-sweeps after the initial pass. `0` reproduces the legacy behavior: a
569    /// single part exhausting its inner retries fails the entire upload.
570    max_extra_rounds: u32,
571    /// Backoff before the first re-sweep; doubles each subsequent round.
572    base_delay: Duration,
573}
574
575impl Default for RetryRounds {
576    fn default() -> Self {
577        Self {
578            max_extra_rounds: 3,
579            base_delay: Duration::from_secs(2),
580        }
581    }
582}
583
584impl RetryRounds {
585    /// Backoff before `round` (1-based for re-sweeps): capped exponential.
586    fn delay_for(self, round: u32) -> Duration {
587        let shift = round.saturating_sub(1).min(16);
588        self.base_delay.saturating_mul(1u32 << shift)
589    }
590}
591
592/// In-flight cap for a given 0-based round: halve per round, never below 1. A
593/// saturated, jittery uplink (the failure mode this whole loop targets) resets
594/// fewer connections when fewer are in flight, so later rounds back off.
595fn round_in_flight(base_in_flight: usize, round: u32) -> usize {
596    let shift = round.min(usize::BITS - 1);
597    (base_in_flight >> shift).max(1)
598}
599
600/// Whether an upload error is *terminal* — guaranteed to reproduce on a
601/// re-sweep because it reflects a server-contract or sizing violation rather
602/// than a transient network condition. Terminal errors fail the upload
603/// immediately; everything else (transport resets, timeouts, storage 4xx/5xx,
604/// mint failures) stays retryable, so a flaky link is never mistaken for a
605/// permanent fault. Kept deliberately narrow: only errors that are deterministic
606/// in the part's own inputs belong here.
607fn is_terminal(err: &UploadError) -> bool {
608    matches!(
609        err,
610        UploadError::MalformedSession(_) | UploadError::SizeOverflow { .. }
611    )
612}
613
614/// Upload every part, surviving transient per-part failures without discarding
615/// the parts that already succeeded.
616///
617/// A round runs the still-pending parts through a `JoinSet` bounded by the
618/// round's in-flight cap, recording each success and **collecting** (not
619/// propagating) each failure. If any parts remain, it waits per [`RetryRounds`]
620/// and re-sweeps just those, at reduced concurrency, until they all land or the
621/// rounds are exhausted. Completed parts' ETags persist across rounds, so the
622/// work already done is never thrown away — the bug this replaces aborted the
623/// whole upload the moment one part exhausted its inner retries.
624///
625/// `upload_part` performs one part's transfer (including its own inner transport
626/// retries) and MUST be idempotent: re-running a part overwrites it in storage
627/// (S3 `UploadPart` by number), so a re-swept part is safe.
628async fn upload_parts_resilient<F, Fut>(
629    plans: Vec<PartPlan>,
630    base_in_flight: usize,
631    rounds: RetryRounds,
632    upload_part: F,
633) -> Result<Vec<models::FinalizeUploadPart>, UploadError>
634where
635    F: Fn(PartPlan) -> Fut + Clone + Send + Sync + 'static,
636    Fut: std::future::Future<Output = Result<models::FinalizeUploadPart, UploadError>>
637        + Send
638        + 'static,
639{
640    let total_parts = plans.len();
641    // `results` is indexed by `plan.index`, so every plan's index must fall in
642    // `0..total_parts` (both callers build consecutive 0-based plans). Enforce
643    // it so a future caller passing a sparse/offset set fails loudly in tests
644    // rather than panicking or writing the wrong slot.
645    debug_assert!(
646        plans.iter().all(|p| p.index < total_parts),
647        "PartPlan.index must be within 0..plans.len()"
648    );
649    let mut results: Vec<Option<models::FinalizeUploadPart>> = vec![None; total_parts];
650    let mut remaining = plans;
651    let mut last_err: Option<UploadError> = None;
652
653    for round in 0..=rounds.max_extra_rounds {
654        if remaining.is_empty() {
655            break;
656        }
657        if round > 0 {
658            // A genuinely later attempt on a fresh window — the point of the
659            // outer loop, distinct from the inner retries that already ran.
660            tokio::time::sleep(rounds.delay_for(round)).await;
661        }
662
663        let in_flight = round_in_flight(base_in_flight, round);
664        let mut pending = std::mem::take(&mut remaining).into_iter();
665        let mut failed: Vec<PartPlan> = Vec::new();
666        let mut join_set: tokio::task::JoinSet<
667            Result<(usize, models::FinalizeUploadPart), (PartPlan, UploadError)>,
668        > = tokio::task::JoinSet::new();
669
670        loop {
671            while join_set.len() < in_flight {
672                let Some(plan) = pending.next() else { break };
673                let upload_part = upload_part.clone();
674                join_set.spawn(async move {
675                    let index = plan.index;
676                    match upload_part(plan.clone()).await {
677                        Ok(part) => Ok((index, part)),
678                        Err(e) => Err((plan, e)),
679                    }
680                });
681            }
682            match join_set.join_next().await {
683                Some(Ok(Ok((index, part)))) => results[index] = Some(part),
684                Some(Ok(Err((plan, e)))) => {
685                    // A clearly-terminal error (server-contract / sizing
686                    // violation) reproduces identically on every re-sweep, so
687                    // fail fast rather than burning the whole round budget on it.
688                    // Anything network-ish stays retryable — we never regress
689                    // resilience by mistaking a flaky link for a permanent fault.
690                    if is_terminal(&e) {
691                        join_set.abort_all();
692                        return Err(e);
693                    }
694                    // Record the failure and keep draining the rest — do NOT
695                    // abort the other in-flight parts. This part is re-swept in
696                    // the next round.
697                    failed.push(plan);
698                    last_err = Some(e);
699                }
700                Some(Err(join_err)) => {
701                    join_set.abort_all();
702                    return Err(UploadError::Io(std::io::Error::other(format!(
703                        "part upload task failed: {join_err}"
704                    ))));
705                }
706                None => break,
707            }
708        }
709        remaining = failed;
710    }
711
712    if !remaining.is_empty() {
713        // Rounds exhausted with parts still failing — surface the last
714        // underlying error so the caller's normal error mapping applies.
715        return Err(last_err
716            .unwrap_or_else(|| UploadError::Io(std::io::Error::other("multipart upload failed"))));
717    }
718
719    Ok(results.into_iter().flatten().collect())
720}
721
722/// Multipart path: slice the file into `part_size`-byte chunks (the last is the
723/// remainder), `PUT` each chunk to its `part_urls[i - 1]` with bounded
724/// concurrency, and collect `(part_number, e_tag)` per part.
725///
726/// `max_concurrency` is the caller's ceiling on in-flight parts; the effective
727/// count also honors a peak-memory budget derived from the server's actual
728/// `part_size` (see [`effective_in_flight`]).
729///
730/// Returns the parts sorted ascending by part number, ready for finalize.
731async fn upload_multipart(
732    configuration: &Configuration,
733    session: &models::UploadSessionResponse,
734    path: &Path,
735    total: u64,
736    max_concurrency: usize,
737    progress: Option<&UploadProgress>,
738) -> Result<Vec<models::FinalizeUploadPart>, UploadError> {
739    let part_urls = session.part_urls.clone().flatten().ok_or_else(|| {
740        UploadError::MalformedSession("multipart upload missing `part_urls`".to_owned())
741    })?;
742    let part_size = session.part_size.flatten().ok_or_else(|| {
743        UploadError::MalformedSession("multipart upload missing `part_size`".to_owned())
744    })?;
745    if part_size <= 0 {
746        return Err(UploadError::MalformedSession(format!(
747            "multipart upload has non-positive `part_size` {part_size}"
748        )));
749    }
750    let part_size = part_size as u64;
751
752    if part_urls.is_empty() {
753        return Err(UploadError::MalformedSession(
754            "multipart upload has empty `part_urls`".to_owned(),
755        ));
756    }
757
758    // The URL count must match the number of `part_size`-byte chunks the file
759    // splits into (last is the remainder). Too many URLs and we'd PUT a
760    // zero-length trailing part; too few and we'd finalize an incomplete list.
761    // Both mean a session inconsistent with our declared size, so fail loudly.
762    let expected_parts = total.div_ceil(part_size).max(1);
763    if part_urls.len() as u64 != expected_parts {
764        return Err(UploadError::MalformedSession(format!(
765            "multipart upload returned {} part URLs but the file ({total} bytes) \
766             splits into {expected_parts} parts of {part_size} bytes",
767            part_urls.len()
768        )));
769    }
770
771    // Peak buffered memory is in_flight * part_size; bound in-flight by both the
772    // caller's max_concurrency and the memory budget, using the SERVER's actual
773    // part size (the same value we slice by below).
774    let in_flight_cap = effective_in_flight(max_concurrency, part_size);
775
776    // Aggregate progress across parts via a shared counter; each part adds its
777    // own byte count once it lands (on success only — a re-swept part that
778    // failed an earlier round did not count, so bytes are never double-counted).
779    let done = Arc::new(AtomicU64::new(0));
780
781    // One plan per part. The last part carries the remainder; earlier parts are
782    // exactly `part_size`. A part starting at/after EOF (only possible for a
783    // zero-length file) is skipped rather than PUT as a zero-length object.
784    let mut plans: Vec<PartPlan> = Vec::with_capacity(part_urls.len());
785    for index in 0..part_urls.len() {
786        let offset = index as u64 * part_size;
787        if offset >= total && total > 0 {
788            continue;
789        }
790        let len = part_size.min(total.saturating_sub(offset));
791        plans.push(PartPlan {
792            index,
793            part_number: (index + 1) as i32,
794            offset,
795            len,
796        });
797    }
798
799    // Per-part uploader: a positioned read of exactly this part's byte range (so
800    // a re-read on retry never shares a cursor) then a header-isolated `PUT`.
801    // Captures only `Arc`s and `Copy` values, so the closure is
802    // `Clone + Send + Sync + 'static` and `upload_parts_resilient` can re-run it
803    // across rounds and concurrent tasks.
804    let part_urls = Arc::new(part_urls);
805    let headers = Arc::new(session.headers.clone());
806    let path = Arc::new(path.to_path_buf());
807    let retry = configuration.retry; // RetryPolicy is Copy.
808    let progress = progress.cloned();
809
810    let uploader = move |plan: PartPlan| {
811        let part_urls = Arc::clone(&part_urls);
812        let headers = Arc::clone(&headers);
813        let path = Arc::clone(&path);
814        let done = Arc::clone(&done);
815        let progress = progress.clone();
816        async move {
817            let url = part_urls[plan.index].clone();
818            let chunk = read_range(&path, plan.offset, plan.len).await?;
819            let resp = put_to_storage(
820                &retry,
821                &url,
822                &headers,
823                chunk,
824                plan.len,
825                Some(plan.part_number),
826            )
827            .await?;
828            let e_tag = parse_etag(resp.headers(), plan.part_number)?;
829            if let Some(progress) = progress.as_ref() {
830                let now = done.fetch_add(plan.len, Ordering::SeqCst) + plan.len;
831                progress(now, total);
832            }
833            Ok(models::FinalizeUploadPart {
834                e_tag,
835                part_number: plan.part_number,
836            })
837        }
838    };
839
840    // Resilient outer loop: a single part's transient failure no longer aborts
841    // the whole upload — it is re-swept on a later round while completed parts
842    // keep their ETags. `upload_parts_resilient` returns the parts ascending by
843    // part number with no duplicates.
844    upload_parts_resilient(plans, in_flight_cap, RetryRounds::default(), uploader).await
845}
846
847/// Streaming (just-in-time) multipart path: the session was opened WITHOUT a
848/// declared size, so the server minted no part URLs up front. We still know the
849/// local file's size, so the part count is fixed by the server's echoed
850/// `part_size`.
851///
852/// Each part mints a FRESH presigned URL immediately before its `PUT` (via
853/// `POST /v1/uploads/{id}/parts`), so a URL can never expire mid-transfer on a
854/// slow upload — and a part re-swept by [`upload_parts_resilient`] simply
855/// re-mints. This replaces the earlier batched pre-mint pipeline and its
856/// one-shot on-`403` re-mint: per-part minting is simpler and fully resilient,
857/// and with bounded concurrency the extra mint round-trip overlaps other parts'
858/// in-flight `PUT`s rather than serializing.
859///
860/// The deliberate cost is mint *request volume*: one `POST /parts` per part
861/// (up to [`TARGET_MAX_PARTS`]) instead of the old batched ≤100-per-call. We
862/// accept it because pre-minting a batch ahead is what made slow uploads fail —
863/// buffered URLs age in the queue and can expire before their part's `PUT` is
864/// reached on a constrained link. Minting each URL immediately before use keeps
865/// its age minimal, which is the whole point on the slow links this hardens.
866///
867/// Returns the parts sorted ascending by part number, ready for finalize.
868async fn upload_multipart_streaming(
869    configuration: &Configuration,
870    session: &models::UploadSessionResponse,
871    path: &Path,
872    total: u64,
873    max_concurrency: usize,
874    progress: Option<&UploadProgress>,
875) -> Result<Vec<models::FinalizeUploadPart>, UploadError> {
876    let part_size = session.part_size.flatten().ok_or_else(|| {
877        UploadError::MalformedSession("streaming upload missing `part_size`".to_owned())
878    })?;
879    if part_size <= 0 {
880        return Err(UploadError::MalformedSession(format!(
881            "streaming upload has non-positive `part_size` {part_size}"
882        )));
883    }
884    let part_size = part_size as u64;
885
886    // Slice by the SERVER's echoed part size (never our hint); the last part is
887    // the remainder. We know the file size, so the part count is fixed up front
888    // even though the server does not.
889    let expected_parts = total.div_ceil(part_size).max(1) as usize;
890
891    // Peak buffered memory is in_flight * part_size; bound in-flight by both the
892    // caller's max_concurrency and the memory budget (same as the eager path).
893    let in_flight_cap = effective_in_flight(max_concurrency, part_size);
894
895    // One plan per part (same shape as the known-size path). A part starting
896    // at/after EOF (only possible for a zero-length file) is skipped rather than
897    // PUT as a zero-length object.
898    let mut plans: Vec<PartPlan> = Vec::with_capacity(expected_parts);
899    for index in 0..expected_parts {
900        let offset = index as u64 * part_size;
901        if offset >= total && total > 0 {
902            continue;
903        }
904        let len = part_size.min(total.saturating_sub(offset));
905        plans.push(PartPlan {
906            index,
907            part_number: (index + 1) as i32,
908            offset,
909            len,
910        });
911    }
912
913    // Per-part uploader: mint a fresh URL for THIS part immediately before
914    // uploading it, then PUT. Captures only `Arc`s and `Copy` values, so the
915    // closure is `Clone + Send + Sync + 'static` and `upload_parts_resilient`
916    // can re-run it across rounds and concurrent tasks; a re-swept part re-mints
917    // a fresh URL, so expiry is impossible.
918    let config = Arc::new(configuration.clone());
919    let upload_id = Arc::new(session.upload_id.clone());
920    let finalize_token = Arc::new(session.finalize_token.clone());
921    let headers = Arc::new(session.headers.clone());
922    let path = Arc::new(path.to_path_buf());
923    let retry = configuration.retry;
924    let done = Arc::new(AtomicU64::new(0));
925    let progress = progress.cloned();
926
927    let uploader = move |plan: PartPlan| {
928        let config = Arc::clone(&config);
929        let upload_id = Arc::clone(&upload_id);
930        let finalize_token = Arc::clone(&finalize_token);
931        let headers = Arc::clone(&headers);
932        let path = Arc::clone(&path);
933        let done = Arc::clone(&done);
934        let progress = progress.clone();
935        async move {
936            let minted = apis::uploads_api::mint_upload_parts_handler(
937                &config,
938                &upload_id,
939                &finalize_token,
940                models::MintUploadPartsRequest::new(vec![plan.part_number]),
941            )
942            .await
943            .map_err(UploadError::MintParts)?;
944            let url = minted
945                .parts
946                .into_iter()
947                .find(|p| p.part_number == plan.part_number)
948                .map(|p| p.url)
949                .ok_or_else(|| {
950                    UploadError::MalformedSession(format!(
951                        "mint returned no URL for part {}",
952                        plan.part_number
953                    ))
954                })?;
955
956            let chunk = read_range(&path, plan.offset, plan.len).await?;
957            let resp = put_to_storage(
958                &retry,
959                &url,
960                &headers,
961                chunk,
962                plan.len,
963                Some(plan.part_number),
964            )
965            .await?;
966            let e_tag = parse_etag(resp.headers(), plan.part_number)?;
967            if let Some(progress) = progress.as_ref() {
968                let now = done.fetch_add(plan.len, Ordering::SeqCst) + plan.len;
969                progress(now, total);
970            }
971            Ok(models::FinalizeUploadPart {
972                e_tag,
973                part_number: plan.part_number,
974            })
975        }
976    };
977
978    // Same resilient outer loop as the known-size path: a transient part failure
979    // is re-swept on a later round (re-minting a fresh URL) instead of aborting
980    // the whole upload; completed parts keep their ETags.
981    upload_parts_resilient(plans, in_flight_cap, RetryRounds::default(), uploader).await
982}
983
984/// Extract and validate the storage `ETag` from a part `PUT` response. Rejects a
985/// missing OR empty/whitespace-only header: finalize needs a real ETag per part,
986/// and an empty value would be carried into the completion request only to fail
987/// (or silently corrupt) it later. Treated as [`UploadError::MissingETag`], so a
988/// re-sweep can re-`PUT` the part and pick up a real ETag.
989fn parse_etag(
990    headers: &reqwest::header::HeaderMap,
991    part_number: i32,
992) -> Result<String, UploadError> {
993    let etag = headers
994        .get(reqwest::header::ETAG)
995        .and_then(|v| v.to_str().ok())
996        .map(|s| s.to_owned())
997        .ok_or(UploadError::MissingETag { part_number })?;
998    if etag.trim().is_empty() {
999        return Err(UploadError::MissingETag { part_number });
1000    }
1001    Ok(etag)
1002}
1003
1004/// Read exactly `len` bytes starting at `offset` from `path`. A positioned read
1005/// (seek + read_exact) so multipart part tasks never share a cursor and a retry
1006/// re-reads the same range cleanly.
1007async fn read_range(path: &Path, offset: u64, len: u64) -> Result<bytes::Bytes, UploadError> {
1008    use tokio::io::{AsyncReadExt, AsyncSeekExt};
1009
1010    let mut file = tokio::fs::File::open(path).await?;
1011    file.seek(std::io::SeekFrom::Start(offset)).await?;
1012    let mut buf = vec![0u8; len as usize];
1013    file.read_exact(&mut buf).await?;
1014    Ok(bytes::Bytes::from(buf))
1015}
1016
1017/// Connect-phase timeout for storage `PUT`s. Bounds only TCP+TLS establishment
1018/// (not the transfer), so it is safe for both the bounded multipart parts and
1019/// the unbounded single-`PUT` whole-file path. Generous: a healthy connect is
1020/// sub-second, so 30 s only trips a genuinely dead/black-holed endpoint.
1021const STORAGE_CONNECT_TIMEOUT: Duration = Duration::from_secs(30);
1022
1023/// Fixed slack added to every per-part timeout for connect/TLS, request
1024/// queueing, and the response round-trip, independent of part size.
1025const PART_TIMEOUT_BASE: Duration = Duration::from_secs(60);
1026
1027/// Throughput floor used to size the per-part timeout. A part is only aborted if
1028/// it cannot sustain even this rate — 64 KiB/s (≈512 kbit/s), well below any
1029/// link on which an upload is worth attempting — so a legitimately slow but
1030/// progressing transfer is never killed; only a true stall is.
1031const PART_TIMEOUT_MIN_BYTES_PER_SEC: u64 = 64 * 1024;
1032
1033/// Operational ceiling on the per-part timeout. Without it a huge part (e.g. a
1034/// 5 GiB part on a multi-TB upload) would compute a ~22-hour timeout, so a
1035/// stalled giant part would hang for the better part of a day before the outer
1036/// loop could re-sweep it. 30 minutes still comfortably covers a legitimately
1037/// slow large part while keeping stall recovery bounded.
1038const PART_TIMEOUT_MAX: Duration = Duration::from_secs(30 * 60);
1039
1040/// Generous per-part total `PUT` timeout, scaled to the part size: a fixed base
1041/// plus the time the part would take at the throughput floor, capped at
1042/// [`PART_TIMEOUT_MAX`]. Examples: an 8 MiB part → ~188 s; a 64 MiB part →
1043/// ~18 min; anything above ~111 MiB → the 30 min cap. The goal is to catch a
1044/// stalled connection (which would otherwise hang the upload forever) without
1045/// aborting a healthy slow link — the outer [`upload_parts_resilient`] loop then
1046/// re-sweeps the timed-out part.
1047fn part_put_timeout(content_length: u64) -> Duration {
1048    (PART_TIMEOUT_BASE + Duration::from_secs(content_length / PART_TIMEOUT_MIN_BYTES_PER_SEC))
1049        .min(PART_TIMEOUT_MAX)
1050}
1051
1052/// `PUT` a body to a presigned storage URL with strict header isolation.
1053///
1054/// Attaches NONE of the SDK's auth/workspace/session/user-agent headers — a
1055/// presigned URL already carries its authorization, and an extra signed-ish
1056/// header makes S3-compatible storage return `403`. Only an explicit
1057/// `Content-Length` and the server-provided `headers` map (replayed verbatim;
1058/// currently always empty) are sent. A `Content-Type` is set ONLY when the
1059/// `headers` map includes one, so reqwest never auto-appends a charset.
1060///
1061/// Sent on the dedicated, header-bare [`storage_client`] with a generous,
1062/// part-size-scaled request timeout (see [`part_put_timeout`]) so a stalled
1063/// connection fails — into the outer retry loop — instead of hanging forever,
1064/// while a legitimately slow but progressing part is never aborted. The body
1065/// buffers in memory so it clones cleanly across retries via
1066/// [`crate::http::execute_retrying`]. Part `PUT`s are retryable: storage
1067/// overwrites a part by number, so a retried part is idempotent. `retry` is the
1068/// SDK's retry policy (carried on `Configuration`), used only for the retry
1069/// timing here.
1070async fn put_to_storage(
1071    retry: &crate::query::RetryPolicy,
1072    url: &str,
1073    headers: &HashMap<String, String>,
1074    body: bytes::Bytes,
1075    content_length: u64,
1076    part_number: Option<i32>,
1077) -> Result<reqwest::Response, UploadError> {
1078    let client = storage_client();
1079
1080    let mut req_builder = client
1081        .request(reqwest::Method::PUT, url)
1082        .header(reqwest::header::CONTENT_LENGTH, content_length);
1083
1084    // Replay the server-provided headers verbatim. Currently always empty; this
1085    // is the only place a Content-Type may be set, so reqwest can't auto-append
1086    // a charset.
1087    for (name, value) in headers {
1088        req_builder = req_builder.header(name.as_str(), value.as_str());
1089    }
1090
1091    // A buffered Bytes body clones cleanly, so 429 / pre-response-reset retries
1092    // in `execute_retrying` can re-send it.
1093    req_builder = req_builder.body(reqwest::Body::from(body));
1094
1095    // Per-PART total timeout, scaled to the part size. Bounds a single part so a
1096    // silently black-holed connection (no RST, write just stalls — which a
1097    // read/idle timeout would not catch) fails instead of hanging the whole
1098    // upload forever. NOT applied to the single-`PUT` whole-file path, which is
1099    // legitimately unbounded. `try_clone` in `execute_retrying` preserves this
1100    // per-request timeout, so every inner attempt gets a fresh full budget.
1101    req_builder = req_builder.timeout(part_put_timeout(content_length));
1102
1103    let req = req_builder.build().map_err(UploadError::Storage)?;
1104    crate::http_log::log_request(&req);
1105    let resp = crate::http::execute_retrying(&client, req, retry)
1106        .await
1107        .map_err(UploadError::Storage)?;
1108
1109    let status = resp.status();
1110    crate::http_log::log_response_status(status);
1111    if status.is_client_error() || status.is_server_error() {
1112        let body = resp.text().await.unwrap_or_default();
1113        crate::http_log::log_response_body(&body);
1114        return Err(UploadError::StorageStatus {
1115            status,
1116            part_number,
1117            body,
1118        });
1119    }
1120    Ok(resp)
1121}
1122
1123/// `PUT` a streaming body to a presigned storage URL with the same strict
1124/// header isolation as [`put_to_storage`] (no SDK auth/scope headers; explicit
1125/// `Content-Length`; `Content-Type` only from the server `headers` map).
1126///
1127/// Used by the single-`PUT` path so progress is byte-granular. A streamed body
1128/// is not clonable, so this is a SINGLE attempt with no 429/reset retry — unlike
1129/// the buffered, retryable [`put_to_storage`] used per multipart part.
1130async fn put_stream_to_storage<S>(
1131    url: &str,
1132    headers: &HashMap<String, String>,
1133    body: S,
1134    content_length: u64,
1135) -> Result<reqwest::Response, UploadError>
1136where
1137    S: futures_core::Stream<Item = std::io::Result<bytes::Bytes>> + Send + 'static,
1138{
1139    let client = storage_client();
1140
1141    let mut req_builder = client
1142        .request(reqwest::Method::PUT, url)
1143        // Explicit Content-Length so the body is sized (not chunked) — storage
1144        // can reject an oversized upload up front, and reqwest honors it as the
1145        // framing for a wrapped stream.
1146        .header(reqwest::header::CONTENT_LENGTH, content_length);
1147
1148    for (name, value) in headers {
1149        req_builder = req_builder.header(name.as_str(), value.as_str());
1150    }
1151
1152    req_builder = req_builder.body(reqwest::Body::wrap_stream(body));
1153
1154    let req = req_builder.build().map_err(UploadError::Storage)?;
1155    crate::http_log::log_request(&req);
1156    // A streamed body can't be cloned, so send once (no retry helper).
1157    let resp = client.execute(req).await.map_err(UploadError::Storage)?;
1158
1159    let status = resp.status();
1160    crate::http_log::log_response_status(status);
1161    if status.is_client_error() || status.is_server_error() {
1162        let body = resp.text().await.unwrap_or_default();
1163        crate::http_log::log_response_body(&body);
1164        return Err(UploadError::StorageStatus {
1165            status,
1166            part_number: None,
1167            body,
1168        });
1169    }
1170    Ok(resp)
1171}
1172
1173/// The dedicated, process-wide reqwest client used for storage `PUT`s.
1174///
1175/// Deliberately NOT `configuration.client`: a host app may have installed
1176/// default headers (auth / workspace / `User-Agent` / `Content-Type`) on the
1177/// SDK's main client, which reqwest would then apply to the storage `PUT` —
1178/// making S3-compatible storage return `403 SignatureDoesNotMatch`. This client
1179/// is built bare: no default headers, and no request timeout (a large upload
1180/// legitimately takes minutes). It is built once and reused.
1181///
1182/// Trade-off: TLS / proxy / connection-pool settings on the SDK's main client
1183/// do NOT apply to storage `PUT`s — they go through this independent client.
1184/// That is intentional; storage transfers must be header-isolated, and a
1185/// host-configured proxy for the API host is not assumed to front object
1186/// storage.
1187fn storage_client() -> reqwest::Client {
1188    static STORAGE_CLIENT: std::sync::OnceLock<reqwest::Client> = std::sync::OnceLock::new();
1189    STORAGE_CLIENT
1190        .get_or_init(|| {
1191            reqwest::Client::builder()
1192                // No `default_headers` and no client-wide request `timeout` (the
1193                // single-`PUT` whole-file path is legitimately unbounded; the
1194                // multipart path bounds each part per-request — see
1195                // `part_put_timeout`). A connect timeout is safe for both: it
1196                // bounds only connection establishment, not the transfer, so a
1197                // dead endpoint fails fast into the retry/outer loop.
1198                .connect_timeout(STORAGE_CONNECT_TIMEOUT)
1199                .build()
1200                // Falls back to a plain default client if the builder somehow
1201                // fails (e.g. no TLS backend); still header-bare.
1202                .unwrap_or_default()
1203        })
1204        .clone()
1205}
1206
1207#[cfg(test)]
1208mod tests {
1209    use super::*;
1210
1211    /// The part count a given hint would produce for a file of `size`.
1212    fn part_count(size: u64, part: u64) -> u64 {
1213        size.div_ceil(part)
1214    }
1215
1216    #[test]
1217    fn auto_part_size_keeps_8mib_for_normal_files() {
1218        // Empty and small files default to 8 MiB.
1219        assert_eq!(auto_part_size_hint(0), DEFAULT_PART_SIZE);
1220        assert_eq!(auto_part_size_hint(1), DEFAULT_PART_SIZE);
1221        assert_eq!(auto_part_size_hint(100 * MIB), DEFAULT_PART_SIZE);
1222        assert_eq!(auto_part_size_hint(1024 * MIB), DEFAULT_PART_SIZE); // 1 GiB
1223                                                                        // Right at the boundary: 8 MiB * 9000 parts = 72 GiB still fits 8 MiB.
1224        let boundary = DEFAULT_PART_SIZE * TARGET_MAX_PARTS;
1225        assert_eq!(auto_part_size_hint(boundary), DEFAULT_PART_SIZE);
1226    }
1227
1228    #[test]
1229    fn auto_part_size_scales_up_for_very_large_files_and_caps_parts() {
1230        // Beyond ~72 GiB the hint must grow above 8 MiB.
1231        let big = 200 * 1024 * MIB; // 200 GiB
1232        let hint = auto_part_size_hint(big);
1233        assert!(
1234            hint > DEFAULT_PART_SIZE,
1235            "hint should scale above 8 MiB for a 200 GiB file, got {hint}"
1236        );
1237        // Hint is a whole number of MiB.
1238        assert_eq!(hint % MIB, 0, "hint must be a whole MiB, got {hint}");
1239        // Part count stays at or under the target ceiling.
1240        assert!(
1241            part_count(big, hint) <= TARGET_MAX_PARTS,
1242            "part count {} must be <= {TARGET_MAX_PARTS}",
1243            part_count(big, hint)
1244        );
1245        // And always within storage's accepted range.
1246        assert!((MIN_PART_SIZE..=MAX_PART_SIZE).contains(&hint));
1247    }
1248
1249    #[test]
1250    fn auto_part_size_clamps_to_max_for_enormous_files() {
1251        // A file so large the count-driven size would exceed 5 GiB clamps to the
1252        // 5 GiB ceiling (the part count then necessarily exceeds the soft target,
1253        // which is fine — it's a hint and the server has the final say).
1254        let enormous = 100 * 1024 * 1024 * MIB; // 100 PiB
1255        assert_eq!(auto_part_size_hint(enormous), MAX_PART_SIZE);
1256    }
1257
1258    #[test]
1259    fn effective_in_flight_capped_by_max_concurrency_for_small_parts() {
1260        // 8 MiB parts: budget allows 256/8 = 32, so max_concurrency wins.
1261        assert_eq!(effective_in_flight(12, 8 * MIB), 12);
1262        assert_eq!(effective_in_flight(10, 8 * MIB), 10);
1263        // A tiny part size still can't exceed max_concurrency.
1264        assert_eq!(effective_in_flight(12, MIB), 12);
1265    }
1266
1267    #[test]
1268    fn effective_in_flight_reduced_by_memory_budget_for_large_parts() {
1269        // 64 MiB parts: budget allows 256/64 = 4, below max_concurrency.
1270        assert_eq!(effective_in_flight(12, 64 * MIB), 4);
1271        // 128 MiB parts: 256/128 = 2.
1272        assert_eq!(effective_in_flight(12, 128 * MIB), 2);
1273    }
1274
1275    #[test]
1276    fn effective_in_flight_honors_explicit_low_concurrency() {
1277        // An explicit max_concurrency of 1 means serial uploads — NOT raised to a
1278        // floor of 2. (Regression guard for the Codex finding.)
1279        assert_eq!(effective_in_flight(1, 8 * MIB), 1);
1280        // 0 is normalized to 1 (you can't run zero in flight), not to 2.
1281        assert_eq!(effective_in_flight(0, 8 * MIB), 1);
1282        // 2 stays 2.
1283        assert_eq!(effective_in_flight(2, 8 * MIB), 2);
1284    }
1285
1286    #[test]
1287    fn effective_in_flight_floors_at_1_for_huge_parts_and_handles_zero() {
1288        // A part larger than the whole budget still keeps at least 1 in flight
1289        // (the budget can't bound below a single part).
1290        assert_eq!(effective_in_flight(12, UPLOAD_MEMORY_BUDGET * 4), 1);
1291        // Zero part size doesn't divide-by-zero (treated as 1 byte): the budget
1292        // then allows a huge count, so max_concurrency wins.
1293        assert_eq!(effective_in_flight(12, 0), 12);
1294    }
1295}
1296
1297#[cfg(test)]
1298mod resilient_retry_tests {
1299    use super::*;
1300    use std::collections::HashMap;
1301    use std::sync::atomic::AtomicUsize;
1302    use std::sync::Mutex;
1303
1304    fn plan(n: i32) -> PartPlan {
1305        PartPlan {
1306            index: (n - 1) as usize,
1307            part_number: n,
1308            offset: (n as u64 - 1) * 16,
1309            len: 16,
1310        }
1311    }
1312    fn plans(count: i32) -> Vec<PartPlan> {
1313        (1..=count).map(plan).collect()
1314    }
1315    fn no_delay(max_extra_rounds: u32) -> RetryRounds {
1316        RetryRounds {
1317            max_extra_rounds,
1318            base_delay: Duration::ZERO,
1319        }
1320    }
1321
1322    /// A transport-free stand-in for the real per-part uploader. It records
1323    /// attempts per part and can be told to fail the first K attempts of
1324    /// specific parts (modelling a part whose inner transport retries were
1325    /// exhausted by a network blip) before succeeding. Also tracks peak
1326    /// in-flight concurrency to verify the cap is honored.
1327    #[derive(Clone)]
1328    struct FakeUploader {
1329        fail: Arc<Mutex<HashMap<i32, usize>>>, // part_number -> remaining forced failures
1330        attempts: Arc<Mutex<HashMap<i32, usize>>>,
1331        in_flight: Arc<AtomicUsize>,
1332        peak_in_flight: Arc<AtomicUsize>,
1333    }
1334
1335    impl FakeUploader {
1336        fn new(fail: HashMap<i32, usize>) -> Self {
1337            Self {
1338                fail: Arc::new(Mutex::new(fail)),
1339                attempts: Arc::new(Mutex::new(HashMap::new())),
1340                in_flight: Arc::new(AtomicUsize::new(0)),
1341                peak_in_flight: Arc::new(AtomicUsize::new(0)),
1342            }
1343        }
1344        fn attempts_for(&self, n: i32) -> usize {
1345            *self.attempts.lock().unwrap().get(&n).unwrap_or(&0)
1346        }
1347        fn peak(&self) -> usize {
1348            self.peak_in_flight.load(Ordering::SeqCst)
1349        }
1350
1351        fn call(
1352            &self,
1353            plan: PartPlan,
1354        ) -> impl std::future::Future<Output = Result<models::FinalizeUploadPart, UploadError>>
1355               + Send
1356               + 'static {
1357            let fail = Arc::clone(&self.fail);
1358            let attempts = Arc::clone(&self.attempts);
1359            let in_flight = Arc::clone(&self.in_flight);
1360            let peak = Arc::clone(&self.peak_in_flight);
1361            async move {
1362                let cur = in_flight.fetch_add(1, Ordering::SeqCst) + 1;
1363                peak.fetch_max(cur, Ordering::SeqCst);
1364                // Force overlap so peak-in-flight reflects real concurrency.
1365                tokio::task::yield_now().await;
1366                *attempts
1367                    .lock()
1368                    .unwrap()
1369                    .entry(plan.part_number)
1370                    .or_insert(0) += 1;
1371                let should_fail = {
1372                    let mut f = fail.lock().unwrap();
1373                    match f.get_mut(&plan.part_number) {
1374                        Some(remaining) if *remaining > 0 => {
1375                            *remaining -= 1;
1376                            true
1377                        }
1378                        _ => false,
1379                    }
1380                };
1381                in_flight.fetch_sub(1, Ordering::SeqCst);
1382                if should_fail {
1383                    Err(UploadError::Io(std::io::Error::other(
1384                        "simulated connection reset",
1385                    )))
1386                } else {
1387                    Ok(models::FinalizeUploadPart {
1388                        e_tag: format!("etag-{}", plan.part_number),
1389                        part_number: plan.part_number,
1390                    })
1391                }
1392            }
1393        }
1394    }
1395
1396    // ---- pure policy ----
1397
1398    #[test]
1399    fn round_in_flight_halves_each_round_min_one() {
1400        assert_eq!(round_in_flight(8, 0), 8);
1401        assert_eq!(round_in_flight(8, 1), 4);
1402        assert_eq!(round_in_flight(8, 2), 2);
1403        assert_eq!(round_in_flight(8, 3), 1);
1404        assert_eq!(round_in_flight(8, 99), 1);
1405        assert_eq!(round_in_flight(1, 3), 1);
1406    }
1407
1408    #[test]
1409    fn delay_for_grows_exponentially() {
1410        let r = RetryRounds {
1411            max_extra_rounds: 3,
1412            base_delay: Duration::from_secs(2),
1413        };
1414        assert_eq!(r.delay_for(1), Duration::from_secs(2));
1415        assert_eq!(r.delay_for(2), Duration::from_secs(4));
1416        assert_eq!(r.delay_for(3), Duration::from_secs(8));
1417    }
1418
1419    #[test]
1420    fn part_put_timeout_is_generous_and_scales_with_part_size() {
1421        // 8 MiB part: 60s base + 8MiB / 64KiB/s = 60 + 128 = 188s. Comfortably
1422        // above the ~3s an 8 MiB part takes on a healthy link, so a legit slow
1423        // transfer is never aborted; only a true stall trips it.
1424        assert_eq!(part_put_timeout(8 * 1024 * 1024), Duration::from_secs(188));
1425        // 64 MiB part stays generous (~18 min).
1426        assert_eq!(
1427            part_put_timeout(64 * 1024 * 1024),
1428            Duration::from_secs(60 + 1024)
1429        );
1430        // A tiny/empty part still gets the full fixed base.
1431        assert_eq!(part_put_timeout(0), Duration::from_secs(60));
1432        // Monotonic in part size (below the cap).
1433        assert!(part_put_timeout(32 * 1024 * 1024) > part_put_timeout(8 * 1024 * 1024));
1434        // A huge part is capped at the 30 min operational ceiling rather than the
1435        // ~22.8 h the raw formula would yield, so stall recovery stays bounded.
1436        assert_eq!(part_put_timeout(5 * 1024 * 1024 * 1024), PART_TIMEOUT_MAX);
1437        assert_eq!(part_put_timeout(u64::MAX), PART_TIMEOUT_MAX);
1438    }
1439
1440    #[test]
1441    fn terminal_errors_are_only_contract_violations() {
1442        assert!(is_terminal(&UploadError::MalformedSession("bad".into())));
1443        assert!(is_terminal(&UploadError::SizeOverflow {
1444            what: "x",
1445            value: 1,
1446        }));
1447        // Network-ish failures must stay retryable so the outer loop re-sweeps.
1448        assert!(!is_terminal(&UploadError::Io(std::io::Error::other(
1449            "reset"
1450        ))));
1451        assert!(!is_terminal(&UploadError::MissingETag { part_number: 1 }));
1452        assert!(!is_terminal(&UploadError::StorageStatus {
1453            status: reqwest::StatusCode::INTERNAL_SERVER_ERROR,
1454            part_number: Some(1),
1455            body: String::new(),
1456        }));
1457    }
1458
1459    #[test]
1460    fn parse_etag_rejects_missing_and_blank() {
1461        use reqwest::header::{HeaderMap, HeaderValue, ETAG};
1462        let mut ok = HeaderMap::new();
1463        ok.insert(ETAG, HeaderValue::from_static("\"etag-7\""));
1464        assert_eq!(parse_etag(&ok, 7).unwrap(), "\"etag-7\"");
1465
1466        // Missing header.
1467        assert!(matches!(
1468            parse_etag(&HeaderMap::new(), 7),
1469            Err(UploadError::MissingETag { part_number: 7 })
1470        ));
1471        // Present but empty / whitespace-only — must be rejected, not finalized.
1472        for blank in ["", "   "] {
1473            let mut h = HeaderMap::new();
1474            h.insert(ETAG, HeaderValue::from_str(blank).unwrap());
1475            assert!(
1476                matches!(parse_etag(&h, 7), Err(UploadError::MissingETag { .. })),
1477                "blank ETag {blank:?} must be rejected"
1478            );
1479        }
1480    }
1481
1482    // ---- REPRODUCE: legacy behavior = one pass, no outer rounds ----
1483
1484    #[tokio::test]
1485    async fn repro_single_part_blip_sinks_whole_upload_without_rounds() {
1486        // Part 3 fails once. With NO extra rounds (the legacy abort-on-first-
1487        // exhaustion behavior) that single transient failure fails the entire
1488        // upload, discarding the work done on parts 1, 2, 4, 5.
1489        let fake = FakeUploader::new(HashMap::from([(3, 1)]));
1490        let f = fake.clone();
1491        let res = upload_parts_resilient(plans(5), 4, no_delay(0), move |p| f.call(p)).await;
1492        assert!(
1493            res.is_err(),
1494            "a single transient part failure should sink the upload under legacy (0-round) semantics"
1495        );
1496    }
1497
1498    // ---- FIX: outer rounds re-sweep only the failed parts ----
1499
1500    #[tokio::test]
1501    async fn fix_single_part_blip_recovers_on_a_later_round() {
1502        let fake = FakeUploader::new(HashMap::from([(3, 1)]));
1503        let f = fake.clone();
1504        let res = upload_parts_resilient(plans(5), 4, no_delay(3), move |p| f.call(p))
1505            .await
1506            .expect("the flaky part should recover on a later round");
1507        // All five parts present, ascending, with the right ETags.
1508        let nums: Vec<i32> = res.iter().map(|p| p.part_number).collect();
1509        assert_eq!(nums, vec![1, 2, 3, 4, 5]);
1510        assert_eq!(res[2].e_tag, "etag-3");
1511        // The flaky part was attempted twice (round 0 fail, round 1 success);
1512        // every healthy part exactly once — completed work is never redone.
1513        assert_eq!(fake.attempts_for(3), 2);
1514        for n in [1, 2, 4, 5] {
1515            assert_eq!(fake.attempts_for(n), 1, "part {n} must not be re-uploaded");
1516        }
1517    }
1518
1519    #[tokio::test]
1520    async fn fix_multiple_flaky_parts_all_recover() {
1521        let fake = FakeUploader::new(HashMap::from([(2, 2), (5, 1), (7, 3)]));
1522        let f = fake.clone();
1523        let res = upload_parts_resilient(plans(8), 4, no_delay(3), move |p| f.call(p))
1524            .await
1525            .expect("all parts should recover within the round budget");
1526        assert_eq!(res.len(), 8);
1527        assert_eq!(fake.attempts_for(2), 3); // 2 fails + success
1528        assert_eq!(fake.attempts_for(7), 4); // 3 fails + success
1529        assert_eq!(fake.attempts_for(5), 2);
1530    }
1531
1532    #[tokio::test]
1533    async fn permanent_failure_surfaced_after_exhausting_rounds() {
1534        // Part 4 always fails (more failures than rounds). After the initial
1535        // pass plus `max_extra_rounds` re-sweeps the upload gives up — but only
1536        // after exactly 1 + max_extra_rounds attempts of that part, and without
1537        // ever re-uploading the healthy parts.
1538        let fake = FakeUploader::new(HashMap::from([(4, 99)]));
1539        let f = fake.clone();
1540        let res = upload_parts_resilient(plans(5), 4, no_delay(2), move |p| f.call(p)).await;
1541        assert!(res.is_err());
1542        assert_eq!(fake.attempts_for(4), 3, "1 initial pass + 2 re-sweeps");
1543        for n in [1, 2, 3, 5] {
1544            assert_eq!(fake.attempts_for(n), 1);
1545        }
1546    }
1547
1548    #[tokio::test]
1549    async fn happy_path_uploads_each_part_exactly_once() {
1550        let fake = FakeUploader::new(HashMap::new());
1551        let f = fake.clone();
1552        let res = upload_parts_resilient(plans(6), 4, no_delay(3), move |p| f.call(p))
1553            .await
1554            .unwrap();
1555        assert_eq!(res.len(), 6);
1556        for n in 1..=6 {
1557            assert_eq!(fake.attempts_for(n), 1);
1558        }
1559    }
1560
1561    #[tokio::test]
1562    async fn concurrency_never_exceeds_base_cap() {
1563        let fake = FakeUploader::new(HashMap::new());
1564        let f = fake.clone();
1565        upload_parts_resilient(plans(20), 3, no_delay(3), move |p| f.call(p))
1566            .await
1567            .unwrap();
1568        assert!(
1569            fake.peak() <= 3,
1570            "peak in-flight {} exceeded the cap of 3",
1571            fake.peak()
1572        );
1573    }
1574
1575    #[tokio::test]
1576    async fn terminal_error_fails_fast_without_resweeping() {
1577        // A terminal error (server-contract violation) reproduces on every
1578        // re-sweep, so it must fail the upload immediately — NOT be retried for
1579        // all rounds the way a transient failure is.
1580        let p2_attempts = Arc::new(AtomicUsize::new(0));
1581        let counter = Arc::clone(&p2_attempts);
1582        let res = upload_parts_resilient(plans(4), 4, no_delay(3), move |plan: PartPlan| {
1583            let counter = Arc::clone(&counter);
1584            async move {
1585                if plan.part_number == 2 {
1586                    counter.fetch_add(1, Ordering::SeqCst);
1587                    Err(UploadError::MalformedSession("contract violation".into()))
1588                } else {
1589                    Ok(models::FinalizeUploadPart {
1590                        e_tag: format!("etag-{}", plan.part_number),
1591                        part_number: plan.part_number,
1592                    })
1593                }
1594            }
1595        })
1596        .await;
1597        assert!(matches!(res, Err(UploadError::MalformedSession(_))));
1598        assert_eq!(
1599            p2_attempts.load(Ordering::SeqCst),
1600            1,
1601            "a terminal error must be attempted once, never re-swept across rounds"
1602        );
1603    }
1604}