s3_wire/operation/multipart/

managed.rs

use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::time::Duration;

use bytes::Bytes;
use http::HeaderMap;

use super::{CompletedPart, MultipartError, UploadId};
use crate::operation::{ChecksumAlgorithm, ChecksumType, ObjectKey, RequestIds};

const MIN_PART_SIZE: u64 = 5 * 1024 * 1024;
const MAX_PART_SIZE: u64 = 5 * 1024 * 1024 * 1024;
const MAX_CONCURRENCY: usize = 64;

/// Validated resource and time bounds for one managed multipart upload.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct MultipartOptions {
    part_size: u64,
    concurrency: usize,
    transfer_timeout: Duration,
    cleanup_timeout: Duration,
}

impl MultipartOptions {
    /// Creates multipart options with bounded transfer and cleanup defaults.
    ///
    /// # Errors
    ///
    /// Returns an error unless `part_size` is between 5 MiB and 5 GiB and
    /// `concurrency` is between 1 and 64.
    pub fn new(part_size: u64, concurrency: usize) -> Result<Self, crate::error::S3Error> {
        if !(MIN_PART_SIZE..=MAX_PART_SIZE).contains(&part_size) {
            return Err(crate::error::S3Error::configuration(
                "multipart part size must be between 5 MiB and 5 GiB",
            ));
        }
        if !(1..=MAX_CONCURRENCY).contains(&concurrency) {
            return Err(crate::error::S3Error::configuration(
                "multipart concurrency must be between 1 and 64",
            ));
        }
        part_size
            .checked_mul(u64::try_from(concurrency).map_err(|_| {
                crate::error::S3Error::configuration("multipart concurrency does not fit in u64")
            })?)
            .ok_or_else(|| {
                crate::error::S3Error::configuration("multipart buffered byte bound overflow")
            })?;
        Ok(Self {
            part_size,
            concurrency,
            transfer_timeout: Duration::from_secs(5 * 60),
            cleanup_timeout: Duration::from_secs(30),
        })
    }

    /// Sets the deadline for preparation, upload, and completion.
    ///
    /// # Errors
    ///
    /// Returns an error when `timeout` is zero.
    pub fn with_transfer_timeout(
        mut self,
        timeout: Duration,
    ) -> Result<Self, crate::error::S3Error> {
        if timeout.is_zero() {
            return Err(crate::error::S3Error::configuration(
                "multipart transfer timeout must be greater than zero",
            ));
        }
        validate_deadline(timeout, "multipart transfer timeout")?;
        self.transfer_timeout = timeout;
        Ok(self)
    }

    /// Sets the separate deadline for quiescing in-flight parts and aborting.
    ///
    /// # Errors
    ///
    /// Returns an error when `timeout` is zero.
    pub fn with_cleanup_timeout(
        mut self,
        timeout: Duration,
    ) -> Result<Self, crate::error::S3Error> {
        if timeout.is_zero() {
            return Err(crate::error::S3Error::configuration(
                "multipart cleanup timeout must be greater than zero",
            ));
        }
        validate_deadline(timeout, "multipart cleanup timeout")?;
        self.cleanup_timeout = timeout;
        Ok(self)
    }

    /// Returns the size of each part except the final part.
    pub const fn part_size(self) -> u64 {
        self.part_size
    }

    /// Returns the maximum number of part requests in flight.
    pub const fn concurrency(self) -> usize {
        self.concurrency
    }

    /// Returns the derived maximum bytes retained by in-flight part buffers.
    pub fn maximum_buffered_bytes(self) -> u64 {
        self.part_size * u64::try_from(self.concurrency).expect("validated concurrency fits in u64")
    }

    /// Returns the transfer deadline duration.
    pub const fn transfer_timeout(self) -> Duration {
        self.transfer_timeout
    }

    /// Returns the part-quiescing and abort cleanup deadline duration.
    pub const fn cleanup_timeout(self) -> Duration {
        self.cleanup_timeout
    }
}

fn validate_deadline(timeout: Duration, name: &str) -> Result<(), crate::error::S3Error> {
    std::time::Instant::now()
        .checked_add(timeout)
        .map(|_| ())
        .ok_or_else(|| {
            crate::error::S3Error::configuration(format!(
                "{name} is too large to represent as a deadline"
            ))
        })
}

impl Default for MultipartOptions {
    fn default() -> Self {
        Self::new(8 * 1024 * 1024, 4).expect("default multipart options are valid")
    }
}

/// Request to initiate a multipart upload.
#[derive(Clone, derive_more::Debug, Eq, PartialEq)]
pub struct CreateMultipartUploadRequest {
    /// Destination object key.
    pub key: ObjectKey,
    /// Optional object media type.
    pub content_type: Option<String>,
    /// User-defined object metadata.
    #[debug("{:?}", self.user_metadata.keys().collect::<Vec<_>>())]
    pub user_metadata: BTreeMap<String, String>,
    /// Checksum algorithm applied to uploaded parts.
    pub checksum_algorithm: Option<ChecksumAlgorithm>,
    /// How S3 should derive the completed object's checksum from its parts.
    pub checksum_type: Option<ChecksumType>,
    /// Additional request headers. Values are signed and repeated values are preserved.
    /// Generated-name collisions are errors; signing- and transport-owned headers are rejected.
    #[debug("{:?}", "<redacted>")]
    pub headers: HeaderMap,
}

impl CreateMultipartUploadRequest {
    /// Constructs a request without optional metadata.
    pub fn new(key: ObjectKey) -> Self {
        Self {
            key,
            content_type: None,
            user_metadata: BTreeMap::new(),
            checksum_algorithm: None,
            checksum_type: None,
            headers: HeaderMap::new(),
        }
    }

    /// Replaces the request's additional headers.
    pub fn with_headers(mut self, headers: HeaderMap) -> Self {
        self.headers = headers;
        self
    }
}

/// Result of initiating a multipart upload.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct CreateMultipartUploadOutput {
    /// Bucket reported by the service.
    pub bucket: Option<String>,
    /// Object key reported by the service.
    pub key: ObjectKey,
    /// Opaque upload identifier required by later operations.
    upload_id: UploadId,
    /// Checksum algorithm selected by the service.
    pub checksum_algorithm: Option<String>,
    /// Checksum aggregation selected by S3.
    pub checksum_type: Option<ChecksumType>,
    /// Service request identifiers.
    pub request_ids: RequestIds,
}

impl CreateMultipartUploadOutput {
    pub(crate) fn new(
        bucket: Option<String>,
        key: ObjectKey,
        upload_id: UploadId,
        checksum_algorithm: Option<String>,
        checksum_type: Option<ChecksumType>,
    ) -> Self {
        Self {
            bucket,
            key,
            upload_id,
            checksum_algorithm,
            checksum_type,
            request_ids: RequestIds::default(),
        }
    }

    /// Returns the validated upload identifier required by later operations.
    pub const fn upload_id(&self) -> &UploadId {
        &self.upload_id
    }
}

/// Explicit state needed to resume or clean up an in-progress multipart upload.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct MultipartUpload {
    key: ObjectKey,
    upload_id: UploadId,
    /// Parts that have completed successfully, in ascending order.
    completed_parts: Vec<CompletedPart>,
}

pub(crate) enum MultipartUploadSource {
    Bytes(Bytes),
    File(PathBuf),
}

/// Phase-specific headers for a managed multipart upload.
///
/// Content metadata, tagging, ACL, and SSE-KMS headers typically belong on
/// `create`. SSE-C headers must usually be repeated on `create` and
/// `upload_part`. Completion checksums and conditions belong on `complete`,
/// while requester-pays, expected-owner, and abort conditions belong on
/// `abort`. AWS support varies by bucket type and feature.
#[derive(Clone, Default, derive_more::Debug, Eq, PartialEq)]
pub struct ManagedMultipartHeaders {
    /// Headers sent when creating the upload. Values are signed and repeated values are preserved.
    /// Generated-name collisions are errors; signing- and transport-owned headers are rejected.
    #[debug("{:?}", "<redacted>")]
    pub create: HeaderMap,
    /// Headers cloned into every part. Values are signed and repeated values are preserved.
    /// Generated-name collisions are errors; signing- and transport-owned headers are rejected.
    #[debug("{:?}", "<redacted>")]
    pub upload_part: HeaderMap,
    /// Completion headers. Values are signed and repeated values are preserved.
    /// Generated-name collisions are errors; signing- and transport-owned headers are rejected.
    #[debug("{:?}", "<redacted>")]
    pub complete: HeaderMap,
    /// Abort headers. Values are signed and repeated values are preserved.
    /// Generated-name collisions are errors; signing- and transport-owned headers are rejected.
    #[debug("{:?}", "<redacted>")]
    pub abort: HeaderMap,
}

/// Request for a bounded, automatically cleaned-up multipart upload.
///
/// In-memory sources retain the caller's complete byte buffer for the duration
/// of the operation. File sources retain at most the derived in-flight part
/// byte bound in memory; their immutable snapshot is disk-backed.
#[derive(derive_more::Debug)]
pub struct ManagedMultipartUploadRequest {
    pub(crate) key: ObjectKey,
    #[debug(
        "{:?}",
        match self.source {
            MultipartUploadSource::Bytes(_) => "bytes",
            MultipartUploadSource::File(_) => "file",
        }
    )]
    pub(crate) source: MultipartUploadSource,
    pub(crate) content_type: Option<String>,
    #[debug("{:?}", self.user_metadata.keys().collect::<Vec<_>>())]
    pub(crate) user_metadata: BTreeMap<String, String>,
    pub(crate) options: MultipartOptions,
    pub(crate) checksum_algorithm: Option<ChecksumAlgorithm>,
    /// Phase-specific additional headers. Values are signed and repeated values are preserved.
    /// Generated-name collisions are errors; signing- and transport-owned headers are rejected.
    pub headers: ManagedMultipartHeaders,
}

impl ManagedMultipartUploadRequest {
    /// Constructs a request with a replayable in-memory body.
    pub fn from_bytes(key: ObjectKey, bytes: impl Into<Bytes>) -> Self {
        Self {
            key,
            source: MultipartUploadSource::Bytes(bytes.into()),
            content_type: None,
            user_metadata: BTreeMap::new(),
            options: MultipartOptions::default(),
            checksum_algorithm: None,
            headers: ManagedMultipartHeaders::default(),
        }
    }

    /// Constructs a request with a replayable regular-file body.
    ///
    /// The file is copied into a private disk-backed snapshot before S3 creates
    /// the multipart upload, so later changes to the source path have no effect.
    pub fn from_path(key: ObjectKey, path: impl AsRef<Path>) -> Self {
        Self {
            key,
            source: MultipartUploadSource::File(path.as_ref().to_owned()),
            content_type: None,
            user_metadata: BTreeMap::new(),
            options: MultipartOptions::default(),
            checksum_algorithm: None,
            headers: ManagedMultipartHeaders::default(),
        }
    }

    /// Replaces the upload's phase-specific additional headers.
    pub fn with_headers(mut self, headers: ManagedMultipartHeaders) -> Self {
        self.headers = headers;
        self
    }

    /// Sets the object's media type.
    pub fn with_content_type(mut self, content_type: impl Into<String>) -> Self {
        self.content_type = Some(content_type.into());
        self
    }

    /// Adds user-defined object metadata.
    pub fn with_metadata(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
        self.user_metadata.insert(name.into(), value.into());
        self
    }

    /// Applies validated resource and time bounds to this upload.
    pub fn with_options(mut self, options: MultipartOptions) -> Self {
        self.options = options;
        self
    }

    /// Calculates and sends this checksum for every uploaded part.
    ///
    /// # Errors
    ///
    /// Returns an error when local calculation is unavailable for the selected
    /// algorithm. SHA-1 may still be supplied through primitive multipart APIs.
    pub fn with_checksum_algorithm(
        mut self,
        algorithm: ChecksumAlgorithm,
    ) -> Result<Self, crate::operation::ChecksumCalculationError> {
        if algorithm == ChecksumAlgorithm::Sha1 {
            return Err(crate::operation::ChecksumCalculationError { algorithm });
        }
        self.checksum_algorithm = Some(algorithm);
        Ok(self)
    }

    /// Returns this upload's resource and time bounds.
    pub const fn options(&self) -> MultipartOptions {
        self.options
    }

    /// Returns the destination object key.
    pub const fn key(&self) -> &ObjectKey {
        &self.key
    }
}

impl MultipartUpload {
    /// Starts tracking a newly created upload.
    pub fn new(key: ObjectKey, upload_id: UploadId) -> Self {
        Self {
            key,
            upload_id,
            completed_parts: Vec::new(),
        }
    }

    /// Records a completed part while retaining canonical ascending order.
    pub fn record_part(&mut self, part: CompletedPart) -> Result<(), MultipartError> {
        match self
            .completed_parts
            .binary_search_by_key(&part.part_number(), CompletedPart::part_number)
        {
            Ok(_) => Err(MultipartError::DuplicatePart(part.part_number())),
            Err(index) => {
                self.completed_parts.insert(index, part);
                Ok(())
            }
        }
    }

    /// Returns completed parts in validated ascending order.
    pub fn completed_parts(&self) -> &[CompletedPart] {
        &self.completed_parts
    }

    /// Returns the destination object key.
    pub const fn key(&self) -> &ObjectKey {
        &self.key
    }

    /// Returns the non-empty service-issued upload identifier.
    pub const fn upload_id(&self) -> &UploadId {
        &self.upload_id
    }
}