refactor(mpsc): replace the legacy queue backend

CHANGELOG.md

@@ -4,6 +4,14 @@ All notable changes to this project will be documented in this file.
 
 ## Unreleased
 
+### Bug fixes
+
+* Release MPSC receiver wakers when the receiver is dropped, avoiding retained tasks and ownership cycles when a waker holds a sender.
+
+### Improvements
+
+* Reduce unbounded MPSC synchronization overhead by transferring messages in batches and coordinating receiver notifications with queued messages; release large empty batch allocations while retaining small buffers for reuse.
+
 ## v0.7.2
 
 ### Improvements

asyncband/src/internal/atomic_waker.rs

@@ -210,8 +210,10 @@ impl AtomicWaker {
         }
     }
 
+    /// Removes the registered waker if this call acquires the slot. A concurrent registration or
+    /// wake may instead take responsibility for notifying it.
     #[inline]
-    fn take(&self) -> Option<Waker> {
+    pub fn take(&self) -> Option<Waker> {
         // ORDERING: When this reads WAITING, Acquire receives the registered waker published by the
         // previous owner. Release publishes the condition update that the caller performed before
         // calling wake, including when a registering thread already owns the slot.

asyncband/src/internal/cache_padded.rs

@@ -0,0 +1,55 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+// Use conservative architecture estimates, not a guarantee about every CPU's cache line.
+// Keep 128 bytes for large ARM/PowerPC lines and adjacent-line prefetching on x86-64,
+// 256 bytes for s390x, and at least 64 bytes elsewhere.
+#[cfg_attr(target_arch = "s390x", repr(align(256)))]
+#[cfg_attr(
+    any(
+        target_arch = "aarch64",
+        target_arch = "arm64ec",
+        target_arch = "powerpc64",
+        target_arch = "x86_64",
+    ),
+    repr(align(128))
+)]
+#[cfg_attr(
+    not(any(
+        target_arch = "s390x",
+        target_arch = "aarch64",
+        target_arch = "arm64ec",
+        target_arch = "powerpc64",
+        target_arch = "x86_64",
+    )),
+    repr(align(64))
+)]
+pub struct CachePadded<T>(T);
+
+impl<T> CachePadded<T> {
+    pub const fn new(value: T) -> Self {
+        Self(value)
+    }
+}
+
+impl<T> std::ops::Deref for CachePadded<T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}

asyncband/src/internal/mod.rs

@@ -70,6 +70,9 @@ pub(crate) mod atomic_waker;
 #[allow(dead_code)]
 pub(crate) mod arena;
 
+#[cfg(feature = "mpsc")]
+pub(crate) mod cache_padded;
+
 #[cfg(any(feature = "latch", feature = "once"))]
 pub(crate) mod countdown;
 
@@ -92,6 +95,8 @@ pub(crate) mod value_cell;
     feature = "waitgroup",
     feature = "watch",
 ))]
+// Some primitives use only shared access, leaving `Mutex::get_mut` unused.
+#[allow(dead_code)]
 pub(crate) mod mutex;
 
 #[cfg(any(

asyncband/src/internal/mutex.rs

@@ -34,4 +34,8 @@ impl<T> Mutex<T> {
     pub fn lock(&self) -> std::sync::MutexGuard<'_, T> {
         self.0.lock().unwrap_or_else(PoisonError::into_inner)
     }
+
+    pub fn get_mut(&mut self) -> &mut T {
+        self.0.get_mut().unwrap_or_else(PoisonError::into_inner)
+    }
 }

asyncband/src/mpsc/bounded.rs

@@ -19,15 +19,15 @@
 //! tasks with backpressure control.
 
 use std::fmt;
-use std::future::Future;
 use std::future::poll_fn;
-use std::pin::pin;
 use std::sync::Arc;
 use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering;
 use std::task::Context;
 use std::task::Poll;
+use std::task::ready;
 
+use self::ring::Ring;
 use super::RecvError;
 use super::SendError;
 use super::TryRecvError;
@@ -36,6 +36,8 @@ use crate::internal::atomic_waker::AtomicWaker;
 use crate::internal::semaphore::Acquire;
 use crate::internal::semaphore::Semaphore;
 
+mod ring;
+
 /// Creates a bounded mpsc channel with room for `buffer` queued messages.
 ///
 /// [`BoundedSender::send`] waits for capacity when the buffer is full. Receiving a message releases
@@ -48,42 +50,38 @@ use crate::internal::semaphore::Semaphore;
 pub fn bounded<T>(buffer: usize) -> (BoundedSender<T>, BoundedReceiver<T>) {
     assert!(buffer > 0, "mpsc bounded channel requires buffer > 0");
     let state = Arc::new(BoundedState {
+        buffer: Ring::new(buffer),
         senders: AtomicUsize::new(1),
-        tx_permits: Semaphore::new(0),
+        send_waiters: Semaphore::new(0),
         rx_waker: AtomicWaker::new(),
     });
-    let (sender, receiver) = std::sync::mpsc::sync_channel(buffer);
     let sender = BoundedSender {
         state: state.clone(),
-        sender: Some(sender),
-    };
-    let receiver = BoundedReceiver {
-        state: state.clone(),
-        receiver: Some(receiver),
     };
+    let receiver = BoundedReceiver { state };
     (sender, receiver)
 }
 
-struct BoundedState {
+struct BoundedState<T> {
+    buffer: Ring<T>,
     senders: AtomicUsize,
-    tx_permits: Semaphore,
+    // Notifications grant retries; only the ring determines whether buffer capacity is available.
+    send_waiters: Semaphore,
     rx_waker: AtomicWaker,
 }
 
 /// The sending endpoint of a bounded mpsc channel.
 ///
 /// Instances are created by the [`bounded`] function.
 pub struct BoundedSender<T> {
-    state: Arc<BoundedState>,
-    sender: Option<std::sync::mpsc::SyncSender<T>>,
+    state: Arc<BoundedState<T>>,
 }
 
 impl<T> Clone for BoundedSender<T> {
     fn clone(&self) -> Self {
         self.state.senders.fetch_add(1, Ordering::Release);
         BoundedSender {
             state: self.state.clone(),
-            sender: self.sender.clone(),
         }
     }
 }
@@ -96,9 +94,6 @@ impl<T> fmt::Debug for BoundedSender<T> {
 
 impl<T> Drop for BoundedSender<T> {
     fn drop(&mut self) {
-        // Dropping the final underlying sender disconnects the channel.
-        drop(self.sender.take());
-
         match self.state.senders.fetch_sub(1, Ordering::AcqRel) {
             1 => {
                 // Wake the receiver so it can observe the channel's disconnected state.
@@ -142,7 +137,7 @@ impl<T> BoundedSender<T> {
                 };
 
                 loop {
-                    let poll = pin!(&mut self.acquire).poll(cx);
+                    let poll = self.acquire.poll_once(cx.waker());
 
                     value = match self.sender.try_send(value) {
                         Ok(()) => return Poll::Ready(Ok(())),
@@ -153,7 +148,7 @@ impl<T> BoundedSender<T> {
                     };
 
                     if poll.is_ready() {
-                        self.acquire = self.sender.state.tx_permits.poll_acquire(1);
+                        self.acquire = self.sender.state.send_waiters.poll_acquire(1);
                     } else {
                         self.value = Some(value);
                         return Poll::Pending;
@@ -162,7 +157,7 @@ impl<T> BoundedSender<T> {
             }
         }
 
-        let acquire = self.state.tx_permits.poll_acquire(1);
+        let acquire = self.state.send_waiters.poll_acquire(1);
         let mut send = SendState {
             sender: self,
             value: Some(value),
@@ -192,34 +187,19 @@ impl<T> BoundedSender<T> {
     /// assert_eq!(tx.try_send(30), Err(TrySendError::Disconnected(30)));
     /// ```
     pub fn try_send(&self, value: T) -> Result<(), TrySendError<T>> {
-        // INVARIANT: A shared borrow of the endpoint cannot overlap its destructor.
-        let sender = self.sender.as_ref().unwrap();
-        match sender.try_send(value) {
-            Ok(()) => {
-                self.state.rx_waker.wake();
-
-                Ok(())
-            }
-            Err(std::sync::mpsc::TrySendError::Full(value)) => Err(TrySendError::Full(value)),
-            Err(std::sync::mpsc::TrySendError::Disconnected(value)) => {
-                Err(TrySendError::Disconnected(value))
-            }
-        }
+        self.state.buffer.try_push(value)?;
+        self.state.rx_waker.wake();
+        Ok(())
     }
 }
 
 /// The receiving endpoint of a bounded mpsc channel.
 ///
 /// Instances are created by the [`bounded`] function.
 pub struct BoundedReceiver<T> {
-    state: Arc<BoundedState>,
-    receiver: Option<std::sync::mpsc::Receiver<T>>,
+    state: Arc<BoundedState<T>>,
 }
 
-/// The only `!Sync` field `receiver` is protected by `&mut self` in `recv` and `try_recv`.
-/// That is, `BoundedReceiver` can only be accessed by one thread at a time.
-unsafe impl<T: Send> Sync for BoundedReceiver<T> {}
-
 impl<T> fmt::Debug for BoundedReceiver<T> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("BoundedReceiver").finish_non_exhaustive()
@@ -228,13 +208,20 @@ impl<T> fmt::Debug for BoundedReceiver<T> {
 
 impl<T> Drop for BoundedReceiver<T> {
     fn drop(&mut self) {
-        drop(self.receiver.take());
-        self.state.tx_permits.notify_all();
+        // A registered waker may own a sender; release it to break that ownership cycle.
+        let receiver_waker = self.state.rx_waker.take();
+        // SAFETY: Only this non-cloneable receiver consumes the queue, through exclusive borrows.
+        unsafe { self.state.buffer.disconnect_receiver() };
+        self.state.send_waiters.notify_all();
+        drop(receiver_waker);
     }
 }
 
 impl<T> BoundedReceiver<T> {
-    /// Attempts to receive the next queued value without waiting.
+    /// Attempts to receive the next queued value without waiting for a new message.
+    ///
+    /// A producer already publishing a queued message may delay this call until publication
+    /// finishes. Use [`Self::recv`] to yield asynchronously while publication is in progress.
     ///
     /// Receiving a value frees one buffer slot. An empty channel returns [`TryRecvError::Empty`]
     /// while at least one sender remains, or [`TryRecvError::Disconnected`] after every sender has
@@ -257,18 +244,33 @@ impl<T> BoundedReceiver<T> {
     /// assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected));
     /// ```
     pub fn try_recv(&mut self) -> Result<T, TryRecvError> {
-        // INVARIANT: A mutable borrow of the endpoint cannot overlap its destructor.
-        let receiver = self.receiver.as_ref().unwrap();
-        match receiver.try_recv() {
-            Ok(v) => {
-                self.state.tx_permits.release_if_nonempty(1);
-                Ok(v)
+        loop {
+            if let Poll::Ready(result) = self.try_recv_once() {
+                return result;
             }
-            Err(std::sync::mpsc::TryRecvError::Disconnected) => Err(TryRecvError::Disconnected),
-            Err(std::sync::mpsc::TryRecvError::Empty) => Err(TryRecvError::Empty),
+            std::thread::yield_now();
         }
     }
 
+    fn try_recv_once(&mut self) -> Poll<Result<T, TryRecvError>> {
+        // SAFETY: Only this non-cloneable receiver consumes the queue, through exclusive borrows.
+        let value = if let Some(value) = ready!(unsafe { self.state.buffer.pop() }) {
+            value
+        } else if self.state.senders.load(Ordering::Acquire) == 0 {
+            // The final sender can enqueue between the first empty observation and decrementing
+            // the sender count, so check the queue again before reporting disconnection.
+            // SAFETY: The exclusive receiver borrow still guarantees a single consumer.
+            let Some(value) = ready!(unsafe { self.state.buffer.pop() }) else {
+                return Poll::Ready(Err(TryRecvError::Disconnected));
+            };
+            value
+        } else {
+            return Poll::Ready(Err(TryRecvError::Empty));
+        };
+        self.state.send_waiters.release_if_nonempty(1);
+        Poll::Ready(Ok(value))
+    }
+
     /// Waits for and receives the next value, freeing one buffer slot.
     ///
     /// If no value is queued, this method waits until a sender adds one or the last sender is
@@ -303,16 +305,20 @@ impl<T> BoundedReceiver<T> {
     }
 
     fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Result<T, RecvError>> {
-        match self.try_recv() {
-            Ok(v) => Poll::Ready(Ok(v)),
-            Err(TryRecvError::Disconnected) => Poll::Ready(Err(RecvError::Disconnected)),
-            Err(TryRecvError::Empty) => {
+        match self.try_recv_once() {
+            Poll::Ready(Ok(v)) => Poll::Ready(Ok(v)),
+            Poll::Ready(Err(TryRecvError::Disconnected)) => {
+                Poll::Ready(Err(RecvError::Disconnected))
+            }
+            Poll::Pending | Poll::Ready(Err(TryRecvError::Empty)) => {
                 self.state.rx_waker.register(cx.waker());
 
-                match self.try_recv() {
-                    Ok(v) => Poll::Ready(Ok(v)),
-                    Err(TryRecvError::Disconnected) => Poll::Ready(Err(RecvError::Disconnected)),
-                    Err(TryRecvError::Empty) => Poll::Pending,
+                match self.try_recv_once() {
+                    Poll::Ready(Ok(v)) => Poll::Ready(Ok(v)),
+                    Poll::Ready(Err(TryRecvError::Disconnected)) => {
+                        Poll::Ready(Err(RecvError::Disconnected))
+                    }
+                    Poll::Pending | Poll::Ready(Err(TryRecvError::Empty)) => Poll::Pending,
                 }
             }
         }