Search code examples
rustrust-tokio

Rust async_trait return `impl Stream`?


I want to write a mock for an existing async class, without changing the original class to much.

To demonstrate my problem I took the source code from Guillaume Endignoux Blog.

I want to change it to an async trait:

use std::time::{Duration, Instant};

use async_trait::async_trait;
use futures::{StreamExt as _, Stream, stream};
use lazy_static::lazy_static;
use rand::distributions::{Uniform, Distribution as _};
use tokio::time::sleep;

lazy_static! {
    static ref START_TIME: Instant = Instant::now();
}

struct Pages;

#[tokio::main]
async fn main() {
    println!("First 10 pages:\n{:?}", Pages::get_n_pages(10).await);
}

#[async_trait]
trait Page {
    async fn get_page(i: usize) -> Vec<usize>;
    async fn get_n_pages(n: usize) -> Vec<Vec<usize>>;
    fn get_pages() -> impl Stream<Item = Vec<usize>>;
}

#[async_trait]
impl Page for Pages {

    async fn get_page(i: usize) -> Vec<usize> {
        let millis = Uniform::from(0..10).sample(&mut rand::thread_rng());
        println!(
            "[{}] # get_page({}) will complete in {} ms",
            START_TIME.elapsed().as_millis(),
            i,
            millis
        );

        sleep(Duration::from_millis(millis)).await;
        println!(
            "[{}] # get_page({}) completed",
            START_TIME.elapsed().as_millis(),
            i
        );

        (10 * i..10 * (i + 1)).collect()
    }

    async fn get_n_pages(n: usize) -> Vec<Vec<usize>> {
        Self::get_pages().take(n).collect().await
    }

    fn get_pages() -> impl Stream<Item = Vec<usize>> {
        stream::iter(0..).then(|i| Self::get_page(i))
    }
}

But then I get the message: `impl Trait` only allowed in function and inherent method return types, not in `impl` method return.

So I change it to dyn, but then I get the message: doesn't have a size known at compile-time.

Is there a way to write async traits that return a kind of impl Stream<Item = Vec<usize>>?


Solution

  • There is an unstable feature allowing impl Trait in return positions of traits, return_position_impl_trait_in_trait. With it you can use your initial syntax at the cost of running nightly:

    #![feature(return_position_impl_trait_in_trait)]
    
    use async_trait::async_trait;
    use futures::{StreamExt as _, Stream, stream};
    
    #[async_trait]
    impl Page for Pages {
        //...
        fn get_pages() -> impl Stream<Item = Vec<usize>> {
            stream::iter(0..).then(|i| Self::get_page(i))
        }
    }
    

    Playground