batch submissions to the runtime optimistically

Author: aajtodd

aws-sdk-s3-transfer-manager/src/runtime/managed.rs

@@ -24,6 +24,7 @@ use tokio_util::sync::CancellationToken;
 use super::topology::Cpu;
 use super::Topology;
 use super::{ExecutionRuntime, RuntimeComponents, ScheduledWork};
+use crate::runtime::sync::SubmissionGuard;
 use crate::scheduler::Scheduler;
 use crate::transfer::{TransferId, WorkOutcome};
 
@@ -314,26 +315,28 @@ impl ManagedThreadRuntimeBuilder {
 }
 
 impl ExecutionRuntime for ManagedThreadRuntime {
-    fn dispatch(&self, mut work: ScheduledWork) {
-        let thread_id = self.router.select(&self.threads);
-        let th = &self.threads[thread_id.0];
-        th.in_flight.fetch_add(1, Ordering::Relaxed);
-
-        let scheduler = self.scheduler.clone();
-        let in_flight = Arc::clone(&th.in_flight);
-
-        th.runtime_handle.spawn(async move {
-            let result = execute_work(&mut work, &scheduler).await;
-            in_flight.fetch_sub(1, Ordering::Relaxed);
-            match result {
-                ExecuteResult::Completed(outcome, elapsed) => {
-                    scheduler.on_completion(work, outcome, elapsed);
+    fn dispatch(&self, batch: &mut SubmissionGuard<'_, ScheduledWork>) {
+        for mut work in batch.drain() {
+            let thread_id = self.router.select(&self.threads);
+            let th = &self.threads[thread_id.0];
+            th.in_flight.fetch_add(1, Ordering::Relaxed);
+
+            let scheduler = self.scheduler.clone();
+            let in_flight = Arc::clone(&th.in_flight);
+
+            th.runtime_handle.spawn(async move {
+                let result = execute_work(&mut work, &scheduler).await;
+                in_flight.fetch_sub(1, Ordering::Relaxed);
+                match result {
+                    ExecuteResult::Completed(outcome, elapsed) => {
+                        scheduler.on_completion(work, outcome, elapsed);
+                    }
+                    ExecuteResult::Panicked => {
+                        scheduler.on_panic(work);
+                    }
                 }
-                ExecuteResult::Panicked => {
-                    scheduler.on_panic(work);
-                }
-            }
-        });
+            });
+        }
     }
 
     fn shutdown(&self) {

aws-sdk-s3-transfer-manager/src/runtime/mod.rs

@@ -16,10 +16,11 @@ pub(crate) use managed::ManagedThreadRuntime;
 mod topology;
 pub(crate) use topology::{Cpu, NumaNode, Topology};
 
-mod sync;
+pub(crate) mod sync;
 
 use aws_smithy_runtime_api::client::http::SharedHttpClient;
 
+use crate::runtime::sync::SubmissionGuard;
 use crate::scheduler::descriptor::TransferDescriptor;
 use crate::transfer::{IoRequest, TransferId};
 
@@ -41,8 +42,8 @@ pub(crate) struct ScheduledWork {
 ///
 /// The scheduler decides WHAT to run and WHEN. The runtime decides WHERE and HOW.
 pub(crate) trait ExecutionRuntime: Send + Sync + std::fmt::Debug {
-    /// Dispatch an IO request for execution.
-    fn dispatch(&self, work: ScheduledWork);
+    /// Dispatch a batch of IO requests for execution.
+    fn dispatch(&self, batch: &mut SubmissionGuard<'_, ScheduledWork>);
 
     /// Shut down the runtime, draining in-flight work.
     fn shutdown(&self);

aws-sdk-s3-transfer-manager/src/runtime/sync/mod.rs

@@ -6,4 +6,4 @@
 //! Synchronization primitives for the runtime.
 
 mod submission;
-pub(crate) use submission::{SubmissionGuard, SubmissionQueue};
+pub(crate) use submission::{Submission, SubmissionGuard, SubmissionQueue};

aws-sdk-s3-transfer-manager/src/runtime/tokio_mt.rs

@@ -15,6 +15,7 @@ use futures_util::FutureExt;
 mod worker_pool;
 
 use super::{ExecutionRuntime, RuntimeComponents, ScheduledWork};
+use crate::runtime::sync::SubmissionGuard;
 use crate::scheduler::Scheduler;
 use crate::transfer::{TransferId, WorkOutcome};
 use worker_pool::WorkerPool;
@@ -79,10 +80,12 @@ impl TokioMultiThreadRuntime {
 }
 
 impl ExecutionRuntime for TokioMultiThreadRuntime {
-    fn dispatch(&self, work: ScheduledWork) {
+    fn dispatch(&self, batch: &mut SubmissionGuard<'_, ScheduledWork>) {
         self.ensure_workers_started();
         self.ensure_worker_capacity();
-        self.pool.push(work);
+        for work in batch.drain() {
+            self.pool.push(work);
+        }
     }
 
     fn shutdown(&self) {

aws-sdk-s3-transfer-manager/src/scheduler/scheduler.rs

@@ -56,6 +56,7 @@ use crate::transfer::{BoxTransfer, IoRequest, PollWork, TransferId, WorkOutcome}
 use crate::metrics::{IOCounters, IoSample};
 use crate::runtime::ExecutionRuntime;
 
+use crate::runtime::sync::{Submission, SubmissionQueue};
 use crate::runtime::ScheduledWork;
 use crate::scheduler::descriptor::TransferDescriptor;
 use crate::scheduler::ready_set::ReadySet;
@@ -66,6 +67,9 @@ use std::sync::Arc;
 use std::sync::{OnceLock, RwLock};
 use std::time::Duration;
 
+/// Batch size for work generated and submitted in a single round
+const SUBMISSION_QUEUE_SIZE: usize = 64;
+
 /// Event-driven scheduler for coordinating transfer work.
 ///
 /// Clone is cheap (Arc).
@@ -79,6 +83,7 @@ struct SchedulerInner {
     io_counters: Arc<IOCounters>,
     runtime: OnceLock<Arc<dyn ExecutionRuntime>>,
     dispatched: AtomicUsize,
+    submission_queue: SubmissionQueue<ScheduledWork>,
 }
 
 impl std::fmt::Debug for Scheduler {
@@ -115,6 +120,7 @@ impl SchedulerBuilder {
             io_counters: self.io_counters,
             runtime: OnceLock::new(),
             dispatched: AtomicUsize::new(0),
+            submission_queue: SubmissionQueue::new(SUBMISSION_QUEUE_SIZE),
         }));
         let runtime = runtime_factory(scheduler.clone());
         scheduler
@@ -301,7 +307,7 @@ impl Scheduler {
                 item: IoRequest { kind, data },
                 descriptor: desc.clone(),
             };
-            self.dispatch_to_runtime(next);
+            self.dispatch_single(next);
         }
 
         // capacity has freed try to queue up more work
@@ -333,8 +339,20 @@ impl Scheduler {
         self.0.dispatched.load(Ordering::Relaxed) < self.0.controller.target()
     }
 
+    /// Flush the current submission and re-enter for the next batch.
+    fn submit_and_reenter<'a>(
+        &'a self,
+        sub: Submission<'a, ScheduledWork>,
+    ) -> Submission<'a, ScheduledWork> {
+        if let Some(mut guard) = sub.submit() {
+            self.runtime().dispatch(&mut guard);
+        }
+        self.0.submission_queue.enter()
+    }
+
     /// Generate work from ready transfers and dispatch to runtime.
     fn generate_work(&self) {
+        let mut sub = self.0.submission_queue.enter();
         while self.has_capacity() {
             let Some(desc) = self.0.ready_set.pop() else {
                 break;
@@ -349,10 +367,16 @@ impl Scheduler {
                 PollWork::Ready(item) => {
                     desc.work_generated();
                     self.0.ready_set.insert(desc.clone());
-                    self.dispatch_to_runtime(ScheduledWork {
+                    self.0.dispatched.fetch_add(1, Ordering::Relaxed);
+                    desc.work_queued();
+                    let work = ScheduledWork {
                         item,
                         descriptor: desc,
-                    });
+                    };
+                    if let Err(work) = sub.push(work) {
+                        sub = self.submit_and_reenter(sub);
+                        sub.push(work).expect("empty queue after flush");
+                    }
                 }
                 PollWork::Pending => {
                     // re-added on wake as state machine progresses
@@ -363,12 +387,24 @@ impl Scheduler {
                 }
             }
         }
+        if let Some(mut guard) = sub.submit() {
+            self.runtime().dispatch(&mut guard);
+        }
     }
 
-    fn dispatch_to_runtime(&self, work: ScheduledWork) {
+    fn dispatch_single(&self, work: ScheduledWork) {
         self.0.dispatched.fetch_add(1, Ordering::Relaxed);
         work.descriptor.work_queued();
-        self.runtime().dispatch(work);
+        let sub = self.0.submission_queue.enter();
+        if let Err(work) = sub.push(work) {
+            let sub = self.submit_and_reenter(sub);
+            sub.push(work).expect("empty queue after flush");
+            if let Some(mut guard) = sub.submit() {
+                self.runtime().dispatch(&mut guard);
+            }
+        } else if let Some(mut guard) = sub.submit() {
+            self.runtime().dispatch(&mut guard);
+        }
     }
 
     #[allow(dead_code)] // TODO: wire into Handle for graceful shutdown