Only return Poll::Ready from Sender if the channel is empty

Author: thomaseizinger

futures-channel/src/mpsc/mod.rs

@@ -616,6 +616,25 @@ impl<T> BoundedSenderInner<T> {
         self.poll_unparked(Some(cx)).map(Ok)
     }
 
+    /// Polls the channel to determine if it is empty.
+    ///
+    /// # Return value
+    ///
+    /// This method returns:
+    ///
+    /// - `Poll::Ready(()` if there are no messages in the channel;
+    /// - `Poll::Pending` if there are messages in the channel.
+    fn poll_is_empty(&mut self, cx: &mut Context<'_>) -> Poll<()> {
+        let state = decode_state(self.inner.state.load(SeqCst));
+        if state.num_messages == 0 {
+            return Poll::Ready(());
+        }
+
+        // If there are messages in the channel, we must park the task unconditionally.
+        self.sender_task.lock().unwrap().task = Some(cx.waker().clone());
+        return Poll::Pending;
+    }
+
     /// Returns whether the senders send to the same receiver.
     fn same_receiver(&self, other: &Self) -> bool {
         Arc::ptr_eq(&self.inner, &other.inner)
@@ -755,6 +774,23 @@ impl<T> Sender<T> {
         let ptr = self.0.as_ref().map(|inner| inner.ptr());
         ptr.hash(hasher);
     }
+
+    /// Polls the channel to determine if it is empty.
+    ///
+    /// # Return value
+    ///
+    /// This method returns:
+    ///
+    /// - `Poll::Ready(()` if there are no messages in the channel or the [`Receiver`] is disconnected.
+    /// - `Poll::Pending` if there are messages in the channel.
+    pub(crate) fn poll_is_empty(&mut self, cx: &mut Context<'_>) -> Poll<()> {
+        let inner = match self.0.as_mut() {
+            None => return Poll::Ready(()),
+            Some(inner) => inner,
+        };
+
+        inner.poll_is_empty(cx)
+    }
 }
 
 impl<T> UnboundedSender<T> {

futures-channel/src/mpsc/sink_impl.rs

@@ -15,13 +15,7 @@ impl<T> Sink<T> for Sender<T> {
     }
 
     fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
-        match (*self).poll_ready(cx) {
-            Poll::Ready(Err(ref e)) if e.is_disconnected() => {
-                // If the receiver disconnected, we consider the sink to be flushed.
-                Poll::Ready(Ok(()))
-            }
-            x => x,
-        }
+        (*self).poll_is_empty(cx).map(Ok)
     }
 
     fn poll_close(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {

futures-channel/tests/mpsc.rs

@@ -60,6 +60,37 @@ fn send_recv_no_buffer() {
     }));
 }
 
+#[test]
+fn sink_poll_flush() {
+    // Run on a task context
+    block_on(poll_fn(move |cx| {
+        let (tx, rx) = mpsc::channel::<i32>(2);
+        pin_mut!(tx, rx);
+
+        assert!(tx.as_mut().poll_flush(cx).is_ready());
+        assert!(tx.as_mut().poll_ready(cx).is_ready());
+
+        // Send two messages, `poll_flush` should be pending after each of them.
+        assert!(tx.as_mut().start_send(1).is_ok());
+        assert!(tx.as_mut().poll_flush(cx).is_pending());
+
+        assert!(tx.as_mut().start_send(2).is_ok());
+        assert!(tx.as_mut().poll_flush(cx).is_pending());
+
+        // Take first message
+        assert_eq!(rx.as_mut().poll_next(cx), Poll::Ready(Some(1)));
+        assert!(tx.as_mut().poll_ready(cx).is_ready());
+        assert!(tx.as_mut().poll_flush(cx).is_pending());
+
+        // Take second message
+        assert_eq!(rx.as_mut().poll_next(cx), Poll::Ready(Some(2)));
+        assert!(tx.as_mut().poll_ready(cx).is_ready());
+        assert!(tx.as_mut().poll_flush(cx).is_ready());
+
+        Poll::Ready(())
+    }));
+}
+
 #[test]
 fn send_shared_recv() {
     let (mut tx1, rx) = mpsc::channel::<i32>(16);