Fuse the damn oneshot receiver

todo: hide the enums after all

Author: rklaehn

src/lib.rs

@@ -69,6 +69,8 @@ pub mod channel {
     pub mod oneshot {
         use std::{fmt::Debug, future::Future, io, pin::Pin, task};
 
+        use crate::util::FusedOneshotReceiver;
+
         pub fn channel<T>() -> (Sender<T>, Receiver<T>) {
             let (tx, rx) = tokio::sync::oneshot::channel();
             (tx.into(), rx.into())
@@ -152,7 +154,7 @@ pub mod channel {
         impl<T> crate::Sender for Sender<T> {}
 
         pub enum Receiver<T> {
-            Tokio(tokio::sync::oneshot::Receiver<T>),
+            Tokio(FusedOneshotReceiver<T>),
             Boxed(BoxedReceiver<T>),
         }
 
@@ -172,7 +174,7 @@ pub mod channel {
         /// Convert a tokio oneshot receiver to a receiver for this crate
         impl<T> From<tokio::sync::oneshot::Receiver<T>> for Receiver<T> {
             fn from(rx: tokio::sync::oneshot::Receiver<T>) -> Self {
-                Self::Tokio(rx)
+                Self::Tokio(FusedOneshotReceiver(rx))
             }
         }
 

src/util.rs

@@ -290,3 +290,23 @@ mod varint_util {
     }
 }
 pub use varint_util::{AsyncReadVarintExt, WriteVarintExt};
+
+mod fuse_wrapper {
+    use std::{future::Future, pin::Pin, task::{Context, Poll}};
+
+    pub struct FusedOneshotReceiver<T>(pub tokio::sync::oneshot::Receiver<T>);
+
+    impl<T> Future for FusedOneshotReceiver<T> {
+        type Output = std::result::Result<T, tokio::sync::oneshot::error::RecvError>;
+
+        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+            if self.0.is_terminated() {
+                // don't panic when polling a terminated receiver
+                Poll::Pending
+            } else {
+                Future::poll(Pin::new(&mut self.0), cx)
+            }
+        }
+    }
+}
+pub use fuse_wrapper::FusedOneshotReceiver;