perf(mpsc): reduce intermediate copies in bounded receives

CHANGELOG.md

@@ -42,6 +42,7 @@ All notable changes to this project will be documented in this file.
 
 ### Improvements
 
+* Reduce intermediate payload copies in bounded MPSC receives, especially for large inline messages.
 * Reduce the cost of initializing an uncontended `OnceCell` by avoiding its internal wait-queue lock.
 * Reduce the cost of polling a pending oneshot receive again from the same task, which now keeps its registered waker instead of cloning and dropping it.
 * Reduce the state kept by futures waiting for a `Mutex`, `RwLock`, `Semaphore`, `Once`, `OnceCell`, or bounded broadcast capacity, making lock, permit acquisition, and bounded broadcast send futures smaller.

asyncband/src/mpsc/bounded/mod.rs

@@ -24,7 +24,6 @@ use std::task::Waker;
 
 use crate::internal::mutex::Mutex;
 use crate::internal::waitlist::WaitList;
-use crate::mpsc::TryRecvError;
 use crate::mpsc::TrySendError;
 
 mod receiver;
@@ -106,16 +105,6 @@ impl<T> State<T> {
         self.available += 1;
         None
     }
-
-    fn pop(&mut self) -> Result<(T, Option<Waker>), TryRecvError> {
-        if let Some(value) = self.queue.pop_front() {
-            Ok((value, self.release()))
-        } else if self.senders == 0 {
-            Err(TryRecvError::Disconnected)
-        } else {
-            Err(TryRecvError::Empty)
-        }
-    }
 }
 
 struct Waiter {

asyncband/src/mpsc/bounded/receiver.rs

@@ -93,7 +93,14 @@ impl<T> BoundedReceiver<T> {
     /// assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected));
     /// ```
     pub fn try_recv(&mut self) -> Result<T, TryRecvError> {
-        let (value, wake) = self.shared.lock().pop()?;
+        let mut state = self.shared.lock();
+        let value = match state.queue.pop_front() {
+            Some(value) => value,
+            None if state.senders == 0 => return Err(TryRecvError::Disconnected),
+            None => return Err(TryRecvError::Empty),
+        };
+        let wake = state.release();
+        drop(state);
         if let Some(waker) = wake {
             waker.wake();
         }
@@ -180,21 +187,22 @@ impl<T> BoundedReceiver<T> {
 
     fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Result<T, RecvError>> {
         let mut state = self.shared.lock();
-        match state.pop() {
-            Ok((value, wake)) => {
+        match state.queue.pop_front() {
+            Some(value) => {
+                let wake = state.release();
                 drop(state);
                 if let Some(waker) = wake {
                     waker.wake();
                 }
                 Poll::Ready(Ok(value))
             }
-            Err(TryRecvError::Disconnected) => {
+            None if state.senders == 0 => {
                 let old = state.recv_waker.take();
                 drop(state);
                 drop(old);
                 Poll::Ready(Err(RecvError::Disconnected))
             }
-            Err(TryRecvError::Empty) => {
+            None => {
                 let old = register_waker(&mut state.recv_waker, cx.waker());
                 drop(state);
                 drop(old);

benchmarks/README.md

@@ -51,6 +51,18 @@ MPSC task scenarios include both one and four producers, small and large buffers
 3. Put that count in the scenario registration and run all implementations. Neither peers nor later revisions are calibrated independently. Compare the same scenario, work size, toolchain and machine; a change from 300 to 240 ms is a 20% reduction in elapsed time.
 4. If the work size, an input or a timing boundary changes, rerun both sides under the new definition. Treat it as a new baseline. There is no automatic calibration or elapsed-time performance assertion in CI.
 
+## Diagnose bounded MPSC payload movement
+
+The `tasks-4p1c-cap64-4w`, `tasks-4p1c-cap64-4w-inline1k`, and `tasks-4p1c-cap64-4w-boxed1k` scenarios use the same task fixture with different payload representations. Set all three to 26 rounds when investigating [inline payload throughput (#401)](https://github.com/apache/asyncband/issues/401); their registered default round counts differ. Boxed values include allocation and deallocation in the measurement, so that control is not a measurement of queue copying alone.
+
+```shell
+cargo x bench --bench ecosystem -- mpsc/bounded/tasks-4p1c-cap64-4w --rounds 26 --repeat 10 --raw
+cargo x bench --bench ecosystem -- mpsc/bounded/ready-roundtrip-inline1k --repeat 10 --raw
+cargo x bench --bench ecosystem -- mpsc/bounded/ready-roundtrip-boxed1k --repeat 10 --raw
+```
+
+The ready payload cases hand-poll immediately ready sends and receives with the existing reference-counted waker, check the received sequence, and include message construction and validation. They help distinguish payload movement from task scheduling; they do not isolate send time from receive time. Inspect the optimized send and receive code separately or collect stack samples to attribute any difference.
+
 ## Diagnose Asyncband primitives
 
 ```shell

benchmarks/benches/ecosystem/mpsc/bounded.rs

@@ -29,6 +29,7 @@ use crate::mpsc::adapters::Flume;
 use crate::mpsc::adapters::Tokio;
 use crate::mpsc::support::BOUNDED_CAPACITY;
 use crate::mpsc::support::Bounded;
+use crate::mpsc::support::Message;
 use crate::mpsc::support::RepeatedBatch;
 use crate::mpsc::support::RepeatedTasks;
 use crate::mpsc::support::TaskShape;
@@ -52,6 +53,18 @@ fn ready_send_receive<C: BoundedMpsc>(bencher: Bencher) {
     });
 }
 
+fn ready_payload<C: BoundedMpsc<T>, T: Message>(bencher: Bencher) {
+    let mut context = bench_context();
+    let (sender, mut receiver) = C::channel(BOUNDED_CAPACITY);
+    bencher.bench_local(|| {
+        C::send_ready(&sender, black_box(T::new(0)), &mut context);
+        assert_eq!(
+            black_box(C::recv_ready(&mut receiver, &mut context)).sequence(),
+            0
+        );
+    });
+}
+
 fn reused_threads<C: BoundedMpsc>(bencher: Bencher, producer_count: usize) {
     let mut batch = RepeatedBatch::<Bounded<C>>::new(producer_count);
     batch.run();
@@ -67,11 +80,11 @@ fn reused_tasks<C: BoundedMpsc, const CAPACITY: usize>(
     bencher.bench_local(|| batch.run());
 }
 
-fn reused_tasks_inline_1k<C: BoundedMpsc<[u8; 1024]>, const CAPACITY: usize>(
+fn reused_tasks_payload<C: BoundedMpsc<T>, const CAPACITY: usize, T: Message>(
     bencher: Bencher,
     TaskShape { producers, workers }: TaskShape,
 ) {
-    let mut batch = RepeatedTasks::<Bounded<C, CAPACITY, [u8; 1024]>>::new(producers, workers);
+    let mut batch = RepeatedTasks::<Bounded<C, CAPACITY, T>>::new(producers, workers);
     batch.run();
     bencher.bench_local(|| batch.run());
 }
@@ -126,6 +139,20 @@ pub fn register(runner: &mut Runner) {
             ready_send_receive::<C>(b)
         }),
     );
+    runner.compare(
+        "mpsc/bounded/ready-roundtrip-inline1k",
+        1_000_000,
+        implementations!([Asyncband, Tokio, AsyncChannel, Flume], |C, b| {
+            ready_payload::<C, [u8; 1024]>(b)
+        }),
+    );
+    runner.compare(
+        "mpsc/bounded/ready-roundtrip-boxed1k",
+        1_000_000,
+        implementations!([Asyncband, Tokio, AsyncChannel, Flume], |C, b| {
+            ready_payload::<C, Box<[u8; 1024]>>(b)
+        }),
+    );
     runner.compare(
         "mpsc/bounded/threads-4p1c-cap64",
         89,
@@ -144,7 +171,20 @@ pub fn register(runner: &mut Runner) {
         "mpsc/bounded/tasks-4p1c-cap64-4w-inline1k",
         26,
         implementations!([Asyncband, Tokio, AsyncChannel, Flume], |C, b| {
-            reused_tasks_inline_1k::<C, 64>(
+            reused_tasks_payload::<C, 64, [u8; 1024]>(
+                b,
+                TaskShape {
+                    producers: 4,
+                    workers: 4,
+                },
+            )
+        }),
+    );
+    runner.compare(
+        "mpsc/bounded/tasks-4p1c-cap64-4w-boxed1k",
+        26,
+        implementations!([Asyncband, Tokio, AsyncChannel, Flume], |C, b| {
+            reused_tasks_payload::<C, 64, Box<[u8; 1024]>>(
                 b,
                 TaskShape {
                     producers: 4,

benchmarks/benches/ecosystem/mpsc/support.rs

@@ -70,6 +70,16 @@ impl Message for [u8; 1024] {
     }
 }
 
+impl Message for Box<[u8; 1024]> {
+    fn new(sequence: usize) -> Self {
+        Box::new(<[u8; 1024] as Message>::new(sequence))
+    }
+
+    fn sequence(self) -> usize {
+        <[u8; 1024] as Message>::sequence(*self)
+    }
+}
+
 pub trait ConcurrentMpsc: Send + Sync + 'static {
     type Message: Message;
     type Sender: Clone + Send + Sync + 'static;

tests-integration/tests/mpsc_test/backpressure.rs

@@ -16,6 +16,7 @@
 // under the License.
 
 use std::pin::pin;
+use std::task::Poll;
 
 use asyncband::mpsc;
 use tests_integration::WakeCounter;
@@ -110,3 +111,88 @@ fn bounded_receiver_drop_returns_values_to_all_blocked_senders() {
     assert_eq!(first_error.into_inner(), 1);
     assert_eq!(second_error.into_inner(), 2);
 }
+
+#[test]
+fn large_pending_send_preserves_ownership_on_cancel_and_disconnect() {
+    use std::sync::Arc;
+    use std::sync::atomic::AtomicUsize;
+    use std::sync::atomic::Ordering;
+
+    #[derive(Debug)]
+    struct Payload {
+        bytes: [u8; 1024],
+        drops: Arc<AtomicUsize>,
+    }
+
+    impl Drop for Payload {
+        fn drop(&mut self) {
+            self.drops.fetch_add(1, Ordering::Relaxed);
+        }
+    }
+
+    for capacity in [1, 3] {
+        for cancel_after_grant in [false, true] {
+            let drops = Arc::new(AtomicUsize::new(0));
+            let payload = |byte| Payload {
+                bytes: [byte; 1024],
+                drops: drops.clone(),
+            };
+            let (tx, mut rx) = mpsc::bounded(capacity);
+            for _ in 0..capacity {
+                assert!(matches!(
+                    poll_once(pin!(tx.send(payload(1)))),
+                    Poll::Ready(Ok(()))
+                ));
+            }
+            let mut pending = Box::pin(tx.send(payload(2)));
+            assert!(poll_once(pending.as_mut()).is_pending());
+            if cancel_after_grant {
+                assert_eq!(rx.try_recv().unwrap().bytes, [1; 1024]);
+            }
+            drop(pending);
+            assert_eq!(
+                drops.load(Ordering::Relaxed),
+                1 + usize::from(cancel_after_grant)
+            );
+            if !cancel_after_grant {
+                assert_eq!(rx.try_recv().unwrap().bytes, [1; 1024]);
+            }
+            // Cancelling either a queued or granted send returns exactly one slot.
+            tx.try_send(payload(3)).unwrap();
+            assert!(tx.try_reserve().is_err());
+            for _ in 1..capacity {
+                assert_eq!(rx.try_recv().unwrap().bytes, [1; 1024]);
+            }
+            assert_eq!(rx.try_recv().unwrap().bytes, [3; 1024]);
+            assert_eq!(drops.load(Ordering::Relaxed), capacity + 2);
+
+            tx.try_send(payload(4)).unwrap();
+            let mut pending = pin!(tx.send(payload(5)));
+            // Fill any remaining slots so this send must wait.
+            for _ in 1..capacity {
+                tx.try_send(payload(4)).unwrap();
+            }
+            assert!(poll_once(pending.as_mut()).is_pending());
+            drop(rx);
+            let error = match poll_once(pending.as_mut()) {
+                Poll::Ready(Err(error)) => error,
+                other => panic!("expected disconnected send, got {other:?}"),
+            };
+            assert_eq!(error.into_inner().bytes, [5; 1024]);
+            assert_eq!(drops.load(Ordering::Relaxed), capacity * 2 + 3);
+
+            let (tx, mut rx) = mpsc::bounded(1);
+            tx.try_send(payload(6)).unwrap();
+            let mut pending = Box::pin(tx.send(payload(7)));
+            assert!(poll_once(pending.as_mut()).is_pending());
+            assert_eq!(rx.try_recv().unwrap().bytes, [6; 1024]);
+            assert!(matches!(poll_once(pending.as_mut()), Poll::Ready(Ok(()))));
+            drop(pending);
+            drop(tx);
+            let received = expect_ready(poll_once(pin!(rx.recv()))).unwrap();
+            assert_eq!(received.bytes, [7; 1024]);
+            drop(received);
+            assert_eq!(drops.load(Ordering::Relaxed), capacity * 2 + 5);
+        }
+    }
+}