use std::collections::HashMap;
use std::path::Path;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;
use crate::apis::configuration::Configuration;
use crate::apis::{self, Error};
use crate::models;
const MIB: u64 = 1024 * 1024;
pub const DEFAULT_MAX_CONCURRENCY: usize = 10;
pub const DEFAULT_PART_SIZE: u64 = 8 * MIB;
pub const TARGET_MAX_PARTS: u64 = 9000;
pub const MIN_PART_SIZE: u64 = 5 * MIB;
pub const MAX_PART_SIZE: u64 = 5 * 1024 * MIB;
pub const STREAMING_THRESHOLD: u64 = DEFAULT_PART_SIZE;
pub const UPLOAD_MEMORY_BUDGET: u64 = 256 * MIB;
pub fn auto_part_size_hint(declared_size: u64) -> u64 {
let by_count = declared_size.div_ceil(TARGET_MAX_PARTS);
let raw = DEFAULT_PART_SIZE.max(by_count);
let rounded = raw.div_ceil(MIB) * MIB;
rounded.clamp(MIN_PART_SIZE, MAX_PART_SIZE)
}
pub fn effective_in_flight(max_concurrency: usize, part_size: u64) -> usize {
let cap = max_concurrency.max(1);
let by_budget = (UPLOAD_MEMORY_BUDGET / part_size.max(1)).max(1) as usize;
by_budget.min(cap)
}
pub type UploadProgress = Arc<dyn Fn(u64, u64) + Send + Sync>;
#[derive(Default, Clone)]
pub struct UploadOptions {
pub content_type: Option<String>,
pub content_encoding: Option<String>,
pub filename: Option<String>,
pub part_size: Option<u64>,
pub max_concurrency: Option<usize>,
pub progress: Option<UploadProgress>,
}
impl std::fmt::Debug for UploadOptions {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("UploadOptions")
.field("content_type", &self.content_type)
.field("content_encoding", &self.content_encoding)
.field("filename", &self.filename)
.field("part_size", &self.part_size)
.field("max_concurrency", &self.max_concurrency)
.field("progress", &self.progress.as_ref().map(|_| "<callback>"))
.finish()
}
}
#[derive(Debug)]
#[non_exhaustive]
pub enum UploadError {
Io(std::io::Error),
CreateSession(Error<apis::uploads_api::CreateUploadSessionHandlerError>),
Storage(reqwest::Error),
StorageStatus {
status: reqwest::StatusCode,
part_number: Option<i32>,
body: String,
},
MissingETag {
part_number: i32,
},
MalformedSession(String),
SizeOverflow {
what: &'static str,
value: u64,
},
Finalize(Error<apis::uploads_api::FinalizeUploadHandlerError>),
MintParts(Error<apis::uploads_api::MintUploadPartsHandlerError>),
}
impl std::fmt::Display for UploadError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
UploadError::Io(e) => write!(f, "reading the source file failed: {e}"),
UploadError::CreateSession(e) => write!(f, "opening the upload session failed: {e}"),
UploadError::Storage(e) => write!(f, "uploading to storage failed: {e}"),
UploadError::StorageStatus {
status,
part_number,
body,
} => match part_number {
Some(n) => write!(f, "storage rejected part {n} with status {status}: {body}"),
None => write!(
f,
"storage rejected the upload with status {status}: {body}"
),
},
UploadError::MissingETag { part_number } => write!(
f,
"storage returned no ETag for part {part_number}; cannot finalize"
),
UploadError::SizeOverflow { what, value } => {
write!(
f,
"{what} ({value} bytes) exceeds the maximum supported size"
)
}
UploadError::MalformedSession(msg) => {
write!(f, "malformed upload session response: {msg}")
}
UploadError::Finalize(e) => write!(f, "finalizing the upload failed: {e}"),
UploadError::MintParts(e) => write!(f, "minting upload part URLs failed: {e}"),
}
}
}
impl std::error::Error for UploadError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
UploadError::Io(e) => Some(e),
UploadError::CreateSession(e) => Some(e),
UploadError::Storage(e) => Some(e),
UploadError::Finalize(e) => Some(e),
UploadError::MintParts(e) => Some(e),
_ => None,
}
}
}
impl From<std::io::Error> for UploadError {
fn from(e: std::io::Error) -> Self {
UploadError::Io(e)
}
}
pub(crate) async fn upload_file(
configuration: &Configuration,
path: &Path,
opts: UploadOptions,
) -> Result<models::FinalizeUploadResponse, UploadError> {
let metadata = tokio::fs::metadata(path).await?;
let total = metadata.len();
let filename = opts
.filename
.clone()
.or_else(|| path.file_name().map(|n| n.to_string_lossy().into_owned()));
let part_size_hint = opts.part_size.unwrap_or_else(|| auto_part_size_hint(total));
let part_size_hint_i64 =
i64::try_from(part_size_hint).map_err(|_| UploadError::SizeOverflow {
what: "part_size",
value: part_size_hint,
})?;
let declared_size_bytes = if total > STREAMING_THRESHOLD {
None
} else {
let size = i64::try_from(total).map_err(|_| UploadError::SizeOverflow {
what: "declared_size_bytes",
value: total,
})?;
Some(Some(size))
};
let create = models::CreateUploadRequest {
content_type: opts.content_type.clone().map(Some),
content_encoding: opts.content_encoding.clone().map(Some),
filename: filename.map(Some),
part_size: Some(Some(part_size_hint_i64)),
declared_size_bytes,
..models::CreateUploadRequest::new()
};
let session = apis::uploads_api::create_upload_session_handler(configuration, create)
.await
.map_err(UploadError::CreateSession)?;
if let Some(ref progress) = opts.progress {
progress(0, total);
}
let parts = match session.mode.as_str() {
"single" => {
upload_single(&session, path, total, opts.progress.as_ref()).await?;
None
}
"multipart" => {
let max_concurrency = opts.max_concurrency.unwrap_or(DEFAULT_MAX_CONCURRENCY);
let parts = if matches!(session.part_urls, Some(Some(_))) {
upload_multipart(
configuration,
&session,
path,
total,
max_concurrency,
opts.progress.as_ref(),
)
.await?
} else {
upload_multipart_streaming(
configuration,
&session,
path,
total,
max_concurrency,
opts.progress.as_ref(),
)
.await?
};
Some(parts)
}
other => {
return Err(UploadError::MalformedSession(format!(
"unknown upload mode `{other}`"
)))
}
};
let finalize_body = Some(
parts
.map(|parts| models::FinalizeUploadRequest {
parts: Some(Some(parts)),
})
.unwrap_or_default(),
);
let mut finalize_config = configuration.clone();
finalize_config.retry.max_retries = 0;
apis::uploads_api::finalize_upload_handler(
&finalize_config,
&session.upload_id,
&session.finalize_token,
finalize_body,
)
.await
.map_err(UploadError::Finalize)
}
async fn upload_single(
session: &models::UploadSessionResponse,
path: &Path,
total: u64,
progress: Option<&UploadProgress>,
) -> Result<(), UploadError> {
let url =
session.url.clone().flatten().ok_or_else(|| {
UploadError::MalformedSession("single upload missing `url`".to_owned())
})?;
let file = tokio::fs::File::open(path).await?;
let body = progress_stream(file, total, progress.cloned());
put_stream_to_storage(&url, &session.headers, body, total).await?;
if let Some(progress) = progress {
progress(total, total);
}
Ok(())
}
fn progress_stream(
file: tokio::fs::File,
total: u64,
progress: Option<UploadProgress>,
) -> ProgressStream {
use tokio_util::codec::{BytesCodec, FramedRead};
ProgressStream {
inner: FramedRead::new(file, BytesCodec::new()),
done: 0,
total,
progress,
}
}
struct ProgressStream {
inner: tokio_util::codec::FramedRead<tokio::fs::File, tokio_util::codec::BytesCodec>,
done: u64,
total: u64,
progress: Option<UploadProgress>,
}
impl futures_core::Stream for ProgressStream {
type Item = std::io::Result<bytes::Bytes>;
fn poll_next(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Option<Self::Item>> {
use std::task::Poll;
let this = self.get_mut();
match std::pin::Pin::new(&mut this.inner).poll_next(cx) {
Poll::Ready(Some(Ok(chunk))) => {
let chunk = chunk.freeze();
this.done = (this.done + chunk.len() as u64).min(this.total);
if let Some(ref progress) = this.progress {
progress(this.done, this.total);
}
Poll::Ready(Some(Ok(chunk)))
}
Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))),
Poll::Ready(None) => Poll::Ready(None),
Poll::Pending => Poll::Pending,
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
struct PartPlan {
index: usize,
part_number: i32,
offset: u64,
len: u64,
}
#[derive(Clone, Copy, Debug)]
struct RetryRounds {
max_extra_rounds: u32,
base_delay: Duration,
}
impl Default for RetryRounds {
fn default() -> Self {
Self {
max_extra_rounds: 3,
base_delay: Duration::from_secs(2),
}
}
}
impl RetryRounds {
fn delay_for(self, round: u32) -> Duration {
let shift = round.saturating_sub(1).min(16);
self.base_delay.saturating_mul(1u32 << shift)
}
}
fn round_in_flight(base_in_flight: usize, round: u32) -> usize {
let shift = round.min(usize::BITS - 1);
(base_in_flight >> shift).max(1)
}
fn is_terminal(err: &UploadError) -> bool {
matches!(
err,
UploadError::MalformedSession(_) | UploadError::SizeOverflow { .. }
)
}
async fn upload_parts_resilient<F, Fut>(
plans: Vec<PartPlan>,
base_in_flight: usize,
rounds: RetryRounds,
upload_part: F,
) -> Result<Vec<models::FinalizeUploadPart>, UploadError>
where
F: Fn(PartPlan) -> Fut + Clone + Send + Sync + 'static,
Fut: std::future::Future<Output = Result<models::FinalizeUploadPart, UploadError>>
+ Send
+ 'static,
{
let total_parts = plans.len();
debug_assert!(
plans.iter().all(|p| p.index < total_parts),
"PartPlan.index must be within 0..plans.len()"
);
let mut results: Vec<Option<models::FinalizeUploadPart>> = vec![None; total_parts];
let mut remaining = plans;
let mut last_err: Option<UploadError> = None;
for round in 0..=rounds.max_extra_rounds {
if remaining.is_empty() {
break;
}
if round > 0 {
tokio::time::sleep(rounds.delay_for(round)).await;
}
let in_flight = round_in_flight(base_in_flight, round);
let mut pending = std::mem::take(&mut remaining).into_iter();
let mut failed: Vec<PartPlan> = Vec::new();
let mut join_set: tokio::task::JoinSet<
Result<(usize, models::FinalizeUploadPart), (PartPlan, UploadError)>,
> = tokio::task::JoinSet::new();
loop {
while join_set.len() < in_flight {
let Some(plan) = pending.next() else { break };
let upload_part = upload_part.clone();
join_set.spawn(async move {
let index = plan.index;
match upload_part(plan.clone()).await {
Ok(part) => Ok((index, part)),
Err(e) => Err((plan, e)),
}
});
}
match join_set.join_next().await {
Some(Ok(Ok((index, part)))) => results[index] = Some(part),
Some(Ok(Err((plan, e)))) => {
if is_terminal(&e) {
join_set.abort_all();
return Err(e);
}
failed.push(plan);
last_err = Some(e);
}
Some(Err(join_err)) => {
join_set.abort_all();
return Err(UploadError::Io(std::io::Error::other(format!(
"part upload task failed: {join_err}"
))));
}
None => break,
}
}
remaining = failed;
}
if !remaining.is_empty() {
return Err(last_err
.unwrap_or_else(|| UploadError::Io(std::io::Error::other("multipart upload failed"))));
}
Ok(results.into_iter().flatten().collect())
}
async fn upload_multipart(
configuration: &Configuration,
session: &models::UploadSessionResponse,
path: &Path,
total: u64,
max_concurrency: usize,
progress: Option<&UploadProgress>,
) -> Result<Vec<models::FinalizeUploadPart>, UploadError> {
let part_urls = session.part_urls.clone().flatten().ok_or_else(|| {
UploadError::MalformedSession("multipart upload missing `part_urls`".to_owned())
})?;
let part_size = session.part_size.flatten().ok_or_else(|| {
UploadError::MalformedSession("multipart upload missing `part_size`".to_owned())
})?;
if part_size <= 0 {
return Err(UploadError::MalformedSession(format!(
"multipart upload has non-positive `part_size` {part_size}"
)));
}
let part_size = part_size as u64;
if part_urls.is_empty() {
return Err(UploadError::MalformedSession(
"multipart upload has empty `part_urls`".to_owned(),
));
}
let expected_parts = total.div_ceil(part_size).max(1);
if part_urls.len() as u64 != expected_parts {
return Err(UploadError::MalformedSession(format!(
"multipart upload returned {} part URLs but the file ({total} bytes) \
splits into {expected_parts} parts of {part_size} bytes",
part_urls.len()
)));
}
let in_flight_cap = effective_in_flight(max_concurrency, part_size);
let done = Arc::new(AtomicU64::new(0));
let mut plans: Vec<PartPlan> = Vec::with_capacity(part_urls.len());
for index in 0..part_urls.len() {
let offset = index as u64 * part_size;
if offset >= total && total > 0 {
continue;
}
let len = part_size.min(total.saturating_sub(offset));
plans.push(PartPlan {
index,
part_number: (index + 1) as i32,
offset,
len,
});
}
let part_urls = Arc::new(part_urls);
let headers = Arc::new(session.headers.clone());
let path = Arc::new(path.to_path_buf());
let retry = configuration.retry; let progress = progress.cloned();
let uploader = move |plan: PartPlan| {
let part_urls = Arc::clone(&part_urls);
let headers = Arc::clone(&headers);
let path = Arc::clone(&path);
let done = Arc::clone(&done);
let progress = progress.clone();
async move {
let url = part_urls[plan.index].clone();
let chunk = read_range(&path, plan.offset, plan.len).await?;
let resp = put_to_storage(
&retry,
&url,
&headers,
chunk,
plan.len,
Some(plan.part_number),
)
.await?;
let e_tag = parse_etag(resp.headers(), plan.part_number)?;
if let Some(progress) = progress.as_ref() {
let now = done.fetch_add(plan.len, Ordering::SeqCst) + plan.len;
progress(now, total);
}
Ok(models::FinalizeUploadPart {
e_tag,
part_number: plan.part_number,
})
}
};
upload_parts_resilient(plans, in_flight_cap, RetryRounds::default(), uploader).await
}
async fn upload_multipart_streaming(
configuration: &Configuration,
session: &models::UploadSessionResponse,
path: &Path,
total: u64,
max_concurrency: usize,
progress: Option<&UploadProgress>,
) -> Result<Vec<models::FinalizeUploadPart>, UploadError> {
let part_size = session.part_size.flatten().ok_or_else(|| {
UploadError::MalformedSession("streaming upload missing `part_size`".to_owned())
})?;
if part_size <= 0 {
return Err(UploadError::MalformedSession(format!(
"streaming upload has non-positive `part_size` {part_size}"
)));
}
let part_size = part_size as u64;
let expected_parts = total.div_ceil(part_size).max(1) as usize;
let in_flight_cap = effective_in_flight(max_concurrency, part_size);
let mut plans: Vec<PartPlan> = Vec::with_capacity(expected_parts);
for index in 0..expected_parts {
let offset = index as u64 * part_size;
if offset >= total && total > 0 {
continue;
}
let len = part_size.min(total.saturating_sub(offset));
plans.push(PartPlan {
index,
part_number: (index + 1) as i32,
offset,
len,
});
}
let config = Arc::new(configuration.clone());
let upload_id = Arc::new(session.upload_id.clone());
let finalize_token = Arc::new(session.finalize_token.clone());
let headers = Arc::new(session.headers.clone());
let path = Arc::new(path.to_path_buf());
let retry = configuration.retry;
let done = Arc::new(AtomicU64::new(0));
let progress = progress.cloned();
let uploader = move |plan: PartPlan| {
let config = Arc::clone(&config);
let upload_id = Arc::clone(&upload_id);
let finalize_token = Arc::clone(&finalize_token);
let headers = Arc::clone(&headers);
let path = Arc::clone(&path);
let done = Arc::clone(&done);
let progress = progress.clone();
async move {
let minted = apis::uploads_api::mint_upload_parts_handler(
&config,
&upload_id,
&finalize_token,
models::MintUploadPartsRequest::new(vec![plan.part_number]),
)
.await
.map_err(UploadError::MintParts)?;
let url = minted
.parts
.into_iter()
.find(|p| p.part_number == plan.part_number)
.map(|p| p.url)
.ok_or_else(|| {
UploadError::MalformedSession(format!(
"mint returned no URL for part {}",
plan.part_number
))
})?;
let chunk = read_range(&path, plan.offset, plan.len).await?;
let resp = put_to_storage(
&retry,
&url,
&headers,
chunk,
plan.len,
Some(plan.part_number),
)
.await?;
let e_tag = parse_etag(resp.headers(), plan.part_number)?;
if let Some(progress) = progress.as_ref() {
let now = done.fetch_add(plan.len, Ordering::SeqCst) + plan.len;
progress(now, total);
}
Ok(models::FinalizeUploadPart {
e_tag,
part_number: plan.part_number,
})
}
};
upload_parts_resilient(plans, in_flight_cap, RetryRounds::default(), uploader).await
}
fn parse_etag(
headers: &reqwest::header::HeaderMap,
part_number: i32,
) -> Result<String, UploadError> {
let etag = headers
.get(reqwest::header::ETAG)
.and_then(|v| v.to_str().ok())
.map(|s| s.to_owned())
.ok_or(UploadError::MissingETag { part_number })?;
if etag.trim().is_empty() {
return Err(UploadError::MissingETag { part_number });
}
Ok(etag)
}
async fn read_range(path: &Path, offset: u64, len: u64) -> Result<bytes::Bytes, UploadError> {
use tokio::io::{AsyncReadExt, AsyncSeekExt};
let mut file = tokio::fs::File::open(path).await?;
file.seek(std::io::SeekFrom::Start(offset)).await?;
let mut buf = vec![0u8; len as usize];
file.read_exact(&mut buf).await?;
Ok(bytes::Bytes::from(buf))
}
const STORAGE_CONNECT_TIMEOUT: Duration = Duration::from_secs(30);
const PART_TIMEOUT_BASE: Duration = Duration::from_secs(60);
const PART_TIMEOUT_MIN_BYTES_PER_SEC: u64 = 64 * 1024;
const PART_TIMEOUT_MAX: Duration = Duration::from_secs(30 * 60);
fn part_put_timeout(content_length: u64) -> Duration {
(PART_TIMEOUT_BASE + Duration::from_secs(content_length / PART_TIMEOUT_MIN_BYTES_PER_SEC))
.min(PART_TIMEOUT_MAX)
}
async fn put_to_storage(
retry: &crate::query::RetryPolicy,
url: &str,
headers: &HashMap<String, String>,
body: bytes::Bytes,
content_length: u64,
part_number: Option<i32>,
) -> Result<reqwest::Response, UploadError> {
let client = storage_client();
let mut req_builder = client
.request(reqwest::Method::PUT, url)
.header(reqwest::header::CONTENT_LENGTH, content_length);
for (name, value) in headers {
req_builder = req_builder.header(name.as_str(), value.as_str());
}
req_builder = req_builder.body(reqwest::Body::from(body));
req_builder = req_builder.timeout(part_put_timeout(content_length));
let req = req_builder.build().map_err(UploadError::Storage)?;
crate::http_log::log_request(&req);
let resp = crate::http::execute_retrying(&client, req, retry)
.await
.map_err(UploadError::Storage)?;
let status = resp.status();
crate::http_log::log_response_status(status);
if status.is_client_error() || status.is_server_error() {
let body = resp.text().await.unwrap_or_default();
crate::http_log::log_response_body(&body);
return Err(UploadError::StorageStatus {
status,
part_number,
body,
});
}
Ok(resp)
}
async fn put_stream_to_storage<S>(
url: &str,
headers: &HashMap<String, String>,
body: S,
content_length: u64,
) -> Result<reqwest::Response, UploadError>
where
S: futures_core::Stream<Item = std::io::Result<bytes::Bytes>> + Send + 'static,
{
let client = storage_client();
let mut req_builder = client
.request(reqwest::Method::PUT, url)
.header(reqwest::header::CONTENT_LENGTH, content_length);
for (name, value) in headers {
req_builder = req_builder.header(name.as_str(), value.as_str());
}
req_builder = req_builder.body(reqwest::Body::wrap_stream(body));
let req = req_builder.build().map_err(UploadError::Storage)?;
crate::http_log::log_request(&req);
let resp = client.execute(req).await.map_err(UploadError::Storage)?;
let status = resp.status();
crate::http_log::log_response_status(status);
if status.is_client_error() || status.is_server_error() {
let body = resp.text().await.unwrap_or_default();
crate::http_log::log_response_body(&body);
return Err(UploadError::StorageStatus {
status,
part_number: None,
body,
});
}
Ok(resp)
}
fn storage_client() -> reqwest::Client {
static STORAGE_CLIENT: std::sync::OnceLock<reqwest::Client> = std::sync::OnceLock::new();
STORAGE_CLIENT
.get_or_init(|| {
reqwest::Client::builder()
.connect_timeout(STORAGE_CONNECT_TIMEOUT)
.build()
.unwrap_or_default()
})
.clone()
}
#[cfg(test)]
mod tests {
use super::*;
fn part_count(size: u64, part: u64) -> u64 {
size.div_ceil(part)
}
#[test]
fn auto_part_size_keeps_8mib_for_normal_files() {
assert_eq!(auto_part_size_hint(0), DEFAULT_PART_SIZE);
assert_eq!(auto_part_size_hint(1), DEFAULT_PART_SIZE);
assert_eq!(auto_part_size_hint(100 * MIB), DEFAULT_PART_SIZE);
assert_eq!(auto_part_size_hint(1024 * MIB), DEFAULT_PART_SIZE); let boundary = DEFAULT_PART_SIZE * TARGET_MAX_PARTS;
assert_eq!(auto_part_size_hint(boundary), DEFAULT_PART_SIZE);
}
#[test]
fn auto_part_size_scales_up_for_very_large_files_and_caps_parts() {
let big = 200 * 1024 * MIB; let hint = auto_part_size_hint(big);
assert!(
hint > DEFAULT_PART_SIZE,
"hint should scale above 8 MiB for a 200 GiB file, got {hint}"
);
assert_eq!(hint % MIB, 0, "hint must be a whole MiB, got {hint}");
assert!(
part_count(big, hint) <= TARGET_MAX_PARTS,
"part count {} must be <= {TARGET_MAX_PARTS}",
part_count(big, hint)
);
assert!((MIN_PART_SIZE..=MAX_PART_SIZE).contains(&hint));
}
#[test]
fn auto_part_size_clamps_to_max_for_enormous_files() {
let enormous = 100 * 1024 * 1024 * MIB; assert_eq!(auto_part_size_hint(enormous), MAX_PART_SIZE);
}
#[test]
fn effective_in_flight_capped_by_max_concurrency_for_small_parts() {
assert_eq!(effective_in_flight(12, 8 * MIB), 12);
assert_eq!(effective_in_flight(10, 8 * MIB), 10);
assert_eq!(effective_in_flight(12, MIB), 12);
}
#[test]
fn effective_in_flight_reduced_by_memory_budget_for_large_parts() {
assert_eq!(effective_in_flight(12, 64 * MIB), 4);
assert_eq!(effective_in_flight(12, 128 * MIB), 2);
}
#[test]
fn effective_in_flight_honors_explicit_low_concurrency() {
assert_eq!(effective_in_flight(1, 8 * MIB), 1);
assert_eq!(effective_in_flight(0, 8 * MIB), 1);
assert_eq!(effective_in_flight(2, 8 * MIB), 2);
}
#[test]
fn effective_in_flight_floors_at_1_for_huge_parts_and_handles_zero() {
assert_eq!(effective_in_flight(12, UPLOAD_MEMORY_BUDGET * 4), 1);
assert_eq!(effective_in_flight(12, 0), 12);
}
}
#[cfg(test)]
mod resilient_retry_tests {
use super::*;
use std::collections::HashMap;
use std::sync::atomic::AtomicUsize;
use std::sync::Mutex;
fn plan(n: i32) -> PartPlan {
PartPlan {
index: (n - 1) as usize,
part_number: n,
offset: (n as u64 - 1) * 16,
len: 16,
}
}
fn plans(count: i32) -> Vec<PartPlan> {
(1..=count).map(plan).collect()
}
fn no_delay(max_extra_rounds: u32) -> RetryRounds {
RetryRounds {
max_extra_rounds,
base_delay: Duration::ZERO,
}
}
#[derive(Clone)]
struct FakeUploader {
fail: Arc<Mutex<HashMap<i32, usize>>>, attempts: Arc<Mutex<HashMap<i32, usize>>>,
in_flight: Arc<AtomicUsize>,
peak_in_flight: Arc<AtomicUsize>,
}
impl FakeUploader {
fn new(fail: HashMap<i32, usize>) -> Self {
Self {
fail: Arc::new(Mutex::new(fail)),
attempts: Arc::new(Mutex::new(HashMap::new())),
in_flight: Arc::new(AtomicUsize::new(0)),
peak_in_flight: Arc::new(AtomicUsize::new(0)),
}
}
fn attempts_for(&self, n: i32) -> usize {
*self.attempts.lock().unwrap().get(&n).unwrap_or(&0)
}
fn peak(&self) -> usize {
self.peak_in_flight.load(Ordering::SeqCst)
}
fn call(
&self,
plan: PartPlan,
) -> impl std::future::Future<Output = Result<models::FinalizeUploadPart, UploadError>>
+ Send
+ 'static {
let fail = Arc::clone(&self.fail);
let attempts = Arc::clone(&self.attempts);
let in_flight = Arc::clone(&self.in_flight);
let peak = Arc::clone(&self.peak_in_flight);
async move {
let cur = in_flight.fetch_add(1, Ordering::SeqCst) + 1;
peak.fetch_max(cur, Ordering::SeqCst);
tokio::task::yield_now().await;
*attempts
.lock()
.unwrap()
.entry(plan.part_number)
.or_insert(0) += 1;
let should_fail = {
let mut f = fail.lock().unwrap();
match f.get_mut(&plan.part_number) {
Some(remaining) if *remaining > 0 => {
*remaining -= 1;
true
}
_ => false,
}
};
in_flight.fetch_sub(1, Ordering::SeqCst);
if should_fail {
Err(UploadError::Io(std::io::Error::other(
"simulated connection reset",
)))
} else {
Ok(models::FinalizeUploadPart {
e_tag: format!("etag-{}", plan.part_number),
part_number: plan.part_number,
})
}
}
}
}
#[test]
fn round_in_flight_halves_each_round_min_one() {
assert_eq!(round_in_flight(8, 0), 8);
assert_eq!(round_in_flight(8, 1), 4);
assert_eq!(round_in_flight(8, 2), 2);
assert_eq!(round_in_flight(8, 3), 1);
assert_eq!(round_in_flight(8, 99), 1);
assert_eq!(round_in_flight(1, 3), 1);
}
#[test]
fn delay_for_grows_exponentially() {
let r = RetryRounds {
max_extra_rounds: 3,
base_delay: Duration::from_secs(2),
};
assert_eq!(r.delay_for(1), Duration::from_secs(2));
assert_eq!(r.delay_for(2), Duration::from_secs(4));
assert_eq!(r.delay_for(3), Duration::from_secs(8));
}
#[test]
fn part_put_timeout_is_generous_and_scales_with_part_size() {
assert_eq!(part_put_timeout(8 * 1024 * 1024), Duration::from_secs(188));
assert_eq!(
part_put_timeout(64 * 1024 * 1024),
Duration::from_secs(60 + 1024)
);
assert_eq!(part_put_timeout(0), Duration::from_secs(60));
assert!(part_put_timeout(32 * 1024 * 1024) > part_put_timeout(8 * 1024 * 1024));
assert_eq!(part_put_timeout(5 * 1024 * 1024 * 1024), PART_TIMEOUT_MAX);
assert_eq!(part_put_timeout(u64::MAX), PART_TIMEOUT_MAX);
}
#[test]
fn terminal_errors_are_only_contract_violations() {
assert!(is_terminal(&UploadError::MalformedSession("bad".into())));
assert!(is_terminal(&UploadError::SizeOverflow {
what: "x",
value: 1,
}));
assert!(!is_terminal(&UploadError::Io(std::io::Error::other(
"reset"
))));
assert!(!is_terminal(&UploadError::MissingETag { part_number: 1 }));
assert!(!is_terminal(&UploadError::StorageStatus {
status: reqwest::StatusCode::INTERNAL_SERVER_ERROR,
part_number: Some(1),
body: String::new(),
}));
}
#[test]
fn parse_etag_rejects_missing_and_blank() {
use reqwest::header::{HeaderMap, HeaderValue, ETAG};
let mut ok = HeaderMap::new();
ok.insert(ETAG, HeaderValue::from_static("\"etag-7\""));
assert_eq!(parse_etag(&ok, 7).unwrap(), "\"etag-7\"");
assert!(matches!(
parse_etag(&HeaderMap::new(), 7),
Err(UploadError::MissingETag { part_number: 7 })
));
for blank in ["", " "] {
let mut h = HeaderMap::new();
h.insert(ETAG, HeaderValue::from_str(blank).unwrap());
assert!(
matches!(parse_etag(&h, 7), Err(UploadError::MissingETag { .. })),
"blank ETag {blank:?} must be rejected"
);
}
}
#[tokio::test]
async fn repro_single_part_blip_sinks_whole_upload_without_rounds() {
let fake = FakeUploader::new(HashMap::from([(3, 1)]));
let f = fake.clone();
let res = upload_parts_resilient(plans(5), 4, no_delay(0), move |p| f.call(p)).await;
assert!(
res.is_err(),
"a single transient part failure should sink the upload under legacy (0-round) semantics"
);
}
#[tokio::test]
async fn fix_single_part_blip_recovers_on_a_later_round() {
let fake = FakeUploader::new(HashMap::from([(3, 1)]));
let f = fake.clone();
let res = upload_parts_resilient(plans(5), 4, no_delay(3), move |p| f.call(p))
.await
.expect("the flaky part should recover on a later round");
let nums: Vec<i32> = res.iter().map(|p| p.part_number).collect();
assert_eq!(nums, vec![1, 2, 3, 4, 5]);
assert_eq!(res[2].e_tag, "etag-3");
assert_eq!(fake.attempts_for(3), 2);
for n in [1, 2, 4, 5] {
assert_eq!(fake.attempts_for(n), 1, "part {n} must not be re-uploaded");
}
}
#[tokio::test]
async fn fix_multiple_flaky_parts_all_recover() {
let fake = FakeUploader::new(HashMap::from([(2, 2), (5, 1), (7, 3)]));
let f = fake.clone();
let res = upload_parts_resilient(plans(8), 4, no_delay(3), move |p| f.call(p))
.await
.expect("all parts should recover within the round budget");
assert_eq!(res.len(), 8);
assert_eq!(fake.attempts_for(2), 3); assert_eq!(fake.attempts_for(7), 4); assert_eq!(fake.attempts_for(5), 2);
}
#[tokio::test]
async fn permanent_failure_surfaced_after_exhausting_rounds() {
let fake = FakeUploader::new(HashMap::from([(4, 99)]));
let f = fake.clone();
let res = upload_parts_resilient(plans(5), 4, no_delay(2), move |p| f.call(p)).await;
assert!(res.is_err());
assert_eq!(fake.attempts_for(4), 3, "1 initial pass + 2 re-sweeps");
for n in [1, 2, 3, 5] {
assert_eq!(fake.attempts_for(n), 1);
}
}
#[tokio::test]
async fn happy_path_uploads_each_part_exactly_once() {
let fake = FakeUploader::new(HashMap::new());
let f = fake.clone();
let res = upload_parts_resilient(plans(6), 4, no_delay(3), move |p| f.call(p))
.await
.unwrap();
assert_eq!(res.len(), 6);
for n in 1..=6 {
assert_eq!(fake.attempts_for(n), 1);
}
}
#[tokio::test]
async fn concurrency_never_exceeds_base_cap() {
let fake = FakeUploader::new(HashMap::new());
let f = fake.clone();
upload_parts_resilient(plans(20), 3, no_delay(3), move |p| f.call(p))
.await
.unwrap();
assert!(
fake.peak() <= 3,
"peak in-flight {} exceeded the cap of 3",
fake.peak()
);
}
#[tokio::test]
async fn terminal_error_fails_fast_without_resweeping() {
let p2_attempts = Arc::new(AtomicUsize::new(0));
let counter = Arc::clone(&p2_attempts);
let res = upload_parts_resilient(plans(4), 4, no_delay(3), move |plan: PartPlan| {
let counter = Arc::clone(&counter);
async move {
if plan.part_number == 2 {
counter.fetch_add(1, Ordering::SeqCst);
Err(UploadError::MalformedSession("contract violation".into()))
} else {
Ok(models::FinalizeUploadPart {
e_tag: format!("etag-{}", plan.part_number),
part_number: plan.part_number,
})
}
}
})
.await;
assert!(matches!(res, Err(UploadError::MalformedSession(_))));
assert_eq!(
p2_attempts.load(Ordering::SeqCst),
1,
"a terminal error must be attempted once, never re-swept across rounds"
);
}
}