Creating a stream of values while calling async fns?
Asked Answered
E

2

4

I can't figure out how to provide a Stream where I await async functions to get the data needed for the values of the stream.

I've tried to implement the the Stream trait directly, but I run into issues because I'd like to use async things like awaiting, the compiler does not want me to call async functions.

I assume that I'm missing some background on what the goal of Stream is and I'm just attacking this incorrectly and perhaps I shouldn't be looking at Stream at all, but I don't know where else to turn. I've seen the other functions in the stream module that could be useful, but I'm unsure how I could store any state and use these functions.

As a slightly simplified version of my actual goal, I want to provide a stream of 64-byte Vecs from an AsyncRead object (i.e. TCP stream), but also store a little state inside whatever logic ends up producing values for the stream, in this example, a counter.

pub struct Receiver<T>
where
    T: AsyncRead + Unpin,
{
    readme: T,
    num: u64,
}

// ..code for a simple `new() -> Self` function..

impl<T> Stream for Receiver<T>
where
    T: AsyncRead + Unpin,
{
    type Item = Result<Vec<u8>, io::Error>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let mut buf: [u8; 64] = [0; 64];
        match self.readme.read_exact(&mut buf).await {
            Ok(()) => {
                self.num += 1;
                Poll::Ready(Some(Ok(buf.to_vec())))
            }
            Err(e) => Poll::Ready(Some(Err(e))),
        }
    }
}

This fails to build, saying

error[E0728]: `await` is only allowed inside `async` functions and blocks

I'm using rustc 1.36.0-nightly (d35181ad8 2019-05-20) and my Cargo.toml looks like this:

[dependencies]
futures-preview = { version = "0.3.0-alpha.16", features = ["compat", "io-compat"] }
pin-utils = "0.1.0-alpha.4"
Exegetic answered 22/6, 2019 at 23:51 Comment(0)
E
0

Answer copy/pasted from the reddit post by user Matthias247:

It's unfortunately not possible at the moment - Streams have to be implemented by hand and can not utilize async fn. Whether it's possible to change this in the future is unclear.

You can work around it by defining a different Stream trait which makes use of Futures like:

trait Stream<T> { 
   type NextFuture: Future<Output=T>;

   fn next(&mut self) -> Self::NextFuture; 
}

This article and this futures-rs issue have more information around it.

Exegetic answered 23/6, 2019 at 6:45 Comment(0)
M
0

You can do it with gen-stream crate:

#![feature(generators, generator_trait, gen_future)]

use {
    futures::prelude::*,
    gen_stream::{gen_await, GenTryStream},
    pin_utils::unsafe_pinned,
    std::{
        io,
        marker::PhantomData,
        pin::Pin,
        sync::{
            atomic::{AtomicU64, Ordering},
            Arc,
        },
        task::{Context, Poll},
    },
};

pub type Inner = Pin<Box<dyn Stream<Item = Result<Vec<u8>, io::Error>> + Send>>;

pub struct Receiver<T> {
    inner: Inner,
    pub num: Arc<AtomicU64>,
    _marker: PhantomData<T>,
}

impl<T> Receiver<T> {
    unsafe_pinned!(inner: Inner);
}

impl<T> From<T> for Receiver<T>
where
    T: AsyncRead + Unpin + Send + 'static,
{
    fn from(mut readme: T) -> Self {
        let num = Arc::new(AtomicU64::new(0));

        Self {
            inner: Box::pin(GenTryStream::from({
                let num = num.clone();
                static move || loop {
                    let mut buf: [u8; 64] = [0; 64];
                    match gen_await!(readme.read_exact(&mut buf)) {
                        Ok(()) => {
                            num.fetch_add(1, Ordering::Relaxed);
                            yield Poll::Ready(buf.to_vec())
                        }
                        Err(e) => return Err(e),
                    }
                }
            })),
            num,
            _marker: PhantomData,
        }
    }
}

impl<T> Stream for Receiver<T>
where
    T: AsyncRead + Unpin,
{
    type Item = Result<Vec<u8>, io::Error>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.inner().poll_next(cx)
    }
}
Maricela answered 23/6, 2019 at 20:15 Comment(0)

© 2022 - 2024 — McMap. All rights reserved.