s3_wire/client/
multipart_upload.rs

1//! Bounded orchestration for managed multipart uploads.
2
3use std::future::Future;
4use std::path::PathBuf;
5use std::pin::Pin;
6use std::sync::Arc;
7use std::task::{Context, Poll};
8use std::time::Duration;
9
10use bytes::Bytes;
11use futures_util::stream::{FuturesUnordered, StreamExt as _};
12use tokio::sync::watch;
13use tokio::task::JoinHandle;
14
15use super::S3Client;
16use super::request::OperationDeadline;
17use crate::error::S3Error;
18use crate::operation::{
19    AbortMultipartUploadRequest, CompleteMultipartUploadOutput, CompleteMultipartUploadRequest,
20    CompletedPart, CreateMultipartUploadRequest, ManagedMultipartUploadRequest, MultipartUpload,
21    MultipartUploadSource, ObjectKey, PartNumber, UploadId, UploadPartRequest,
22};
23use crate::stream::{ByteStream, FileSnapshot};
24
25impl S3Client {
26    /// Uploads a replayable bytes or file source using bounded multipart requests.
27    ///
28    /// Part size, concurrency, total in-flight part bytes, and deadlines are
29    /// derived from the request's [`crate::MultipartOptions`]. File input is
30    /// snapshotted before an upload is created. Dropping this future cancels
31    /// outstanding parts and leaves an owned cleanup task to quiesce transmitted
32    /// requests before aborting the upload. A normal failure waits for cleanup;
33    /// if cleanup cannot be confirmed,
34    /// [`S3Error::cleanup_failure`] exposes that error while preserving the
35    /// original failure.
36    ///
37    /// # Errors
38    ///
39    /// Returns an error if the source cannot be prepared, exceeds S3's 10,000
40    /// part limit, a part fails, cancellation occurs, completion fails, or abort
41    /// cleanup fails after another error.
42    pub async fn multipart_upload(
43        &self,
44        request: ManagedMultipartUploadRequest,
45    ) -> Result<CompleteMultipartUploadOutput, S3Error> {
46        let runtime = tokio::runtime::Handle::try_current().map_err(S3Error::transport)?;
47        let transfer_deadline = OperationDeadline::new(request.options().transfer_timeout());
48        let cleanup_timeout = request.options().cleanup_timeout();
49        let cancellation = Cancellation::new();
50        let worker_cancellation = cancellation.clone();
51        let client = self.clone();
52        let handle = runtime.spawn(async move {
53            run_multipart_upload(
54                client,
55                request,
56                worker_cancellation,
57                transfer_deadline,
58                cleanup_timeout,
59            )
60            .await
61        });
62        OwnedMultipartTask {
63            cancellation,
64            handle,
65        }
66        .await
67    }
68}
69
70struct OwnedMultipartTask {
71    cancellation: Cancellation,
72    handle: JoinHandle<Result<CompleteMultipartUploadOutput, S3Error>>,
73}
74
75impl Future for OwnedMultipartTask {
76    type Output = Result<CompleteMultipartUploadOutput, S3Error>;
77
78    fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
79        match Pin::new(&mut self.handle).poll(context) {
80            Poll::Ready(Ok(result)) => Poll::Ready(result),
81            Poll::Ready(Err(error)) => Poll::Ready(Err(S3Error::transport(error))),
82            Poll::Pending => Poll::Pending,
83        }
84    }
85}
86
87impl Drop for OwnedMultipartTask {
88    fn drop(&mut self) {
89        self.cancellation.cancel();
90    }
91}
92
93#[derive(Clone)]
94struct Cancellation {
95    sender: watch::Sender<bool>,
96}
97
98impl Cancellation {
99    fn new() -> Self {
100        let (sender, _) = watch::channel(false);
101        Self { sender }
102    }
103
104    fn cancel(&self) {
105        self.sender.send_replace(true);
106    }
107
108    fn is_cancelled(&self) -> bool {
109        *self.sender.borrow()
110    }
111
112    async fn cancelled(&self) {
113        let mut receiver = self.sender.subscribe();
114        while !*receiver.borrow() {
115            if receiver.changed().await.is_err() {
116                return;
117            }
118        }
119    }
120}
121
122async fn run_multipart_upload(
123    client: S3Client,
124    request: ManagedMultipartUploadRequest,
125    cancellation: Cancellation,
126    transfer_deadline: OperationDeadline,
127    cleanup_timeout: Duration,
128) -> Result<CompleteMultipartUploadOutput, S3Error> {
129    let options = request.options();
130    let source = tokio::select! {
131        biased;
132        () = cancellation.cancelled() => return Err(cancelled()),
133        result = tokio::time::timeout_at(
134            transfer_deadline.instant(),
135            PreparedMultipartSource::prepare(request.source),
136        ) => result.map_err(|_| transfer_timeout())??,
137    };
138    let part_size = options.part_size();
139    let part_count = source.part_count(part_size)?;
140    if cancellation.is_cancelled() {
141        return Err(cancelled());
142    }
143
144    let mut create = CreateMultipartUploadRequest::new(request.key.clone());
145    create.content_type = request.content_type;
146    create.user_metadata = request.user_metadata;
147    // Once creation starts, allow it to finish even after cancellation so a
148    // successful response cannot be discarded together with the upload ID
149    // needed for cleanup.
150    let created = client
151        .create_multipart_upload_with_deadline(create, &transfer_deadline)
152        .await?;
153    let upload_id = created.upload_id().clone();
154
155    let result = upload_parts(
156        PartUploadContext {
157            client: client.clone(),
158            key: request.key.clone(),
159            upload_id: upload_id.clone(),
160            source: Arc::new(source),
161            part_size,
162            deadline: transfer_deadline,
163        },
164        part_count,
165        options.concurrency(),
166        &cancellation,
167        cleanup_timeout,
168    )
169    .await;
170
171    let upload = match result {
172        Ok(upload) => upload,
173        Err(failure) => {
174            return fail_with_abort(&client, request.key, upload_id, failure).await;
175        }
176    };
177    let completion = match CompleteMultipartUploadRequest::new(
178        upload.key().clone(),
179        upload.upload_id().clone(),
180        upload.completed_parts().to_vec(),
181    ) {
182        Ok(completion) => completion,
183        Err(error) => {
184            return fail_with_abort(
185                &client,
186                request.key,
187                upload_id,
188                ManagedUploadFailure::new(S3Error::integrity(error.to_string()), cleanup_timeout),
189            )
190            .await;
191        }
192    };
193
194    let completion = tokio::select! {
195        biased;
196        result = client.complete_multipart_upload_with_deadline(
197            completion,
198            &transfer_deadline,
199        ) => result,
200        () = cancellation.cancelled() => Err(cancelled()),
201    };
202
203    match completion {
204        Ok(output) => Ok(output),
205        Err(primary) => {
206            fail_with_abort(
207                &client,
208                request.key,
209                upload_id,
210                ManagedUploadFailure::new(primary, cleanup_timeout),
211            )
212            .await
213        }
214    }
215}
216
217#[derive(Clone)]
218struct PartUploadContext {
219    client: S3Client,
220    key: ObjectKey,
221    upload_id: UploadId,
222    source: Arc<PreparedMultipartSource>,
223    part_size: u64,
224    deadline: OperationDeadline,
225}
226
227async fn upload_parts(
228    context: PartUploadContext,
229    part_count: u16,
230    concurrency: usize,
231    cancellation: &Cancellation,
232    cleanup_timeout: Duration,
233) -> Result<MultipartUpload, ManagedUploadFailure> {
234    let mut next_part = 1_u16;
235    let mut pending = FuturesUnordered::new();
236    let mut upload = MultipartUpload::new(context.key.clone(), context.upload_id.clone());
237    let deadline_sleep = tokio::time::sleep_until(context.deadline.instant());
238    tokio::pin!(deadline_sleep);
239
240    loop {
241        while pending.len() < concurrency && next_part <= part_count {
242            pending.push(upload_one_part(context.clone(), next_part));
243            next_part += 1;
244        }
245        if pending.is_empty() {
246            return Ok(upload);
247        }
248        let completed = tokio::select! {
249            () = cancellation.cancelled() => Err(cancelled()),
250            () = &mut deadline_sleep => Err(transfer_timeout()),
251            result = pending.next() => result
252                .ok_or_else(|| S3Error::integrity("multipart scheduler lost an in-flight part"))
253                .and_then(|result| result),
254        };
255        let completed = match completed {
256            Ok(completed) => completed,
257            Err(primary) => {
258                return Err(quiesce_part_requests(primary, &mut pending, cleanup_timeout).await);
259            }
260        };
261        if let Err(error) = upload.record_part(completed) {
262            return Err(quiesce_part_requests(
263                S3Error::integrity(error.to_string()),
264                &mut pending,
265                cleanup_timeout,
266            )
267            .await);
268        }
269    }
270}
271
272struct ManagedUploadFailure {
273    primary: S3Error,
274    cleanup_deadline: OperationDeadline,
275    quiesce_failure: Option<S3Error>,
276}
277
278impl ManagedUploadFailure {
279    fn new(primary: S3Error, cleanup_timeout: Duration) -> Self {
280        Self {
281            primary,
282            cleanup_deadline: OperationDeadline::new(cleanup_timeout),
283            quiesce_failure: None,
284        }
285    }
286}
287
288async fn quiesce_part_requests<F>(
289    primary: S3Error,
290    pending: &mut FuturesUnordered<F>,
291    cleanup_timeout: Duration,
292) -> ManagedUploadFailure
293where
294    F: Future<Output = Result<CompletedPart, S3Error>>,
295{
296    let mut failure = ManagedUploadFailure::new(primary, cleanup_timeout);
297    let settle_until = tokio::time::Instant::now() + cleanup_timeout / 2;
298    while !pending.is_empty() {
299        match tokio::time::timeout_at(settle_until, pending.next()).await {
300            Ok(Some(_)) => {}
301            Ok(None) => break,
302            Err(_) => {
303                failure.quiesce_failure = Some(S3Error::timeout(
304                    crate::error::TimeoutPhase::Operation,
305                    "multipart cleanup could not quiesce in-flight part requests",
306                ));
307                break;
308            }
309        }
310    }
311    failure
312}
313
314async fn fail_with_abort(
315    client: &S3Client,
316    key: ObjectKey,
317    upload_id: UploadId,
318    failure: ManagedUploadFailure,
319) -> Result<CompleteMultipartUploadOutput, S3Error> {
320    let abort = client
321        .abort_multipart_upload_with_deadline(
322            AbortMultipartUploadRequest::new(key, upload_id),
323            &failure.cleanup_deadline,
324        )
325        .await
326        .map(|_| ());
327    let cleanup = match (abort, failure.quiesce_failure) {
328        (Ok(()), None) => Ok(()),
329        (Err(error), None) | (Ok(()), Some(error)) => Err(error),
330        (Err(abort), Some(quiesce)) => Err(abort.with_cleanup_failure(quiesce)),
331    };
332    Err(attach_cleanup_failure(failure.primary, cleanup))
333}
334
335async fn upload_one_part(
336    context: PartUploadContext,
337    part_number: u16,
338) -> Result<CompletedPart, S3Error> {
339    let body = context
340        .source
341        .read_part(context.part_size, part_number)
342        .await?;
343    let number = PartNumber::new(part_number)
344        .ok_or_else(|| S3Error::integrity("generated multipart part number is invalid"))?;
345    let output = context
346        .client
347        .upload_part_with_deadline(
348            UploadPartRequest::new(
349                context.key,
350                context.upload_id,
351                number,
352                ByteStream::from_bytes(body),
353            ),
354            &context.deadline,
355        )
356        .await?;
357    CompletedPart::new(output.part_number.get(), output.e_tag)
358        .map(|part| part.with_checksum(output.checksum))
359        .map_err(|error| S3Error::invalid_response(error.to_string()))
360}
361
362fn attach_cleanup_failure(primary: S3Error, cleanup: Result<(), S3Error>) -> S3Error {
363    match cleanup {
364        Ok(()) => primary,
365        Err(cleanup) => primary.with_cleanup_failure(cleanup),
366    }
367}
368
369fn cancelled() -> S3Error {
370    S3Error::cancellation("multipart upload was cancelled")
371}
372
373fn transfer_timeout() -> S3Error {
374    S3Error::timeout(
375        crate::error::TimeoutPhase::Operation,
376        "managed multipart upload exceeded its transfer deadline",
377    )
378}
379
380enum PreparedMultipartSource {
381    Bytes(Bytes),
382    File(FileSnapshot),
383}
384
385impl PreparedMultipartSource {
386    async fn prepare(source: MultipartUploadSource) -> Result<Self, S3Error> {
387        match source {
388            MultipartUploadSource::Bytes(bytes) => Ok(Self::Bytes(bytes)),
389            MultipartUploadSource::File(path) => Self::snapshot(path).await,
390        }
391    }
392
393    async fn snapshot(path: PathBuf) -> Result<Self, S3Error> {
394        FileSnapshot::create(path, false).await.map(Self::File)
395    }
396
397    fn length(&self) -> Result<u64, S3Error> {
398        match self {
399            Self::Bytes(bytes) => u64::try_from(bytes.len())
400                .map_err(|_| S3Error::configuration("multipart byte length does not fit in u64")),
401            Self::File(snapshot) => Ok(snapshot.length()),
402        }
403    }
404
405    fn part_count(&self, part_size: u64) -> Result<u16, S3Error> {
406        let length = self.length()?;
407        if length == 0 {
408            return Err(S3Error::configuration(
409                "managed multipart uploads require a non-empty source",
410            ));
411        }
412        let count = length.div_ceil(part_size);
413        if count > u64::from(PartNumber::MAX) {
414            return Err(S3Error::configuration(
415                "multipart source exceeds the configured 10,000-part limit",
416            ));
417        }
418        u16::try_from(count)
419            .map_err(|_| S3Error::configuration("multipart part count does not fit in u16"))
420    }
421
422    async fn read_part(&self, part_size: u64, part_number: u16) -> Result<Bytes, S3Error> {
423        let offset = u64::from(part_number - 1)
424            .checked_mul(part_size)
425            .ok_or_else(|| S3Error::integrity("multipart part offset overflow"))?;
426        let remaining = self
427            .length()?
428            .checked_sub(offset)
429            .ok_or_else(|| S3Error::integrity("multipart part offset exceeds source length"))?;
430        let length = remaining.min(part_size);
431        let length_usize = usize::try_from(length)
432            .map_err(|_| S3Error::configuration("multipart part size does not fit in usize"))?;
433        match self {
434            Self::Bytes(bytes) => {
435                let start = usize::try_from(offset).map_err(|_| {
436                    S3Error::configuration("multipart byte offset does not fit in usize")
437                })?;
438                Ok(bytes.slice(start..start + length_usize))
439            }
440            Self::File(snapshot) => snapshot.read_range(offset, length_usize).await,
441        }
442    }
443}
444
445#[cfg(test)]
446mod tests {
447    use super::*;
448    use crate::error::ErrorCategory;
449
450    #[tokio::test]
451    async fn task_drop_signals_owned_cleanup_worker() {
452        let cancellation = Cancellation::new();
453        let worker = cancellation.clone();
454        let (finished, observed) = tokio::sync::oneshot::channel();
455        let handle = tokio::spawn(async move {
456            worker.cancelled().await;
457            let _ = finished.send(());
458            Err(cancelled())
459        });
460        drop(OwnedMultipartTask {
461            cancellation,
462            handle,
463        });
464        observed.await.unwrap();
465    }
466
467    #[test]
468    fn cleanup_failure_does_not_replace_primary_error() {
469        let primary = S3Error::integrity("primary");
470        let cleanup = S3Error::invalid_response("cleanup");
471        let combined = attach_cleanup_failure(primary, Err(cleanup));
472        assert_eq!(combined.category(), ErrorCategory::Integrity);
473        assert_eq!(combined.message(), "primary");
474        assert_eq!(
475            combined.cleanup_failure().map(S3Error::category),
476            Some(ErrorCategory::InvalidResponse)
477        );
478    }
479
480    #[tokio::test]
481    async fn byte_parts_are_exact_and_bounded() {
482        let source = PreparedMultipartSource::Bytes(Bytes::from_static(b"abcdefghij"));
483        assert_eq!(source.part_count(4).unwrap(), 3);
484        assert_eq!(source.read_part(4, 1).await.unwrap(), "abcd");
485        assert_eq!(source.read_part(4, 2).await.unwrap(), "efgh");
486        assert_eq!(source.read_part(4, 3).await.unwrap(), "ij");
487    }
488
489    #[tokio::test]
490    async fn file_source_is_snapshotted_before_parts_are_read() {
491        let directory = tempfile::tempdir().unwrap();
492        let path = directory.path().join("source");
493        tokio::fs::write(&path, b"original").await.unwrap();
494        let source = PreparedMultipartSource::snapshot(path.clone())
495            .await
496            .unwrap();
497        tokio::fs::write(path, b"changed!").await.unwrap();
498        assert_eq!(source.read_part(8, 1).await.unwrap(), "original");
499    }
500
501    #[tokio::test]
502    async fn cleanup_reports_part_requests_that_cannot_quiesce() {
503        let mut pending = FuturesUnordered::new();
504        pending.push(std::future::pending::<Result<CompletedPart, S3Error>>());
505
506        let failure = quiesce_part_requests(
507            S3Error::cancellation("primary"),
508            &mut pending,
509            Duration::from_millis(10),
510        )
511        .await;
512
513        assert!(failure.quiesce_failure.is_some());
514        assert_eq!(failure.primary.category(), ErrorCategory::Cancellation);
515    }
516
517    #[test]
518    fn part_count_rejects_empty_and_excessive_sources() {
519        let empty = PreparedMultipartSource::Bytes(Bytes::new());
520        assert!(empty.part_count(5).is_err());
521        let too_large = PreparedMultipartSource::Bytes(Bytes::from(vec![0_u8; 10_001]));
522        assert!(too_large.part_count(1).is_err());
523    }
524}
s3_wire/client/multipart_upload.rs

s3_wire/client/
multipart_upload.rs
