Repository: alexcrichton/futures-rs (rust)

Issue:
MPSC: Send unless blocks
Using a bounded MPSC channel, would it be possible to send conditionally on the condition that the channel won't block?

Something like: `futures::unsync::mpsc::Sender::send_unless_blocks(&mut self, item: Self::SinkItem) -> Result<(), WouldBlock>`.

Here's my motivation for such an API: I want to send a `Task` handler from a stream that's going to be parked, to an outer receiver that's going to unpark it, but I don't want to send it multiple times in the case the stream is polled multiple times: if the stream won't block, I want to sent the `Task` and return an `Async::NotReady`. If the stream would block, that means that the stream isn't unparked yet, so I want just to return an `Async::NotReady`.

Pull request description:
Add `mpsc::Sender::try_send`
This is an API which serves for sending a message along a mpsc::Sender
without attempting to park the current task, making it safe to invoke
from a helper thread or a non-futures-aware thread.

Even though `try_send` doesn't "park" the current task, it still needs
to internally perform the parking logic. This is due to the fact that
parking is really tracking that a given sender instance is full.

Doing this required refactoring the internals to decouple tracking a
parked sender from having a task::Task value to notify. After this
patch, `SenderTask` has a `is_parked` field which is the authority when
it comes to determining if a sender is blocked or not. The task field
contains the task, if any, that should be notified on unpark.

Closes #478

Patch:
diff --git a/src/sync/mpsc/mod.rs b/src/sync/mpsc/mod.rs
index e8e370ea..ebf98539 100644
--- a/src/sync/mpsc/mod.rs
+++ b/src/sync/mpsc/mod.rs
@@ -96,7 +96,7 @@ pub struct Sender<T> {
     // Handle to the task that is blocked on this sender. This handle is sent
     // to the receiver half in order to be notified when the sender becomes
     // unblocked.
-    sender_task: SenderTask,
+    sender_task: Arc<Mutex<SenderTask>>,
 
     // True if the sender might be blocked. This is an optimization to avoid
     // having to lock the mutex most of the time.
@@ -136,6 +136,18 @@ pub struct UnboundedReceiver<T>(Receiver<T>);
 #[derive(Clone, PartialEq, Eq)]
 pub struct SendError<T>(T);
 
+/// Error type returned from `try_send`
+#[derive(Clone, PartialEq, Eq)]
+pub struct TrySendError<T> {
+    kind: TrySendErrorKind<T>,
+}
+
+#[derive(Clone, PartialEq, Eq)]
+enum TrySendErrorKind<T> {
+    Full(T),
+    Disconnected(T),
+}
+
 impl<T> fmt::Debug for SendError<T> {
     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
         fmt.debug_tuple("SendError")
@@ -164,6 +176,65 @@ impl<T> SendError<T> {
     }
 }
 
+impl<T> fmt::Debug for TrySendError<T> {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+        fmt.debug_tuple("TrySendError")
+            .field(&"...")
+            .finish()
+    }
+}
+
+impl<T> fmt::Display for TrySendError<T> {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+        if self.is_full() {
+            write!(fmt, "send failed because channel is full")
+        } else {
+            write!(fmt, "send failed because receiver is gone")
+        }
+    }
+}
+
+impl<T: Any> Error for TrySendError<T> {
+    fn description(&self) -> &str {
+        if self.is_full() {
+            "send failed because channel is full"
+        } else {
+            "send failed because receiver is gone"
+        }
+    }
+}
+
+impl<T> TrySendError<T> {
+    /// Returns true if this error is a result of the channel being full
+    pub fn is_full(&self) -> bool {
+        use self::TrySendErrorKind::*;
+
+        match self.kind {
+            Full(_) => true,
+            _ => false,
+        }
+    }
+
+    /// Returns true if this error is a result of the receiver being dropped
+    pub fn is_disconnected(&self) -> bool {
+        use self::TrySendErrorKind::*;
+
+        match self.kind {
+            Disconnected(_) => true,
+            _ => false,
+        }
+    }
+
+    /// Returns the message that was attempted to be sent but failed.
+    pub fn into_inner(self) -> T {
+        use self::TrySendErrorKind::*;
+
+        match self.kind {
+            Full(v) | Disconnected(v) => v,
+        }
+    }
+}
+
 #[derive(Debug)]
 struct Inner<T> {
     // Max buffer size of the channel. If `None` then the channel is unbounded.
@@ -177,7 +248,7 @@ struct Inner<T> {
     message_queue: Queue<Option<T>>,
 
     // Atomic, FIFO queue used to send parked task handles to the receiver.
-    parked_queue: Queue<SenderTask>,
+    parked_queue: Queue<Arc<Mutex<SenderTask>>>,
 
     // Number of senders in existence
     num_senders: AtomicUsize,
@@ -224,7 +295,28 @@ const MAX_CAPACITY: usize = !(OPEN_MASK);
 const MAX_BUFFER: usize = MAX_CAPACITY >> 1;
 
 // Sent to the consumer to wake up blocked producers
-type SenderTask = Arc<Mutex<Option<Task>>>;
+#[derive(Debug)]
+struct SenderTask {
+    task: Option<Task>,
+    is_parked: bool,
+}
+
+impl SenderTask {
+    fn new() -> Self {
+        SenderTask {
+            task: None,
+            is_parked: false,
+        }
+    }
+
+    fn notify(&mut self) {
+        self.is_parked = false;
+
+        if let Some(task) = self.task.take() {
+            task.notify();
+        }
+    }
+}
 
 /// Creates an in-memory channel implementation of the `Stream` trait with
 /// bounded capacity.
@@ -278,7 +370,7 @@ fn channel2<T>(buffer: Option<usize>) -> (Sender<T>, Receiver<T>) {
 
     let tx = Sender {
         inner: inner.clone(),
-        sender_task: Arc::new(Mutex::new(None)),
+        sender_task: Arc::new(Mutex::new(SenderTask::new())),
         maybe_parked: false,
     };
 
@@ -296,10 +388,35 @@ fn channel2<T>(buffer: Option<usize>) -> (Sender<T>, Receiver<T>) {
  */
 
 impl<T> Sender<T> {
+    /// Attempts to send a message on this `Sender<T>` without blocking.
+    ///
+    /// This function, unlike `start_send`, is safe to call whether it's being
+    /// called on a task or not. Note that this function, however, will *not*
+    /// attempt to block the current task if the message cannot be sent.
+    ///
+    /// It is not recommended to call this function from inside of a future,
+    /// only from an external thread where you've otherwise arranged to be
+    /// notified when the channel is no longer full.
+    pub fn try_send(&mut self, msg: T) -> Result<(), TrySendError<T>> {
+        // If the sender is currently blocked, reject the message
+        if !self.poll_unparked(false).is_ready() {
+            return Err(TrySendError {
+                kind: TrySendErrorKind::Full(msg),
+            });
+        }
+
+        // The channel has capacity to accept the message, so send it
+        self.do_send(Some(msg), false)
+            .map_err(|SendError(v)| {
+                TrySendError {
+                    kind: TrySendErrorKind::Disconnected(v),
+                }
+            })
+    }
+
     // Do the send without failing
     // None means close
-    fn do_send(&mut self, msg: Option<T>) -> Result<(), SendError<T>> {
-
+    fn do_send(&mut self, msg: Option<T>, do_park: bool) -> Result<(), SendError<T>> {
         // First, increment the number of messages contained by the channel.
         // This operation will also atomically determine if the sender task
         // should be parked.
@@ -334,7 +451,7 @@ impl<T> Sender<T> {
         // maintain internal consistency, a blank message is pushed onto the
         // parked task queue.
         if park_self {
-            self.park(msg.is_some());
+            self.park(do_park);
         }
 
         self.queue_push_and_signal(msg);
@@ -447,7 +564,11 @@ impl<T> Sender<T> {
             None
         };
 
-        *self.sender_task.lock().unwrap() = task;
+        {
+            let mut sender = self.sender_task.lock().unwrap();
+            sender.task = task;
+            sender.is_parked = true;
+        }
 
         // Send handle over queue
         let t = self.sender_task.clone();
@@ -475,17 +596,17 @@ impl<T> Sender<T> {
             return Err(SendError(()));
         }
 
-        Ok(self.poll_unparked())
+        Ok(self.poll_unparked(true))
     }
 
-    fn poll_unparked(&mut self) -> Async<()> {
+    fn poll_unparked(&mut self, do_park: bool) -> Async<()> {
         // First check the `maybe_parked` variable. This avoids acquiring the
         // lock in most cases
         if self.maybe_parked {
             // Get a lock on the task handle
             let mut task = self.sender_task.lock().unwrap();
 
-            if task.is_none() {
+            if !task.is_parked {
                 self.maybe_parked = false;
                 return Async::Ready(())
             }
@@ -496,7 +617,11 @@ impl<T> Sender<T> {
             //
             // Update the task in case the `Sender` has been moved to another
             // task
-            *task = Some(task::current());
+            task.task = if do_park {
+                Some(task::current())
+            } else {
+                None
+            };
 
             Async::NotReady
         } else {
@@ -512,12 +637,12 @@ impl<T> Sink for Sender<T> {
     fn start_send(&mut self, msg: T) -> StartSend<T, SendError<T>> {
         // If the sender is currently blocked, reject the message before doing
         // any work.
-        if !self.poll_unparked().is_ready() {
+        if !self.poll_unparked(true).is_ready() {
             return Ok(AsyncSink::NotReady(msg));
         }
 
         // The channel has capacity to accept the message, so send it.
-        self.do_send(Some(msg))?;
+        self.do_send(Some(msg), true)?;
 
         Ok(AsyncSink::Ready)
     }
@@ -618,7 +743,7 @@ impl<T> Clone for Sender<T> {
             if actual == curr {
                 return Sender {
                     inner: self.inner.clone(),
-                    sender_task: Arc::new(Mutex::new(None)),
+                    sender_task: Arc::new(Mutex::new(SenderTask::new())),
                     maybe_parked: false,
                 };
             }
@@ -634,7 +759,7 @@ impl<T> Drop for Sender<T> {
         let prev = self.inner.num_senders.fetch_sub(1, SeqCst);
 
         if prev == 1 {
-            let _ = self.do_send(None);
+            let _ = self.do_send(None, false);
         }
     }
 }
@@ -674,10 +799,7 @@ impl<T> Receiver<T> {
         loop {
             match unsafe { self.inner.parked_queue.pop() } {
                 PopResult::Data(task) => {
-                    let task = task.lock().unwrap().take();
-                    if let Some(task) = task {
-                        task.notify();
-                    }
+                    task.lock().unwrap().notify();
                 }
                 PopResult::Empty => break,
                 PopResult::Inconsistent => thread::yield_now(),
@@ -719,14 +841,7 @@ impl<T> Receiver<T> {
         loop {
             match unsafe { self.inner.parked_queue.pop() } {
                 PopResult::Data(task) => {
-                    // Do this step first so that the lock is dropped when
-                    // `unpark` is called
-                    let task = task.lock().unwrap().take();
-
-                    if let Some(task) = task {
-                        task.notify();
-                    }
-
+                    task.lock().unwrap().notify();
                     return;
                 }
                 PopResult::Empty => {


Test patch:
diff --git a/tests/mpsc.rs b/tests/mpsc.rs
index f8fe40c5..b2cf5d12 100644
--- a/tests/mpsc.rs
+++ b/tests/mpsc.rs
@@ -420,3 +420,72 @@ fn is_ready<T>(res: &AsyncSink<T>) -> bool {
         _ => false,
     }
 }
+
+#[test]
+fn try_send_1() {
+    const N: usize = 3000;
+    let (mut tx, rx) = mpsc::channel(0);
+
+    let t = thread::spawn(move || {
+        for i in 0..N {
+            loop {
+                if tx.try_send(i).is_ok() {
+                    break
+                }
+            }
+        }
+    });
+    for (i, j) in rx.wait().enumerate() {
+        assert_eq!(i, j.unwrap());
+    }
+    t.join().unwrap();
+}
+
+#[test]
+fn try_send_2() {
+    use std::thread;
+    use std::time::Duration;
+
+    let (mut tx, rx) = mpsc::channel(0);
+
+    tx.try_send("hello").unwrap();
+
+    let th = thread::spawn(|| {
+        lazy(|| {
+            assert!(tx.start_send("fail").unwrap().is_not_ready());
+            Ok::<_, ()>(())
+        }).wait().unwrap();
+
+        tx.send("goodbye").wait().unwrap();
+    });
+
+    // Little sleep to hopefully get the action on the thread to happen first
+    thread::sleep(Duration::from_millis(300));
+
+    let mut rx = rx.wait();
+
+    assert_eq!(rx.next(), Some(Ok("hello")));
+    assert_eq!(rx.next(), Some(Ok("goodbye")));
+    assert!(rx.next().is_none());
+
+    th.join().unwrap();
+}
+
+#[test]
+fn try_send_fail() {
+    let (mut tx, rx) = mpsc::channel(0);
+    let mut rx = rx.wait();
+
+    tx.try_send("hello").unwrap();
+
+    // This should fail
+    assert!(tx.try_send("fail").is_err());
+
+    assert_eq!(rx.next(), Some(Ok("hello")));
+
+    tx.try_send("goodbye").unwrap();
+    drop(tx);
+
+    assert_eq!(rx.next(), Some(Ok("goodbye")));
+    assert!(rx.next().is_none());
+}