move uploads to new scheduler by aajtodd · Pull Request #132 · awslabs/aws-s3-transfer-manager-rs

aajtodd · 2026-03-05T19:06:17Z

Summary

Replace the Tower service-based upload orchestration (ConcurrencyLimitLayer, service_fn, Buffer) with a scheduler-driven state machine. The upload transfer implements the Transfer trait from #131. The scheduler polls the transfer for work; the transfer produces work items lazily, one at a time.

Upload Lifecycle

Multipart upload

orchestrate() → UploadTransfer enqueued to scheduler

poll_work → CreateMPU           [Network]
poll_work → Pending              (waiting for upload_id)
  ... CreateMPU completes, wake ...
poll_work → UploadPart(1)       [DataIO]   ← disk read
poll_work → UploadPart(2)       [DataIO]   ← disk read
  ... DataIO completes → follow-on: UploadPart(1) [Network] ...
  ... all parts complete ...
poll_work → CompleteMPU         [Network]
poll_work → Done

Small file (below MPU threshold)

Skips multipart entirely — produces a single PutObject work item.

Design

Each UploadPart has two phases: a DataIO phase that reads the part from the PartReader, then a Network phase that sends it to S3. The Network phase is a follow-on work item that bypasses the ready set, so the buffer isn't held waiting behind other transfers.

Error handling follows the download transfer's contract:

fail() stores the original error via ctx.set_failed(error) and signals terminal
join() awaits completion_rx from TransferContext, then reads the error from ctx.take_error() on failure or the result from transfer.take_result() on success
No separate result channel — errors and results flow through TransferContext and the transfer struct respectively

What changed

File	Change
`operation/upload/transfer.rs`	New. Upload state machine implementing `Transfer` trait.
`operation/upload/upload.rs`	`orchestrate()` creates `UploadTransfer` and enqueues to scheduler instead of spawning Tower service pipeline.
`operation/upload/handle.rs`	Holds `UploadTransfer` + `completion_rx` directly. `join()` follows download's pattern. Abort waits for scheduler idle, calls `AbortMultipartUpload`.
`operation/upload/service.rs`	Deleted. Tower middleware chain no longer needed.
`scheduler/mod.rs`	Re-exports `StateMachineTerminalReceiver`.

What did not change

Download still uses the old runtime::scheduler and service pattern
upload_objects / download_objects unchanged
middleware/, runtime/ kept for download and multi-object operations
No scheduler changes beyond the re-export
No public API changes

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

landonxjames

Few questions about some of the logic. Nothing major or blocking.

landonxjames · 2026-03-10T18:45:26Z

aws-sdk-s3-transfer-manager/src/operation/upload/transfer.rs

+            return PollWork::Done;
+        }
+
+        let mut state = self.inner.state.lock().unwrap();


Suggested change

let mut state = self.inner.state.lock().unwrap();

let mut state = self.inner.state.lock().expect("Failed to unlock transfer state");

This suggestion applies to a few places, just putting it here so I don't spam it everywhere.

More generally, this Mutex seems like it is going to be locked for every poll_work and execute_* and like that could cause a lot of contention. I wonder if using something that is supposed to be more performant, like parking_lot::Mutex would make sense here?

Probably a premature optimization, but worth evaluating when you get to benchmarking this

landonxjames · 2026-03-10T18:48:57Z

aws-sdk-s3-transfer-manager/src/operation/upload.rs

-    distribute_work(&mut mpu_data, ctx)?;
-    Ok(UploadType::MultipartUpload(mpu_data))
-}
+        // TODO: Relax this constraint - unknown content length implies MPU


Suggested change

// TODO: Relax this constraint - unknown content length implies MPU

// TODO(vnext): Relax this constraint - unknown content length implies MPU

landonxjames · 2026-03-10T19:41:21Z

aws-sdk-s3-transfer-manager/src/operation/upload/transfer.rs

+                }))
+            }
+            Ok(None) => {
+                tracing::warn!("part_reader returned None for part {}", part_number);


Think this should just be a trace log? If content_length is just an upper bound instead of exact (which is how I think size_hint works?) then this will happen sometimes and shouldn't be exposed as something failing.

landonxjames · 2026-03-10T19:46:12Z

aws-sdk-s3-transfer-manager/src/operation/upload/transfer.rs

+
+        let (stream, content_length) = {
+            let mut state = self.inner.state.lock().unwrap();
+            match std::mem::replace(&mut *state, UploadState::Done) {


Is replacing this with Done correct? If poll_work is called in the window what state is set to Done won't it return PollWork::Done, causing the scheduler to remove
the transfer entirely? Seems like it would be possible since you have to re-acquire the lock on the state later in this function?

Good catch, no this is wrong.

landonxjames · 2026-03-10T19:48:13Z

aws-sdk-s3-transfer-manager/src/operation/upload/handle.rs

-        .clone()
-        .unwrap_or_default();
-    match abort_policy {
-        FailedMultipartUploadPolicy::Retain => Ok(AbortedUpload::default()),


We are no longer checking the abort_policy is that intended?

Nope will fix

landonxjames · 2026-03-10T19:50:01Z

aws-sdk-s3-transfer-manager/src/operation/upload/transfer.rs

+
+        let resp = match req
+            .customize()
+            .disable_payload_signing()


We disable this for MPU but not PutObject? Might make sense but would probably want a comment explaining why

Just a miss when converting.

landonxjames · 2026-03-10T19:51:19Z

aws-sdk-s3-transfer-manager/src/operation/upload.rs

+    use crate::types::{ConcurrencyMode, PartSize};

    #[tokio::test]
-    async fn test_basic_mpu() {


Doesn't seem to be a new version of the happy path test_basic_mpu and test_basic_upload_object tests?

ysaito1001 · 2026-03-10T19:57:18Z

Again, to aid code review:

a runtime sequence diagram for MPU

                          MPU Upload Runtime Lifecycle
                          ===========================
 
 builders.rs              upload.rs                transfer.rs              scheduler/          handle.rs
 (initiate)            (Upload::orchestrate)       (UploadTransfer)         (Scheduler)         (UploadHandle)
     │                        │                         │                       │                    │
     │──initiate_with()──────>│                         │                       │                    │
     │                        │ resolve checksum_strategy│                       │                    │
     │                        │ take_body() from input   │                       │                    │
     │                        │ check size_hint upper    │                       │                    │
     │                        │ BucketType::from_bucket_name                     │                    │
     │                        │                         │                       │                    │
     │                        │ TransferContext::new(handle)                     │                    │
     │                        │──────────────────────────> new(ctx, bucket_type, │                    │
     │                        │                           input, stream)        │                    │
     │                        │                         │                       │                    │
     │                        │                         │  state = PendingInit  │                    │
     │                        │                         │  {stream,             │                    │
     │                        │                         │   content_length,     │                    │
     │                        │                         │   init_in_flight:false}│                    │
     │                        │                         │                       │                    │
     │                        │─────────────────────────────enqueue_transfer(───>│                    │
     │                        │                         │  Box<transfer.clone>)  │                    │
     │                        │                         │                       │                    │
     │                        │──────────────────────────────────────────────────────> new(completion_rx,
     │<───UploadHandle────────│                         │                       │        transfer)   │
     │                        │                         │                       │                    │
 
 
                    ┌─── Scheduler poll loop begins ───┐
                    │                                   │
                    ▼                                   │
 
 scheduler/                          transfer.rs
 (Scheduler)                         (UploadTransfer)
     │                                    │
     │──poll_work()──────────────────────>│
     │                                    │ state == PendingInit
     │                                    │ content_length >= mpu_threshold
     │                                    │ set init_in_flight = true
     │<──PollWork::Ready(CreateMPU, ──────│
     │                   Network)         │
     │                                    │
     │──execute(CreateMPU)───────────────>│
     │                                    │ s3_client.create_multipart_upload().send()
     │                                    │ take stream + content_length from PendingInit
     │                                    │ compute part_size, total_parts
     │                                    │ build PartReader
     │                                    │ state = Transferring {
     │                                    │   upload_id, part_reader,
     │                                    │   next_part: 1, total_parts,
     │                                    │   parts_in_flight: 0,
     │                                    │   completed_parts: [],
     │                                    │   response_builder }
     │                                    │ create_mpu_complete.notify_waiters()
     │                                    │ ctx.try_wake()
     │<──WorkOutcome::Success(None)───────│
     │                                    │
     │──poll_work()──────────────────────>│  (repeated for each part 1..total_parts)
     │                                    │ state == Transferring
     │                                    │ next_part++, parts_in_flight++
     │<──PollWork::Ready(UploadPart{N},───│
     │                   DataIO)          │
     │                                    │
     │──execute(UploadPart{N}, DataIO)───>│
     │                                    │ part_reader.next_part()  (disk read)
     │                                    │ store part_data in work item
     │<──WorkOutcome::Success(Some(───────│  ◄── follow-on work item
     │     WorkItem{UploadPart{N},        │
     │              Network}))            │
     │                                    │
     │──execute(UploadPart{N}, Network)──>│
     │                                    │ s3_client.upload_part()
     │                                    │   .upload_id(&upload_id)
     │                                    │   .part_number(N)
     │                                    │   .body(data)
     │                                    │   .disable_payload_signing()
     │                                    │   .send()
     │                                    │ build CompletedPart from response
     │                                    │ push to completed_parts
     │                                    │ parts_in_flight--
     │                                    │ if next_part > total_parts && parts_in_flight == 0:
     │                                    │   state = Completing {
     │                                    │     upload_id, part_reader,
     │                                    │     completed_parts, response_builder,
     │                                    │     complete_in_flight: false }
     │                                    │   ctx.try_wake()
     │<──WorkOutcome::Success(None)───────│
     │                                    │
     │──poll_work()──────────────────────>│
     │                                    │ state == Completing
     │                                    │ set complete_in_flight = true
     │<──PollWork::Ready(CompleteMPU, ────│
     │                   Network)         │
     │                                    │
     │──execute(CompleteMPU)─────────────>│
     │                                    │ sort completed_parts by part_number
     │                                    │ s3_client.complete_multipart_upload()
     │                                    │   .upload_id(&upload_id)
     │                                    │   .multipart_upload(parts)
     │                                    │   .send()
     │                                    │ build UploadOutput via response_builder
     │                                    │ store in self.result
     │                                    │ ctx.set_completed()
     │                                    │ ctx.signal_terminal()
     │<──WorkOutcome::Success(None)───────│
     │                                    │
     │──poll_work()──────────────────────>│
     │                                    │ state == Done
     │<──PollWork::Done───────────────────│

ysaito1001

Left nit comments. Looks good!

ysaito1001 · 2026-03-10T17:41:42Z

aws-sdk-s3-transfer-manager/src/operation/upload/transfer.rs

+    /// State machine for work progression
+    state: Mutex<UploadState>,
+    /// The original request (body taken for processing)
+    request: Arc<UploadInput>,


Will this end up being double Arc, seen from UploadTransfer?

ysaito1001 · 2026-03-10T18:00:13Z

aws-sdk-s3-transfer-manager/src/operation/upload/transfer.rs

+    }
+
+    /// Get the upload_id if MPU was started.
+    pub(crate) fn upload_id(&self) -> Option<String> {


nit: do we always want to clone a returned String, as opposed to Option<&str>?

It's not on a hot path, only for abort, I'm going to leave for now.

ysaito1001 · 2026-03-10T19:57:29Z

aws-sdk-s3-transfer-manager/src/operation/upload.rs

-    distribute_work(&mut mpu_data, ctx)?;
-    Ok(UploadType::MultipartUpload(mpu_data))
-}
+        // TODO: Relax this constraint - unknown content length implies MPU


ysaito1001 · 2026-03-10T20:15:27Z

aws-sdk-s3-transfer-manager/src/operation/upload/transfer.rs

+                        return PollWork::Pending;
+                    }
+                    self.inner.ctx.set_pending();
+                    return PollWork::Pending;


Just to double check, once next_part exceeds total_parts, is returning PollWork::Pending here is expected?

No, in practice this state is never observed I think because we transition to completing before poll_work sees it. I'll update it though to transition to completing as well.

aajtodd added 4 commits March 4, 2026 23:15

move upload transfer over to vnext-scheduler

145dd77

fail harder

677660e

fix comments

11c338e

make upload adhere to state machine contract better

3401389

aajtodd requested a review from a team as a code owner March 5, 2026 19:06

landonxjames approved these changes Mar 10, 2026

View reviewed changes

ysaito1001 approved these changes Mar 10, 2026

View reviewed changes

ysaito1001 reviewed Mar 10, 2026

View reviewed changes

address feedback and issues

1f8669d

aajtodd merged commit 6d0c017 into s3-tm-vnext Mar 11, 2026
10 of 11 checks passed

aajtodd deleted the vnext-upload branch March 11, 2026 17:55

aajtodd mentioned this pull request Mar 18, 2026

migrate download to state machine and new scheduler #133

Merged

aajtodd mentioned this pull request Mar 27, 2026

implement support for adaptive concurrency #135

Open

	let mut state = self.inner.state.lock().unwrap();
	let mut state = self.inner.state.lock().expect("Failed to unlock transfer state");

	// TODO: Relax this constraint - unknown content length implies MPU
	// TODO(vnext): Relax this constraint - unknown content length implies MPU

Conversation

aajtodd commented Mar 5, 2026

Summary

Upload Lifecycle

Multipart upload

Small file (below MPU threshold)

Design

What changed

What did not change

Uh oh!

landonxjames left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ysaito1001 commented Mar 10, 2026

Uh oh!

ysaito1001 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants