Repository: alexcrichton/futures-rs
Pull request #575 (2017-09-09 23:39:54)

Add sync::mpsc::Sender::poll_ready
This is useful a API for situations where it's advantageous to check for
channel capacity before creating or procuring an item to send.

danburkert: Failed due to using struct initialization shorthand syntax in the test; squashed a fix into the stress test commit.
danburkert: > Thanks! Could this also have a clause in the documentation that due to the multiple nature of senders that the return value does not guarantee that the next call to start_send will succeed?

This is actually guaranteed; `poll_ready` will only indicate ready when the sender's reserved slot is vacant.  Without this guarantee the method would be much less useful (at least for the use-case I have in mind for it).  The attached stress test verifies that the guarantee holds by way of asserting on `start_send` returning ready.
carllerche: What @danburkert says is correct. If `poll_ready` returns `Ready`, then the next call to `start_send` is guaranteed to succeed **unless** the receive half closed.
alexcrichton: re `()` vs `SendError<()>` -- my main hesitation with using `()` as an error is that elsewhere in `futures` we use `()` as "this can never happen", which while we want to change that may cause confusion if there's multiple definitions of it.

re the guarantees here, is this something we want to commit to long-term? Such a guarantee may preclude future optimizations and/or implementations? Is it really necessary to provide such a guarantee?
carllerche: @alexcrichton I think returning `SendError<()>` is good (I wasn't arguing against it.

Re: making the guarantee, I think that it should. Being able to provide this guarantee is a pretty fundamental part of the design of the mpsc channel. Also, being able to rely on the guarantee is a requirement for providing sane back pressure support at the `Service` layer.

I think it's fine to have this structure provide the guarantee and if there is an alternate design that can provide different trade-offs, it can be added as a separate structure.
alexcrichton: Ok sounds like we should leave the documentation guarantee and @danburkert mind changing to return `SendError<()>`?

FWIW The `()` in `SendError<()>` I agree isn't exactly ideal, but I'm hoping `SendError<??>` is better than `()`

Patch:
diff --git a/src/sync/mpsc/mod.rs b/src/sync/mpsc/mod.rs
index 00bc4012..e8e370ea 100644
--- a/src/sync/mpsc/mod.rs
+++ b/src/sync/mpsc/mod.rs
@@ -39,7 +39,7 @@
 // Since most of this work is lock-free, once the work starts, it is impossible
 // to safely revert.
 //
-// If the sender is unable to process a send operation, then the the curren
+// If the sender is unable to process a send operation, then the the current
 // task is parked and the handle is sent on the parked task queue.
 //
 // Note that the implementation guarantees that the channel capacity will never
@@ -459,6 +459,25 @@ impl<T> Sender<T> {
         self.maybe_parked = state.is_open;
     }
 
+    /// Polls the channel to determine if there is guaranteed to be capacity to send at least one
+    /// item without waiting.
+    ///
+    /// Returns `Ok(Async::Ready(_))` if there is sufficient capacity, or returns
+    /// `Ok(Async::NotReady)` if the channel is not guaranteed to have capacity. Returns
+    /// `Err(SendError(_))` if the receiver has been dropped.
+    ///
+    /// # Panics
+    ///
+    /// This method will panic if called from outside the context of a task or future.
+    pub fn poll_ready(&mut self) -> Poll<(), SendError<()>> {
+        let state = decode_state(self.inner.state.load(SeqCst));
+        if !state.is_open {
+            return Err(SendError(()));
+        }
+
+        Ok(self.poll_unparked())
+    }
+
     fn poll_unparked(&mut self) -> Async<()> {
         // First check the `maybe_parked` variable. This avoids acquiring the
         // lock in most cases


Test patch:
diff --git a/tests/mpsc.rs b/tests/mpsc.rs
index ecd84109..f8fe40c5 100644
--- a/tests/mpsc.rs
+++ b/tests/mpsc.rs
@@ -1,5 +1,6 @@
 #![cfg(feature = "use_std")]
 
+#[macro_use]
 extern crate futures;
 
 use futures::prelude::*;
@@ -37,11 +38,12 @@ fn send_recv_no_buffer() {
     // Run on a task context
     lazy(move || {
         assert!(tx.poll_complete().unwrap().is_ready());
+        assert!(tx.poll_ready().unwrap().is_ready());
 
         // Send first message
-
         let res = tx.start_send(1).unwrap();
         assert!(is_ready(&res));
+        assert!(tx.poll_ready().unwrap().is_not_ready());
 
         // Send second message
         let res = tx.start_send(2).unwrap();
@@ -49,12 +51,15 @@ fn send_recv_no_buffer() {
 
         // Take the value
         assert_eq!(rx.poll().unwrap(), Async::Ready(Some(1)));
+        assert!(tx.poll_ready().unwrap().is_ready());
 
         let res = tx.start_send(2).unwrap();
         assert!(is_ready(&res));
+        assert!(tx.poll_ready().unwrap().is_not_ready());
 
         // Take the value
         assert_eq!(rx.poll().unwrap(), Async::Ready(Some(2)));
+        assert!(tx.poll_ready().unwrap().is_ready());
 
         Ok::<(), ()>(())
     }).wait().unwrap();
@@ -106,13 +111,14 @@ fn send_recv_threads_no_capacity() {
 
 #[test]
 fn recv_close_gets_none() {
-    let (tx, mut rx) = mpsc::channel::<i32>(10);
+    let (mut tx, mut rx) = mpsc::channel::<i32>(10);
 
     // Run on a task context
     lazy(move || {
         rx.close();
 
         assert_eq!(rx.poll(), Ok(Async::Ready(None)));
+        assert!(tx.poll_ready().is_err());
 
         drop(tx);
 
@@ -348,6 +354,66 @@ fn stress_close_receiver() {
     }
 }
 
+/// Tests that after `poll_ready` indicates capacity a channel can always send without waiting.
+#[test]
+fn stress_poll_ready() {
+    // A task which checks channel capcity using poll_ready, and pushes items onto the channel when
+    // ready.
+    struct SenderTask {
+        sender: mpsc::Sender<u32>,
+        count: u32,
+    }
+    impl Future for SenderTask {
+        type Item = ();
+        type Error = ();
+        fn poll(&mut self) -> Poll<(), ()> {
+            // In a loop, check if the channel is ready. If so, push an item onto the channel
+            // (asserting that it doesn't attempt to block).
+            while self.count > 0 {
+                try_ready!(self.sender.poll_ready().map_err(|_| ()));
+                assert!(self.sender.start_send(self.count).unwrap().is_ready());
+                self.count -= 1;
+            }
+            Ok(Async::Ready(()))
+        }
+    }
+
+    const AMT: u32 = 1000;
+    const NTHREADS: u32 = 8;
+
+    /// Run a stress test using the specified channel capacity.
+    fn stress(capacity: usize) {
+        let (tx, rx) = mpsc::channel(capacity);
+        let mut threads = Vec::new();
+        for _ in 0..NTHREADS {
+            let sender = tx.clone();
+            threads.push(thread::spawn(move || {
+                SenderTask {
+                    sender: sender,
+                    count: AMT,
+                }.wait()
+            }));
+        }
+        drop(tx);
+
+        let mut rx = rx.wait();
+        for _ in 0..AMT * NTHREADS {
+            assert!(rx.next().is_some());
+        }
+
+        assert!(rx.next().is_none());
+
+        for thread in threads {
+            thread.join().unwrap().unwrap();
+        }
+    }
+
+    stress(0);
+    stress(1);
+    stress(8);
+    stress(16);
+}
+
 fn is_ready<T>(res: &AsyncSink<T>) -> bool {
     match *res {
         AsyncSink::Ready => true,