๐ŸŽ“ A1Learner
๐Ÿ”

Rust Streams & Tasks

A stream is an async sequence of items that arrive over time. Each item takes one await to receive.

Items arrive one await at a time

next().await waits for the next item and returns an Option โ€” Some while items remain, None at the end:

let mut s = stream::iter(vec![10, 20]);
println!("first = {:?}", s.next().await);

The Stream trait is Futureโ€™s multi-shot sibling

A stream polls the same way as Future::poll โ€” but it can answer many times before it is done. (Simplified from the futures crate's trait โ€” same shape.)

trait Stream {
    type Item;
    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>)
        -> Poll<Option<Self::Item>>;
}

Hand-write a tiny stream

impl Stream differs from impl Future in one line: the answer type is Poll<Option<Item>>. Some(n) is a value, None is 'the stream is over'.

struct Countdown(u32);

impl Stream for Countdown {
    type Item = u32;

    fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>)
        -> Poll<Option<u32>>
    {
        if self.0 == 0 {
            Poll::Ready(None)
        } else {
            let n = self.0;
            self.0 -= 1;
            Poll::Ready(Some(n))
        }
    }
}

Drain it with next()

next() borrows the stream mutably. It returns a future for one item. It is the async answer to Iterator::next. It is the loop shape from this lesson.

let mut s = Countdown(3);
while let Some(n) = s.next().await {
    println!("tick {n}");
}

Handle several items at once

Sequential next() leaves concurrency on the table. for_each_concurrent runs up to limit of the item futures at the same time. It is the Async Book's Streams chapter answer to 'my stream is slow'.

use futures::stream::{self, StreamExt}; // for iter, for_each_concurrent

stream::iter(1..=4)
    .for_each_concurrent(2, |n| async move { handle(n).await })
    .await;

A deadline for each item

tokio-stream's timeout wraps every item in a Result: Ok(item), or Err when that one wait ran out. A stream holding a timer is not Unpin, so pin! it in place before polling:

use tokio_stream::StreamExt; // for timeout and next

let mut s = tokio_stream::iter(vec![1, 2]).timeout(std::time::Duration::from_millis(50));
tokio::pin!(s);
println!("item = {:?}", s.next().await);

Return a stream from a function

A stream's real type is long and has no name you can write. Box::pin hides it behind one nameable type.

fn countdown() -> Pin<Box<dyn Stream<Item = u32>>> {
    Box::pin(stream::iter(vec![3, 2, 1]))
}
let mut s = countdown();
println!("first = {:?}", s.next().await);

Build a stream from state: unfold

unfold holds one piece of state. Each step returns an item and the next state. None ends the stream. pin_mut! pins the stream in place, so next() may poll it.

let t = stream::unfold(3, |n| async move {
    if n == 0 { None } else { Some((n, n - 1)) }
});
futures::pin_mut!(t);
println!("item = {:?}", t.next().await);

Example

// With tokio-stream: stream::iter(1..=3).map(|n| n * n)
fn squares() -> Vec<i32> { (1..=3).map(|n| n * n).collect() }
fn main() {
    // A stream version yields the same values asynchronously, one await at a time.
    println!("streamed squares = {:?}", squares());
}

Key points

  • Use a JoinSet to look after a changing group of spawned tasks and collect their results
  • Cancelling a task is simple in Rust: drop its future or handle, and the task stops
  • A CancellationToken lets many tasks agree to shut down gracefully together
  • Stream::poll_next mirrors Future::poll exactly: Poll::Ready(Some(item)) yields a value, Poll::Ready(None) ends the stream, Poll::Pending means ask again after a wake.
  • A stream must never be polled again after it returned Ready(None)
  • Use stream::unfold to build a stream from a piece of state.
  • Each step yields one item and the next state.
  • It yields None to end
  • Streams live in the futures crate and in tokio-stream.
  • Messages on a socket, rows from a database, and events from a user are examples of streams.

In this example, the runnable code shows the values first and the stream second. squares() uses an ordinary iterator map to square 1, 2, 3, giving [1, 4, 9]. The comment above it shows the stream version: the same map, but each value would arrive through an .await. The main function prints the collected squares. The shape of the pipeline โ€” source, transform, collect โ€” is identical. Only the timing changed: iterators deliver instantly, streams deliver when the world is ready.

Note

The stream trait is still stabilising, so a crate provides it today.

Try it Yourself

Edit the code, then press Run.

Running is simulated ยท Rust Playground

Exercise

How is a running async task normally cancelled in Rust?