Streams

A stream is an asynchronous series of values. It is the asynchronous equivalent to Rust’s std::iter::Iterator and is represented by the Stream trait. Streams can be iterated in async functions. They can also be transformed using adapters. Tokio provides a number of common adapters on the StreamExt trait.

Tokio provides stream support in a separate crate: tokio-stream.

  1. tokio-stream = "0.1"

Currently, Tokio’s Stream utilities exist in the tokio-stream crate. Once the Stream trait is stabilized in the Rust standard library, Tokio’s stream utilities will be moved into the tokio crate.

Iteration

Currently, the Rust programming language does not support async for loops. Instead, iterating streams is done using a while let loop paired with StreamExt::next().

  1. use tokio_stream::StreamExt;
  3. #[tokio::main]
  4. async fn main() {
  5.     let mut stream = tokio_stream::iter(&[1, 2, 3]);
  7.     while let Some(v) = stream.next().await {
  8.         println!("GOT = {:?}", v);
  9.     }
  10. }

Like iterators, the next() method returns Option<T> where T is the stream’s value type. Receiving None indicates that stream iteration is terminated.

Mini-Redis broadcast

Let’s go over a slightly more complicated example using the Mini-Redis client.

Full code can be found here.

  1. use tokio_stream::StreamExt;
  2. use mini_redis::client;
  4. async fn publish() -> mini_redis::Result<()> {
  5.     let mut client = client::connect("127.0.0.1:6379").await?;
  7.     // Publish some data
  8.     client.publish("numbers", "1".into()).await?;
  9.     client.publish("numbers", "two".into()).await?;
  10.     client.publish("numbers", "3".into()).await?;
  11.     client.publish("numbers", "four".into()).await?;
  12.     client.publish("numbers", "five".into()).await?;
  13.     client.publish("numbers", "6".into()).await?;
  14.     Ok(())
  15. }
  17. async fn subscribe() -> mini_redis::Result<()> {
  18.     let client = client::connect("127.0.0.1:6379").await?;
  19.     let subscriber = client.subscribe(vec!["numbers".to_string()]).await?;
  20.     let messages = subscriber.into_stream();
  22.     tokio::pin!(messages);
  24.     while let Some(msg) = messages.next().await {
  25.         println!("got = {:?}", msg);
  26.     }
  28.     Ok(())
  29. }
  31. #[tokio::main]
  32. async fn main() -> mini_redis::Result<()> {
  33.     tokio::spawn(async {
  34.         publish().await
  35.     });
  37.     subscribe().await?;
  39.     println!("DONE");
  41.     Ok(())
  42. }

A task is spawned to publish messages to the Mini-Redis server on the “numbers” channel. Then, on the main task, we subscribe to the “numbers” channel and display received messages.

After subscribing, into_stream() is called on the returned subscriber. This consumes the Subscriber, returning a stream that yields messages as they arrive. Before we start iterating the messages, note that the stream is pinned to the stack using tokio::pin!. Calling next() on a stream requires the stream to be pinned. The into_stream() function returns a stream that is not pinned, we must explicitly pin it in order to iterate it.

A Rust value is “pinned” when it can no longer be moved in memory. A key property of a pinned value is that pointers can be taken to the pinned data and the caller can be confident the pointer stays valid. This feature is used by async/await to support borrowing data across .await points.

If we forget to pin the stream, we get an error like this:

  1. error[E0277]: `from_generator::GenFuture<[static generator@Subscriber::into_stream::{closure#0} for<'r, 's, 't0, 't1, 't2, 't3, 't4, 't5, 't6> {ResumeTy, &'r mut Subscriber, Subscriber, impl Future, (), std::result::Result<Option<Message>, Box<(dyn std::error::Error + Send + Sync + 't0)>>, Box<(dyn std::error::Error + Send + Sync + 't1)>, &'t2 mut async_stream::yielder::Sender<std::result::Result<Message, Box<(dyn std::error::Error + Send + Sync + 't3)>>>, async_stream::yielder::Sender<std::result::Result<Message, Box<(dyn std::error::Error + Send + Sync + 't4)>>>, std::result::Result<Message, Box<(dyn std::error::Error + Send + Sync + 't5)>>, impl Future, Option<Message>, Message}]>` cannot be unpinned
  2.   --> streams/src/main.rs:29:36
  3.    |
  4. 29 |     while let Some(msg) = messages.next().await {
  5.    |                                    ^^^^ within `tokio_stream::filter::_::__Origin<'_, impl Stream, [closure@streams/src/main.rs:22:17: 25:10]>`, the trait `Unpin` is not implemented for `from_generator::GenFuture<[static generator@Subscriber::into_stream::{closure#0} for<'r, 's, 't0, 't1, 't2, 't3, 't4, 't5, 't6> {ResumeTy, &'r mut Subscriber, Subscriber, impl Future, (), std::result::Result<Option<Message>, Box<(dyn std::error::Error + Send + Sync + 't0)>>, Box<(dyn std::error::Error + Send + Sync + 't1)>, &'t2 mut async_stream::yielder::Sender<std::result::Result<Message, Box<(dyn std::error::Error + Send + Sync + 't3)>>>, async_stream::yielder::Sender<std::result::Result<Message, Box<(dyn std::error::Error + Send + Sync + 't4)>>>, std::result::Result<Message, Box<(dyn std::error::Error + Send + Sync + 't5)>>, impl Future, Option<Message>, Message}]>`
  6.    |
  7.    = note: required because it appears within the type `impl Future`
  8.    = note: required because it appears within the type `async_stream::async_stream::AsyncStream<std::result::Result<Message, Box<(dyn std::error::Error + Send + Sync + 'static)>>, impl Future>`
  9.    = note: required because it appears within the type `impl Stream`
  10.    = note: required because it appears within the type `tokio_stream::filter::_::__Origin<'_, impl Stream, [closure@streams/src/main.rs:22:17: 25:10]>`
  11.    = note: required because of the requirements on the impl of `Unpin` for `tokio_stream::filter::Filter<impl Stream, [closure@streams/src/main.rs:22:17: 25:10]>`
  12.    = note: required because it appears within the type `tokio_stream::map::_::__Origin<'_, tokio_stream::filter::Filter<impl Stream, [closure@streams/src/main.rs:22:17: 25:10]>, [closure@streams/src/main.rs:26:14: 26:40]>`
  13.    = note: required because of the requirements on the impl of `Unpin` for `tokio_stream::map::Map<tokio_stream::filter::Filter<impl Stream, [closure@streams/src/main.rs:22:17: 25:10]>, [closure@streams/src/main.rs:26:14: 26:40]>`
  14.    = note: required because it appears within the type `tokio_stream::take::_::__Origin<'_, tokio_stream::map::Map<tokio_stream::filter::Filter<impl Stream, [closure@streams/src/main.rs:22:17: 25:10]>, [closure@streams/src/main.rs:26:14: 26:40]>>`
  15.    = note: required because of the requirements on the impl of `Unpin` for `tokio_stream::take::Take<tokio_stream::map::Map<tokio_stream::filter::Filter<impl Stream, [closure@streams/src/main.rs:22:17: 25:10]>, [closure@streams/src/main.rs:26:14: 26:40]>>`

If you hit an error message like this, try pinning the value!

Before trying to run this, start the Mini-Redis server:

  1. $ mini-redis-server

Then try running the code. We will see the messages outputted to STDOUT.

  1. got = Ok(Message { channel: "numbers", content: b"1" })
  2. got = Ok(Message { channel: "numbers", content: b"two" })
  3. got = Ok(Message { channel: "numbers", content: b"3" })
  4. got = Ok(Message { channel: "numbers", content: b"four" })
  5. got = Ok(Message { channel: "numbers", content: b"five" })
  6. got = Ok(Message { channel: "numbers", content: b"6" })

Some early messages may be dropped as there is a race between subscribing and publishing. The program never exits. A subscription to a Mini-Redis channel stays active as long as the server is active.

Let’s see how we can work with streams to expand on this program.

Adapters

Functions that take a Stream and return another Stream are often called ‘stream adapters’, as they’re a form of the ‘adapter pattern’. Common stream adapters include map, take, and filter.

Lets update the Mini-Redis so that it will exit. After receiving three messages, stop iterating messages. This is done using take. This adapter limits the stream to yield at most n messages.

  1. let messages = subscriber
  2.     .into_stream()
  3.     .take(3);

Running the program again, we get:

  1. got = Ok(Message { channel: "numbers", content: b"1" })
  2. got = Ok(Message { channel: "numbers", content: b"two" })
  3. got = Ok(Message { channel: "numbers", content: b"3" })

This time the program ends.

Now, let’s limit the stream to single digit numbers. We will check this by checking for the message length. We use the filter adapter to drop any message that does not match the predicate.

  1. let messages = subscriber
  2.     .into_stream()
  3.     .filter(|msg| match msg {
  4.         Ok(msg) if msg.content.len() == 1 => true,
  5.         _ => false,
  6.     })
  7.     .take(3);

Running the program again, we get:

  1. got = Ok(Message { channel: "numbers", content: b"1" })
  2. got = Ok(Message { channel: "numbers", content: b"3" })
  3. got = Ok(Message { channel: "numbers", content: b"6" })

Note that the order in which adapters are applied matters. Calling filter first then take is different than calling take then filter.

Finally, we will tidy up the output by stripping the Ok(Message { ... }) part of the output. This is done with map. Because this is applied after filter, we know the message is Ok, so we can use unwrap().

  1. let messages = subscriber
  2.     .into_stream()
  3.     .filter(|msg| match msg {
  4.         Ok(msg) if msg.content.len() == 1 => true,
  5.         _ => false,
  6.     })
  7.     .map(|msg| msg.unwrap().content)
  8.     .take(3);

Now, the output is:

  1. got = b"1"
  2. got = b"3"
  3. got = b"6"

Another option would be to combine the filter and map steps into a single call using filter_map.

There are more available adapters. See the list here.

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. }