Implementing Stream

The Stream trait is very similar to the Future trait.

  1. use std::pin::Pin;
  2. use std::task::{Context, Poll};
  4. pub trait Stream {
  5.     type Item;
  7.     fn poll_next(
  8.         self: Pin<&mut Self>, 
  9.         cx: &mut Context<'_>
  10.     ) -> Poll<Option<Self::Item>>;
  12.     fn size_hint(&self) -> (usize, Option<usize>) {
  13.         (0, None)
  14.     }
  15. }

The Stream::poll_next() function is much like Future::poll, except it can be called repeatedly to receive many values from the stream. Just as we saw in Async in depth, when a stream is not ready to return a value, Poll::Pending is returned instead. The task’s waker is registered. Once the stream should be polled again, the waker is notified.

The size_hint() method is used the same way as it is with iterators.

Usually, when manually implementing a Stream, it is done by composing futures and other streams. As an example, let’s build off of the Delay future we implemented in Async in depth. We will convert it to a stream that yields () three times at 10 ms intervals

  1. use tokio_stream::Stream;
  2. use std::pin::Pin;
  3. use std::task::{Context, Poll};
  4. use std::time::Duration;
  6. struct Interval {
  7.     rem: usize,
  8.     delay: Delay,
  9. }
  11. impl Stream for Interval {
  12.     type Item = ();
  14.     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>)
  15.         -> Poll<Option<()>>
  16.     {
  17.         if self.rem == 0 {
  18.             // No more delays
  19.             return Poll::Ready(None);
  20.         }
  22.         match Pin::new(&mut self.delay).poll(cx) {
  23.             Poll::Ready(_) => {
  24.                 let when = self.delay.when + Duration::from_millis(10);
  25.                 self.delay = Delay { when };
  26.                 self.rem -= 1;
  27.                 Poll::Ready(Some(()))
  28.             }
  29.             Poll::Pending => Poll::Pending,
  30.         }
  31.     }
  32. }

async-stream

Manually implementing streams using the Stream trait can be tedious. Unfortunately, the Rust programming language does not yet support async/await syntax for defining streams. This is in the works, but not yet ready.

The async-stream crate is available as a temporary solution. This crate provides an async_stream! macro that transforms the input into a stream. Using this crate, the above interval can be implemented like this:

  1. use async_stream::stream;
  2. use std::time::{Duration, Instant};
  4. stream! {
  5.     let mut when = Instant::now();
  6.     for _ in 0..3 {
  7.         let delay = Delay { when };
  8.         delay.await;
  9.         yield ();
  10.         when += Duration::from_millis(10);
  11.     }
  12. }