Merge pull request #11 from Restioson/async_stream

Add stream impl to receiver

Author: zesterer

Cargo.toml

@@ -12,11 +12,12 @@ readme = "README.md"
 
 [features]
 select = []
-async = []
+async = ["futures"]
 default = ["select", "async"]
 
 [dependencies]
 spin = "0.5"
+futures = { version = "^0.3", default-features = false, optional = true }
 
 [dev-dependencies]
 crossbeam-channel = "0.4"

src/async.rs

@@ -6,43 +6,33 @@ use std::{
     task::{Context, Poll},
 };
 use crate::*;
+use futures::{Stream, stream::FusedStream, future::FusedFuture};
 
-/// A future  used to receive a value from the channel.
-pub struct RecvFuture<'a, T> {
-    recv: &'a mut Receiver<T>,
-}
-
-impl<'a, T> RecvFuture<'a, T> {
-    pub(crate) fn new(recv: &mut Receiver<T>) -> RecvFuture<T> {
-        RecvFuture { recv }
-    }
-}
-
-impl<'a, T> Future for RecvFuture<'a, T> {
-    type Output = Result<T, RecvError>;
-
-    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        // On success, set the waker to none to avoid it being woken again in case that is wrong
-        // TODO: `poll_recv` instead to prevent even spinning?
-        let mut buf = self.recv.buffer.borrow_mut();
+impl<T> Receiver<T> {
+    #[inline]
+    fn poll(&self, cx: &mut Context<'_>) -> Poll<Result<T, RecvError>> {
+        let mut buf = self.buffer.borrow_mut();
 
         let res = if let Some(msg) = buf.pop_front() {
             return Poll::Ready(Ok(msg));
         } else {
             self
-                .recv
                 .shared
                 .poll_inner()
                 .map(|mut inner| self
-                    .recv
                     .shared
-                    .try_recv(move || {
-                        // Detach the waker
-                        inner.recv_waker = None;
-                        // Inform the sender that we no longer need waking
-                        inner.listen_mode = 1;
-                        inner
-                    }, &mut buf))
+                    .try_recv(
+                        move || {
+                            // Detach the waker
+                            inner.recv_waker = None;
+                            // Inform the sender that we no longer need waking
+                            inner.listen_mode = 1;
+                            inner
+                        },
+                        &mut buf,
+                        &self.finished,
+                    )
+                )
         };
 
         let poll = match res {
@@ -64,3 +54,46 @@ impl<'a, T> Future for RecvFuture<'a, T> {
         poll
     }
 }
+
+/// A future  used to receive a value from the channel.
+pub struct RecvFuture<'a, T> {
+    recv: &'a mut Receiver<T>,
+}
+
+impl<'a, T> RecvFuture<'a, T> {
+    pub(crate) fn new(recv: &mut Receiver<T>) -> RecvFuture<T> {
+        RecvFuture { recv }
+    }
+}
+
+impl<'a, T> Future for RecvFuture<'a, T> {
+    type Output = Result<T, RecvError>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        self.recv.poll(cx)
+    }
+}
+
+impl<'a, T> FusedFuture for RecvFuture<'a, T> {
+    fn is_terminated(&self) -> bool {
+        self.recv.finished.get()
+    }
+}
+
+impl<T> Stream for Receiver<T> {
+    type Item = T;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        self.poll(cx).map(|ready| ready.ok())
+    }
+    
+    fn size_hint(&self) -> (usize, Option<usize>) {
+	    (self.buffer.borrow().len(), None)
+    }
+}
+
+impl<T> FusedStream for Receiver<T> {
+    fn is_terminated(&self) -> bool {
+        self.finished.get()
+    }
+}

src/lib.rs

@@ -41,6 +41,7 @@ use std::task::Waker;
 use crate::select::Token;
 #[cfg(feature = "async")]
 use crate::r#async::RecvFuture;
+use std::cell::Cell;
 
 /// An error that may be emitted when attempting to send a value into a channel on a sender.
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
@@ -260,6 +261,7 @@ impl<T> Shared<T> {
     }
 
     #[inline]
+    #[cfg(feature = "async")]
     fn poll_inner(&self) -> Option<MutexGuard<'_, Inner<T>>> {
         #[cfg(windows)] { self.inner.try_lock().ok() }
         #[cfg(not(windows))] { self.inner.try_lock() }
@@ -363,6 +365,7 @@ impl<T> Shared<T> {
         &'a self,
         take_inner: impl FnOnce() -> MutexGuard<'a, Inner<T>>,
         buf: &mut VecDeque<T>,
+        finished: &Cell<bool>,
     ) -> Result<T, (MutexGuard<Inner<T>>, TryRecvError)> {
         // Eagerly check the buffer
         if let Some(msg) = buf.pop_front() {
@@ -380,8 +383,10 @@ impl<T> Shared<T> {
                 msg
             },
             // If there's nothing more in the queue, this might be because there are no senders
-            None if inner.sender_count == 0 =>
-                return Err((inner, TryRecvError::Disconnected)),
+            None if inner.sender_count == 0 => {
+                finished.set(true);
+                return Err((inner, TryRecvError::Disconnected));
+            },
             None => return Err((inner, TryRecvError::Empty)),
         };
 
@@ -412,12 +417,13 @@ impl<T> Shared<T> {
     fn recv(
         &self,
         buf: &mut VecDeque<T>,
+        finished: &Cell<bool>,
     ) -> Result<T, RecvError> {
         loop {
             // Attempt to receive a message
             let mut i = 0;
             let inner = loop {
-                match self.try_recv(|| self.wait_inner(), buf) {
+                match self.try_recv(|| self.wait_inner(), buf, finished) {
                     Ok(msg) => return Ok(msg),
                     Err((_, TryRecvError::Disconnected)) => return Err(RecvError::Disconnected),
                     Err((inner, TryRecvError::Empty)) if i == 3 => break inner,
@@ -438,9 +444,10 @@ impl<T> Shared<T> {
         &self,
         deadline: Instant,
         buf: &mut VecDeque<T>,
+        finished: &Cell<bool>,
     ) -> Result<T, RecvTimeoutError> {
         // Attempt a speculative recv. If we are lucky there might be a message in the queue!
-        let mut inner = match self.try_recv(|| self.wait_inner(), buf) {
+        let mut inner = match self.try_recv(|| self.wait_inner(), buf, finished) {
             Ok(msg) => return Ok(msg),
             Err((_, TryRecvError::Disconnected)) => return Err(RecvTimeoutError::Disconnected),
             Err((inner, TryRecvError::Empty)) => inner,
@@ -466,7 +473,7 @@ impl<T> Shared<T> {
             }
 
             // Attempt to receive a message from the queue
-            inner = match self.try_recv(|| self.wait_inner(), buf) {
+            inner = match self.try_recv(|| self.wait_inner(), buf, finished) {
                 Ok(msg) => return Ok(msg),
                 Err((inner, TryRecvError::Empty)) => inner,
                 Err((_, TryRecvError::Disconnected)) => return Err(RecvTimeoutError::Disconnected),
@@ -555,28 +562,35 @@ pub struct Receiver<T> {
     /// Used to prevent Sync being implemented for this type - we never actually use it!
     /// TODO: impl<T> !Sync for Receiver<T> {} when negative traits are stable
     _phantom_cell: UnsafeCell<()>,
+    /// Whether all receivers have disconnected and there are no messages in any buffer
+    finished: Cell<bool>,
 }
 
 impl<T> Receiver<T> {
     /// Wait for an incoming value from the channel associated with this receiver, returning an
     /// error if all channel senders have been dropped.
     pub fn recv(&self) -> Result<T, RecvError> {
-        self.shared.recv(&mut self.buffer.borrow_mut())
+        self.shared.recv(&mut self.buffer.borrow_mut(), &self.finished)
     }
 
     /// Wait for an incoming value from the channel associated with this receiver, returning an
     /// error if all channel senders have been dropped or the timeout has expired.
     pub fn recv_timeout(&self, timeout: Duration) -> Result<T, RecvTimeoutError> {
         self.shared.recv_deadline(
             Instant::now().checked_add(timeout).unwrap(),
-            &mut self.buffer.borrow_mut()
+            &mut self.buffer.borrow_mut(),
+            &self.finished
         )
     }
 
     /// Wait for an incoming value from the channel associated with this receiver, returning an
     /// error if all channel senders have been dropped or the deadline has passed.
     pub fn recv_deadline(&self, deadline: Instant) -> Result<T, RecvTimeoutError> {
-        self.shared.recv_deadline(deadline, &mut self.buffer.borrow_mut())
+        self.shared.recv_deadline(
+            deadline,
+            &mut self.buffer.borrow_mut(),
+            &self.finished
+        )
     }
 
     // Takes `&mut self` to avoid >1 task waiting on this channel
@@ -593,7 +607,11 @@ impl<T> Receiver<T> {
     pub fn try_recv(&self) -> Result<T, TryRecvError> {
         self
             .shared
-            .try_recv(|| self.shared.wait_inner(), &mut self.buffer.borrow_mut())
+            .try_recv(
+                || self.shared.wait_inner(),
+                &mut self.buffer.borrow_mut(),
+                &self.finished
+            )
             .map_err(|(_, err)| err)
     }
 
@@ -716,6 +734,7 @@ pub fn unbounded<T>() -> (Sender<T>, Receiver<T>) {
         Receiver {
             shared,
             buffer: RefCell::new(VecDeque::new()),
+            finished: Cell::new(false),
             _phantom_cell: UnsafeCell::new(())
         },
     )
@@ -750,6 +769,7 @@ pub fn bounded<T>(cap: usize) -> (Sender<T>, Receiver<T>) {
         Receiver {
             shared,
             buffer: RefCell::new(VecDeque::new()),
+            finished: Cell::new(false),
             _phantom_cell: UnsafeCell::new(())
         },
     )