Allow calling UnboundedReceiver::try_next after None

futures-channel/src/mpsc/mod.rs

@@ -1020,9 +1020,6 @@ impl<T> Receiver<T> {
     /// It is not recommended to call this function from inside of a future,
     /// only when you've otherwise arranged to be notified when the channel is
     /// no longer empty.
-    ///
-    /// This function will panic if called after `try_next` or `poll_next` has
-    /// returned `None`.
     pub fn try_next(&mut self) -> Result<Option<T>, TryRecvError> {
         match self.next_message() {
             Poll::Ready(msg) => {
@@ -1033,7 +1030,10 @@ impl<T> Receiver<T> {
     }
 
     fn next_message(&mut self) -> Poll<Option<T>> {
-        let inner = self.inner.as_mut().expect("Receiver::next_message called after `None`");
+        let inner = match self.inner.as_mut() {
+            None => return Poll::Ready(None),
+            Some(inner) => inner,
+        };
         // Pop off a message
         match unsafe { inner.message_queue.pop_spin() } {
             Some(msg) => {
@@ -1173,9 +1173,6 @@ impl<T> UnboundedReceiver<T> {
     /// * `Ok(Some(t))` when message is fetched
     /// * `Ok(None)` when channel is closed and no messages left in the queue
     /// * `Err(e)` when there are no messages available, but channel is not yet closed
-    ///
-    /// This function will panic if called after `try_next` or `poll_next` has
-    /// returned `None`.
     pub fn try_next(&mut self) -> Result<Option<T>, TryRecvError> {
         match self.next_message() {
             Poll::Ready(msg) => {
@@ -1186,7 +1183,10 @@ impl<T> UnboundedReceiver<T> {
     }
 
     fn next_message(&mut self) -> Poll<Option<T>> {
-        let inner = self.inner.as_mut().expect("Receiver::next_message called after `None`");
+        let inner = match self.inner.as_mut() {
+            None => return Poll::Ready(None),
+            Some(inner) => inner,
+        };
         // Pop off a message
         match unsafe { inner.message_queue.pop_spin() } {
             Some(msg) => {

futures-channel/tests/mpsc-close.rs

@@ -276,3 +276,25 @@ fn stress_try_send_as_receiver_closes() {
     bg.join()
         .expect("background thread join");
 }
+
+#[test]
+fn unbounded_try_next_after_none() {
+    let (tx, mut rx) = mpsc::unbounded::<String>();
+    // Drop the sender, close the channel.
+    drop(tx);
+    // Receive the end of channel.
+    assert_eq!(Ok(None), rx.try_next().map_err(|_| ()));
+    // None received, check we can call `try_next` again.
+    assert_eq!(Ok(None), rx.try_next().map_err(|_| ()));
+}
+
+#[test]
+fn bounded_try_next_after_none() {
+    let (tx, mut rx) = mpsc::channel::<String>(17);
+    // Drop the sender, close the channel.
+    drop(tx);
+    // Receive the end of channel.
+    assert_eq!(Ok(None), rx.try_next().map_err(|_| ()));
+    // None received, check we can call `try_next` again.
+    assert_eq!(Ok(None), rx.try_next().map_err(|_| ()));
+}