rmcgibbo / async-priority-channel

Apache License 2.0
11 stars 2 forks source link

Provide `Stream` adaptor for `Receiver` #69

Open lovesegfault opened 1 year ago

lovesegfault commented 1 year ago

It'd be great if there was an easy way to turn a Receiver into a futures::stream::Stream.

rmcgibbo commented 1 year ago

Hi!

Just based on my other commitments, I am almost surely not going to get to this. I'm happy to merge PRs, but that's about all I have time for these days.

lovesegfault commented 1 year ago

I quickly hacked this together, but I'm not sure the implementation is correct:

struct PriorityReceiverStream<I, P: Ord>(async_priority_channel::Receiver<I, P>);

impl<I, P: Ord> PriorityReceiverStream<I, P> {
    fn new(recv: async_priority_channel::Receiver<I, P>) -> Self {
        Self(recv)
    }
}

impl<I: Send, P: Ord + Send> Stream for PriorityReceiverStream<I, P> {
    type Item = (I, P);
    fn poll_next(
        self: Pin<&mut Self>,
        _cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Option<Self::Item>> {
        loop {
            match self.0.try_recv() {
                Ok(msg) => {
                    return std::task::Poll::Ready(Some(msg));
                }
                Err(async_priority_channel::TryRecvError::Closed) => {
                    return std::task::Poll::Ready(None);
                }
                Err(async_priority_channel::TryRecvError::Empty) => {}
            }
        }
    }
}
lovesegfault commented 10 months ago

Much nicer approach:

fn async_receiver_stream<T, P: Ord>(
    chan: async_priority_channel::Receiver<T, P>,
) -> impl Stream<Item = (T, P)> + Unpin {
    Box::pin(futures::stream::unfold(chan, |state| async move {
        match state.recv().await {
            Ok(val) => Some((val, state)),
            Err(async_priority_channel::RecvError) => None,
        }
    }))
}
rmcgibbo commented 10 months ago

For the un-initiated (i.e. me), can you explain (or link to a blog post) about what the benefits of using Stream would be? I guess this allows merging streams using various methods on the StreamExt trait rather that writing explicit loops with select!?

xamgore commented 2 months ago

Here is the implementation of a wrapper from tokio, would stick to their approach.

use crate::Stream;
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::sync::mpsc::UnboundedReceiver;

/// A wrapper around [`tokio::sync::mpsc::UnboundedReceiver`] that implements [`Stream`].
///
/// [`tokio::sync::mpsc::UnboundedReceiver`]: struct@tokio::sync::mpsc::UnboundedReceiver
/// [`Stream`]: trait@crate::Stream
#[derive(Debug)]
pub struct UnboundedReceiverStream<T> {
    inner: UnboundedReceiver<T>,
}

impl<T> UnboundedReceiverStream<T> {
    /// Create a new `UnboundedReceiverStream`.
    pub fn new(recv: UnboundedReceiver<T>) -> Self {
        Self { inner: recv }
    }

    /// Get back the inner `UnboundedReceiver`.
    pub fn into_inner(self) -> UnboundedReceiver<T> {
        self.inner
    }

    /// Closes the receiving half of a channel without dropping it.
    ///
    /// This prevents any further messages from being sent on the channel while
    /// still enabling the receiver to drain messages that are buffered.
    pub fn close(&mut self) {
        self.inner.close();
    }
}

impl<T> Stream for UnboundedReceiverStream<T> {
    type Item = T;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.inner.poll_recv(cx)
    }
}

impl<T> AsRef<UnboundedReceiver<T>> for UnboundedReceiverStream<T> {
    fn as_ref(&self) -> &UnboundedReceiver<T> {
        &self.inner
    }
}

impl<T> AsMut<UnboundedReceiver<T>> for UnboundedReceiverStream<T> {
    fn as_mut(&mut self) -> &mut UnboundedReceiver<T> {
        &mut self.inner
    }
}

impl<T> From<UnboundedReceiver<T>> for UnboundedReceiverStream<T> {
    fn from(recv: UnboundedReceiver<T>) -> Self {
        Self::new(recv)
    }
}