From 7fb1698e8d8fa9d4ce295b63a17c461b3a40dddd Mon Sep 17 00:00:00 2001 From: Jon Gjengset Date: Thu, 2 Apr 2020 11:27:37 -0400 Subject: [PATCH] sync: Add disarm to mpsc::Sender (#2358) Fixes #898. --- tokio/src/sync/mpsc/bounded.rs | 81 +++++++++++++++++++++++++++++++++- tokio/src/sync/mpsc/chan.rs | 11 +++++ tokio/tests/sync_mpsc.rs | 38 ++++++++++++++++ 3 files changed, 129 insertions(+), 1 deletion(-) diff --git a/tokio/src/sync/mpsc/bounded.rs b/tokio/src/sync/mpsc/bounded.rs index ba89fdd72..f83c05322 100644 --- a/tokio/src/sync/mpsc/bounded.rs +++ b/tokio/src/sync/mpsc/bounded.rs @@ -196,11 +196,90 @@ impl Sender { Sender { chan } } - #[doc(hidden)] // TODO: document + /// Returns `Poll::Ready(Ok(()))` when the channel is able to accept another item. + /// + /// If the channel is full, then `Poll::Pending` is returned and the task is notified when a + /// slot becomes available. + /// + /// Once `poll_ready` returns `Poll::Ready(Ok(()))`, a call to `try_send` will succeed unless + /// the channel has since been closed. To provide this guarantee, the channel reserves one slot + /// in the channel for the coming send. This reserved slot is not available to other `Sender` + /// instances, so you need to be careful to not end up with deadlocks by blocking after calling + /// `poll_ready` but before sending an element. + /// + /// If, after `poll_ready` succeeds, you decide you do not wish to send an item after all, you + /// can use [`disarm`](Sender::disarm) to release the reserved slot. + /// + /// Until an item is sent or [`disarm`](Sender::disarm) is called, repeated calls to + /// `poll_ready` will return either `Poll::Ready(Ok(()))` or `Poll::Ready(Err(_))` if channel + /// is closed. pub fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { self.chan.poll_ready(cx).map_err(|_| ClosedError::new()) } + /// Undo a successful call to `poll_ready`. + /// + /// Once a call to `poll_ready` returns `Poll::Ready(Ok(()))`, it holds up one slot in the + /// channel to make room for the coming send. `disarm` allows you to give up that slot if you + /// decide you do not wish to send an item after all. After calling `disarm`, you must call + /// `poll_ready` until it returns `Poll::Ready(Ok(()))` before attempting to send again. + /// + /// Returns `false` if no slot is reserved for this sender (usually because `poll_ready` was + /// not previously called, or did not succeed). + /// + /// # Motivation + /// + /// Since `poll_ready` takes up one of the finite number of slots in a bounded channel, callers + /// need to send an item shortly after `poll_ready` succeeds. If they do not, idle senders may + /// take up all the slots of the channel, and prevent active senders from getting any requests + /// through. Consider this code that forwards from one channel to another: + /// + /// ```rust,ignore + /// loop { + /// ready!(tx.poll_ready(cx))?; + /// if let Some(item) = ready!(rx.poll_recv(cx)) { + /// tx.try_send(item)?; + /// } else { + /// break; + /// } + /// } + /// ``` + /// + /// If many such forwarders exist, and they all forward into a single (cloned) `Sender`, then + /// any number of forwarders may be waiting for `rx.poll_recv` at the same time. While they do, + /// they are effectively each reducing the channel's capacity by 1. If enough of these + /// forwarders are idle, forwarders whose `rx` _do_ have elements will be unable to find a spot + /// for them through `poll_ready`, and the system will deadlock. + /// + /// `disarm` solves this problem by allowing you to give up the reserved slot if you find that + /// you have to block. We can then fix the code above by writing: + /// + /// ```rust,ignore + /// loop { + /// ready!(tx.poll_ready(cx))?; + /// let item = rx.poll_recv(cx); + /// if let Poll::Ready(Ok(_)) = item { + /// // we're going to send the item below, so don't disarm + /// } else { + /// // give up our send slot, we won't need it for a while + /// tx.disarm(); + /// } + /// if let Some(item) = ready!(item) { + /// tx.try_send(item)?; + /// } else { + /// break; + /// } + /// } + /// ``` + pub fn disarm(&mut self) -> bool { + if self.chan.is_ready() { + self.chan.disarm(); + true + } else { + false + } + } + /// Attempts to immediately send a message on this `Sender` /// /// This method differs from [`send`] by returning immediately if the channel's diff --git a/tokio/src/sync/mpsc/chan.rs b/tokio/src/sync/mpsc/chan.rs index 326286610..346639578 100644 --- a/tokio/src/sync/mpsc/chan.rs +++ b/tokio/src/sync/mpsc/chan.rs @@ -190,12 +190,23 @@ where self.inner.semaphore.poll_acquire(cx, &mut self.permit) } + pub(crate) fn disarm(&mut self) { + // TODO: should this error if not acquired? + self.inner.semaphore.drop_permit(&mut self.permit) + } + /// Send a message and notify the receiver. pub(crate) fn try_send(&mut self, value: T) -> Result<(), (T, TrySendError)> { self.inner.try_send(value, &mut self.permit) } } +impl Tx { + pub(crate) fn is_ready(&self) -> bool { + self.permit.is_acquired() + } +} + impl Tx { pub(crate) fn send_unbounded(&self, value: T) -> Result<(), (T, TrySendError)> { self.inner.try_send(value, &mut ()) diff --git a/tokio/tests/sync_mpsc.rs b/tokio/tests/sync_mpsc.rs index 7e5c60e2e..f02d90aa5 100644 --- a/tokio/tests/sync_mpsc.rs +++ b/tokio/tests/sync_mpsc.rs @@ -40,6 +40,44 @@ fn send_recv_with_buffer() { assert!(val.is_none()); } +#[test] +fn disarm() { + let (tx, rx) = mpsc::channel::(2); + let mut tx1 = task::spawn(tx.clone()); + let mut tx2 = task::spawn(tx.clone()); + let mut tx3 = task::spawn(tx.clone()); + let mut tx4 = task::spawn(tx); + let mut rx = task::spawn(rx); + + // We should be able to `poll_ready` two handles without problem + assert_ready_ok!(tx1.enter(|cx, mut tx| tx.poll_ready(cx))); + assert_ready_ok!(tx2.enter(|cx, mut tx| tx.poll_ready(cx))); + + // But a third should not be ready + assert_pending!(tx3.enter(|cx, mut tx| tx.poll_ready(cx))); + + // Using one of the reserved slots should allow a new handle to become ready + tx1.try_send(1).unwrap(); + // We also need to receive for the slot to be free + let _ = assert_ready!(rx.enter(|cx, mut rx| rx.poll_recv(cx))).unwrap(); + // Now there's a free slot! + assert_ready_ok!(tx3.enter(|cx, mut tx| tx.poll_ready(cx))); + assert_pending!(tx4.enter(|cx, mut tx| tx.poll_ready(cx))); + + // Dropping a ready handle should also open up a slot + drop(tx2); + assert_ready_ok!(tx4.enter(|cx, mut tx| tx.poll_ready(cx))); + assert_pending!(tx1.enter(|cx, mut tx| tx.poll_ready(cx))); + + // Explicitly disarming a handle should also open a slot + assert!(tx3.disarm()); + assert_ready_ok!(tx1.enter(|cx, mut tx| tx.poll_ready(cx))); + + // Disarming a non-armed sender does not free up a slot + assert!(!tx3.disarm()); + assert_pending!(tx3.enter(|cx, mut tx| tx.poll_ready(cx))); +} + #[tokio::test] async fn send_recv_stream_with_buffer() { use tokio::stream::StreamExt;