feat: Add a way to reserve many permits on bounded mpsc channel

tokio/src/sync/batch_semaphore.rs

@@ -71,7 +71,7 @@ pub struct AcquireError(());
 pub(crate) struct Acquire<'a> {
     node: Waiter,
     semaphore: &'a Semaphore,
-    num_permits: u32,
+    num_permits: usize,
     queued: bool,
 }
 
@@ -293,7 +293,7 @@ impl Semaphore {
         }
     }
 
-    pub(crate) fn acquire(&self, num_permits: u32) -> Acquire<'_> {
+    pub(crate) fn acquire(&self, num_permits: usize) -> Acquire<'_> {
         Acquire::new(self, num_permits)
     }
 
@@ -371,7 +371,7 @@ impl Semaphore {
     fn poll_acquire(
         &self,
         cx: &mut Context<'_>,
-        num_permits: u32,
+        num_permits: usize,
         node: Pin<&mut Waiter>,
         queued: bool,
     ) -> Poll<Result<(), AcquireError>> {
@@ -380,7 +380,7 @@ impl Semaphore {
         let needed = if queued {
             node.state.load(Acquire) << Self::PERMIT_SHIFT
         } else {
-            (num_permits as usize) << Self::PERMIT_SHIFT
+            num_permits << Self::PERMIT_SHIFT
         };
 
         let mut lock = None;
@@ -506,12 +506,12 @@ impl fmt::Debug for Semaphore {
 
 impl Waiter {
     fn new(
-        num_permits: u32,
+        num_permits: usize,
         #[cfg(all(tokio_unstable, feature = "tracing"))] ctx: trace::AsyncOpTracingCtx,
     ) -> Self {
         Waiter {
             waker: UnsafeCell::new(None),
-            state: AtomicUsize::new(num_permits as usize),
+            state: AtomicUsize::new(num_permits),
             pointers: linked_list::Pointers::new(),
             #[cfg(all(tokio_unstable, feature = "tracing"))]
             ctx,
@@ -591,7 +591,7 @@ impl Future for Acquire<'_> {
 }
 
 impl<'a> Acquire<'a> {
-    fn new(semaphore: &'a Semaphore, num_permits: u32) -> Self {
+    fn new(semaphore: &'a Semaphore, num_permits: usize) -> Self {
         #[cfg(any(not(tokio_unstable), not(feature = "tracing")))]
         return Self {
             node: Waiter::new(num_permits),
@@ -635,7 +635,7 @@ impl<'a> Acquire<'a> {
         });
     }
 
-    fn project(self: Pin<&mut Self>) -> (Pin<&mut Waiter>, &Semaphore, u32, &mut bool) {
+    fn project(self: Pin<&mut Self>) -> (Pin<&mut Waiter>, &Semaphore, usize, &mut bool) {
         fn is_unpin<T: Unpin>() {}
         unsafe {
             // Safety: all fields other than `node` are `Unpin`
@@ -673,7 +673,7 @@ impl Drop for Acquire<'_> {
         // Safety: we have locked the wait list.
         unsafe { waiters.queue.remove(node) };
 
-        let acquired_permits = self.num_permits as usize - self.node.state.load(Acquire);
+        let acquired_permits = self.num_permits - self.node.state.load(Acquire);
         if acquired_permits > 0 {
             self.semaphore.add_permits_locked(acquired_permits, waiters);
         }

tokio/src/sync/mpsc/bounded.rs

@@ -68,6 +68,18 @@ pub struct Permit<'a, T> {
     chan: &'a chan::Tx<T, Semaphore>,
 }
 
+/// An [`Iterator`] of [`Permit`] that can be used to reserve `n` slots in the channel.
+///
+/// `PermitIterator` values are returned by [`Sender::reserve_many()`] and [`Sender::try_reserve_many()`]
+/// and are used to guarantee channel capacity before generating `n` message to send.
+///
+/// [`Sender::reserve_many()`]: Sender::reserve_many
+/// [`Sender::try_reserve_many()`]: Sender::try_reserve_many
+pub struct PermitIterator<'a, T> {
+    chan: &'a chan::Tx<T, Semaphore>,
+    n: usize,
+}
+
 /// Owned permit to send one value into the channel.
 ///
 /// This is identical to the [`Permit`] type, except that it moves the sender
@@ -849,10 +861,68 @@ impl<T> Sender<T> {
     /// }
     /// ```
     pub async fn reserve(&self) -> Result<Permit<'_, T>, SendError<()>> {
-        self.reserve_inner().await?;
+        self.reserve_inner(1).await?;
         Ok(Permit { chan: &self.chan })
     }
 
+    /// Waits for channel capacity. Once capacity to send `n` messages is
+    /// available, it is reserved for the caller.
+    ///
+    /// If the channel is full, the function waits for the number of unreceived
+    /// messages to become `n` less than the channel capacity. Capacity to send `n`
+    /// message is reserved for the caller. A [`PermitIterator`] is returned to track
+    /// the reserved capacity. You can call this [`Iterator`] until it is exhausted to
+    /// get a [`Permit`] and then call [`Permit::send`].
+    ///
+    /// Dropping [`PermitIterator`] without sending all messages releases the capacity back
+    /// to the channel.
+    ///
+    /// [`PermitIterator`]: PermitIterator
+    /// [`Permit`]: Permit
+    /// [`send`]: Permit::send
+    ///
+    /// # Cancel safety
+    ///
+    /// This channel uses a queue to ensure that calls to `send` and `reserve`
+    /// complete in the order they were requested. Cancelling a call to
+    /// `reserve` makes you lose your place in the queue.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use tokio::sync::mpsc;
+    ///
+    /// #[tokio::main]
+    /// async fn main() {
+    ///     let (tx, mut rx) = mpsc::channel(2);
+    ///
+    ///     // Reserve capacity
+    ///     let mut permit = tx.reserve_many(2).await.unwrap();
+    ///
+    ///     // Trying to send directly on the `tx` will fail due to no
+    ///     // available capacity.
+    ///     assert!(tx.try_send(123).is_err());
+    ///
+    ///     // Sending with the permit iterator succeeds
+    ///     permit.next().unwrap().send(456);
+    ///     permit.next().unwrap().send(457);
+    ///
+    ///     // The iterator should now be exhausted
+    ///     assert!(permit.next().is_none());
+    ///     
+    ///     // The value sent on the permit is received
+    ///     assert_eq!(rx.recv().await.unwrap(), 456);
+    ///     assert_eq!(rx.recv().await.unwrap(), 457);
+    /// }
+    /// ```
+    pub async fn reserve_many(&self, n: usize) -> Result<PermitIterator<'_, T>, SendError<()>> {
+        self.reserve_inner(n).await?;
+        Ok(PermitIterator {
+            chan: &self.chan,
+            n,
+        })
+    }
+
     /// Waits for channel capacity, moving the `Sender` and returning an owned
     /// permit. Once capacity to send one message is available, it is reserved
     /// for the caller.
@@ -934,16 +1004,16 @@ impl<T> Sender<T> {
     /// [`send`]: OwnedPermit::send
     /// [`Arc::clone`]: std::sync::Arc::clone
     pub async fn reserve_owned(self) -> Result<OwnedPermit<T>, SendError<()>> {
-        self.reserve_inner().await?;
+        self.reserve_inner(1).await?;
         Ok(OwnedPermit {
             chan: Some(self.chan),
         })
     }
 
-    async fn reserve_inner(&self) -> Result<(), SendError<()>> {
+    async fn reserve_inner(&self, n: usize) -> Result<(), SendError<()>> {
         crate::trace::async_trace_leaf().await;
 
-        match self.chan.semaphore().semaphore.acquire(1).await {
+        match self.chan.semaphore().semaphore.acquire(n).await {
             Ok(()) => Ok(()),
             Err(_) => Err(SendError(())),
         }
@@ -1002,6 +1072,68 @@ impl<T> Sender<T> {
         Ok(Permit { chan: &self.chan })
     }
 
+    /// Tries to acquire `n` slot in the channel without waiting for the slot to become
+    /// available.
+    ///
+    /// If the channel is full this function will return a [`TrySendError`], otherwise
+    /// A [`PermitIterator`] is returned to track the reserved capacity.
+    /// You can call this [`Iterator`] until it is exhausted to
+    /// get a [`Permit`] and then call [`Permit::send`]. This function is similar to
+    /// [`reserve_many`] except it does not await for the slot to become available.
+    ///
+    /// Dropping [`PermitIterator`] without sending a message releases the capacity back
+    /// to the channel.
+    ///
+    /// [`PermitIterator`]: PermitIterator
+    /// [`send`]: Permit::send
+    /// [`reserve_many`]: Sender::reserve_many
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use tokio::sync::mpsc;
+    ///
+    /// #[tokio::main]
+    /// async fn main() {
+    ///     let (tx, mut rx) = mpsc::channel(2);
+    ///
+    ///     // Reserve capacity
+    ///     let mut permit = tx.try_reserve_many(2).unwrap();
+    ///
+    ///     // Trying to send directly on the `tx` will fail due to no
+    ///     // available capacity.
+    ///     assert!(tx.try_send(123).is_err());
+    ///
+    ///     // Trying to reserve an additional slot on the `tx` will
+    ///     // fail because there is no capacity.
+    ///     assert!(tx.try_reserve().is_err());
+    ///
+    ///     // Sending with the permit iterator succeeds
+    ///     permit.next().unwrap().send(456);
+    ///     permit.next().unwrap().send(457);
+    ///
+    ///     // The iterator should now be exhausted
+    ///     assert!(permit.next().is_none());
+    ///
+    ///     // The value sent on the permit is received
+    ///     assert_eq!(rx.recv().await.unwrap(), 456);
+    ///     assert_eq!(rx.recv().await.unwrap(), 457);
+    ///
+    /// }
+    /// ```
+    pub fn try_reserve_many(&self, n: usize) -> Result<PermitIterator<'_, T>, TrySendError<()>> {
+        match self.chan.semaphore().semaphore.try_acquire(n as u32) {
+            Ok(()) => {}
+            Err(TryAcquireError::Closed) => return Err(TrySendError::Closed(())),
+            Err(TryAcquireError::NoPermits) => return Err(TrySendError::Full(())),
+        }
+
+        Ok(PermitIterator {
+            chan: &self.chan,
+            n,
+        })
+    }
+
     /// Tries to acquire a slot in the channel without waiting for the slot to become
     /// available, returning an owned permit.
     ///
@@ -1278,6 +1410,53 @@ impl<T> fmt::Debug for Permit<'_, T> {
     }
 }
 
+// ===== impl PermitIterator =====
+
+impl<'a, T> Iterator for PermitIterator<'a, T> {
+    type Item = Permit<'a, T>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.n == 0 {
+            return None;
+        }
+
+        self.n -= 1;
+        Some(Permit { chan: self.chan })
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let n = self.n;
+        (n, Some(n))
+    }
+}
+impl<T> ExactSizeIterator for PermitIterator<'_, T> {}
+
+impl<T> Drop for PermitIterator<'_, T> {
+    fn drop(&mut self) {
+        use chan::Semaphore;
+
+        let semaphore = self.chan.semaphore();
+
+        // Add the remaining permits back to the semaphore
+        semaphore.add_permits(self.n);
+
+        // If this is the last sender for this channel, wake the receiver so
+        // that it can be notified that the channel is closed.
+        if semaphore.is_closed() && semaphore.is_idle() {
+            self.chan.wake_rx();
+        }
+    }
+}
+
+impl<T> fmt::Debug for PermitIterator<'_, T> {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        fmt.debug_struct("PermitIterator")
+            .field("chan", &self.chan)
+            .field("capacity", &self.n)
+            .finish()
+    }
+}
+
 // ===== impl Permit =====
 
 impl<T> OwnedPermit<T> {

tokio/src/sync/mpsc/mod.rs

@@ -95,7 +95,9 @@
 pub(super) mod block;
 
 mod bounded;
-pub use self::bounded::{channel, OwnedPermit, Permit, Receiver, Sender, WeakSender};
+pub use self::bounded::{
+    channel, OwnedPermit, Permit, PermitIterator, Receiver, Sender, WeakSender,
+};
 
 mod chan;
 

tokio/src/sync/rwlock.rs

@@ -772,7 +772,7 @@ impl<T: ?Sized> RwLock<T> {
     /// ```
     pub async fn write(&self) -> RwLockWriteGuard<'_, T> {
         let acquire_fut = async {
-            self.s.acquire(self.mr).await.unwrap_or_else(|_| {
+            self.s.acquire(self.mr as usize).await.unwrap_or_else(|_| {
                 // The semaphore was closed. but, we never explicitly close it, and we have a
                 // handle to it through the Arc, which means that this can never happen.
                 unreachable!()
@@ -907,7 +907,7 @@ impl<T: ?Sized> RwLock<T> {
         let resource_span = self.resource_span.clone();
 
         let acquire_fut = async {
-            self.s.acquire(self.mr).await.unwrap_or_else(|_| {
+            self.s.acquire(self.mr as usize).await.unwrap_or_else(|_| {
                 // The semaphore was closed. but, we never explicitly close it, and we have a
                 // handle to it through the Arc, which means that this can never happen.
                 unreachable!()

tokio/src/sync/semaphore.rs

@@ -565,7 +565,7 @@ impl Semaphore {
     pub async fn acquire_many(&self, n: u32) -> Result<SemaphorePermit<'_>, AcquireError> {
         #[cfg(all(tokio_unstable, feature = "tracing"))]
         trace::async_op(
-            || self.ll_sem.acquire(n),
+            || self.ll_sem.acquire(n as usize),
             self.resource_span.clone(),
             "Semaphore::acquire_many",
             "poll",
@@ -574,7 +574,7 @@ impl Semaphore {
         .await?;
 
         #[cfg(not(all(tokio_unstable, feature = "tracing")))]
-        self.ll_sem.acquire(n).await?;
+        self.ll_sem.acquire(n as usize).await?;
 
         Ok(SemaphorePermit {
             sem: self,
@@ -764,14 +764,14 @@ impl Semaphore {
     ) -> Result<OwnedSemaphorePermit, AcquireError> {
         #[cfg(all(tokio_unstable, feature = "tracing"))]
         let inner = trace::async_op(
-            || self.ll_sem.acquire(n),
+            || self.ll_sem.acquire(n as usize),
             self.resource_span.clone(),
             "Semaphore::acquire_many_owned",
             "poll",
             true,
         );
         #[cfg(not(all(tokio_unstable, feature = "tracing")))]
-        let inner = self.ll_sem.acquire(n);
+        let inner = self.ll_sem.acquire(n as usize);
 
         inner.await?;
         Ok(OwnedSemaphorePermit {