diff --git a/tokio/src/sync/mpsc/bounded.rs b/tokio/src/sync/mpsc/bounded.rs index c11d2696c..38a9c34f8 100644 --- a/tokio/src/sync/mpsc/bounded.rs +++ b/tokio/src/sync/mpsc/bounded.rs @@ -1275,8 +1275,37 @@ impl Sender { if n > self.max_capacity() { return Err(SendError(())); } - match self.chan.semaphore().semaphore.acquire(n).await { - Ok(()) => Ok(()), + + // If the reservation fails or is cancelled while it holds some but not + // all of the permits it asked for, the inner `Acquire` returns those + // permits to the semaphore on drop, without the receiver notification + // that `Permit` and `PermitIterator` do on release. Wake the receiver in + // that case so a `recv()` blocked on the channel becoming idle can + // observe a close. On success the caller receives the permits and takes + // over that job, so the guard is forgotten below. + struct WakeReceiverOnDrop<'a, T> { + chan: &'a chan::Tx, + } + + impl Drop for WakeReceiverOnDrop<'_, T> { + fn drop(&mut self) { + use chan::Semaphore; + + let semaphore = self.chan.semaphore(); + if semaphore.is_closed() && semaphore.is_idle() { + self.chan.wake_rx(); + } + } + } + + let guard = WakeReceiverOnDrop { chan: &self.chan }; + let result = self.chan.semaphore().semaphore.acquire(n).await; + + match result { + Ok(()) => { + std::mem::forget(guard); + Ok(()) + } Err(_) => Err(SendError(())), } } diff --git a/tokio/tests/sync_mpsc.rs b/tokio/tests/sync_mpsc.rs index f91601313..398ab6339 100644 --- a/tokio/tests/sync_mpsc.rs +++ b/tokio/tests/sync_mpsc.rs @@ -519,6 +519,110 @@ async fn recv_close_gets_none_reserved() { assert!(rx.recv().await.is_none()); } +#[test] +fn failed_reserve_many_wakes_closed_receiver() { + let (tx, mut rx) = mpsc::channel::<()>(2); + + // Hold one of the two slots so the reservation can only take the other one + // and has to queue for the second. + let permit = tx.try_reserve().unwrap(); + + let mut reserve = tokio_test::task::spawn(tx.reserve_many(2)); + assert_pending!(reserve.poll()); + + rx.close(); + + let mut recv = tokio_test::task::spawn(rx.recv()); + assert_pending!(recv.poll()); + + // One slot is still outstanding, so the channel is not idle and the + // receiver must not be woken yet. + drop(permit); + assert!(!recv.is_woken()); + + // The reservation fails on the closed channel and returns the slot it had + // already acquired, leaving the channel idle. The receiver must be woken. + assert_ready_err!(reserve.poll()); + assert!(recv.is_woken()); + assert_ready!(recv.poll()); +} + +#[test] +fn cancelled_reserve_many_wakes_closed_receiver() { + let (tx, mut rx) = mpsc::channel::<()>(2); + + let permit = tx.try_reserve().unwrap(); + + let mut reserve = tokio_test::task::spawn(tx.reserve_many(2)); + assert_pending!(reserve.poll()); + + rx.close(); + + let mut recv = tokio_test::task::spawn(rx.recv()); + assert_pending!(recv.poll()); + + drop(permit); + assert!(!recv.is_woken()); + + // Cancelling the reservation returns the slot it had already acquired, + // leaving the channel idle. The receiver must be woken. + drop(reserve); + assert!(recv.is_woken()); + assert_ready!(recv.poll()); +} + +#[test] +fn failed_reserve_wakes_closed_receiver() { + let (tx, mut rx) = mpsc::channel::<()>(1); + + // Hold the only slot so the reservation below has to queue. + let permit = tx.try_reserve().unwrap(); + + let mut reserve = tokio_test::task::spawn(tx.reserve()); + assert_pending!(reserve.poll()); + + // Returning the slot assigns it to the queued reservation, which is woken + // but not yet polled, so the channel stays non-idle. + drop(permit); + assert!(reserve.is_woken()); + + rx.close(); + + let mut recv = tokio_test::task::spawn(rx.recv()); + assert_pending!(recv.poll()); + + // The reservation fails on the closed channel despite holding the + // assigned slot; returning it leaves the channel idle. The receiver must + // be woken. + assert_ready_err!(reserve.poll()); + assert!(recv.is_woken()); + assert_ready!(recv.poll()); +} + +#[test] +fn cancelled_reserve_wakes_closed_receiver() { + let (tx, mut rx) = mpsc::channel::<()>(1); + + let permit = tx.try_reserve().unwrap(); + + let mut reserve = tokio_test::task::spawn(tx.reserve()); + assert_pending!(reserve.poll()); + + drop(permit); + assert!(reserve.is_woken()); + + rx.close(); + + let mut recv = tokio_test::task::spawn(rx.recv()); + assert_pending!(recv.poll()); + + // Cancelling the reservation returns the slot it was assigned, leaving + // the channel idle. The receiver must be woken. + drop(reserve); + assert!(recv.is_woken()); + assert_ready!(recv.poll()); +} + #[maybe_tokio_test] async fn tx_close_gets_none() { let (_, mut rx) = mpsc::channel::(10);