Struct tokio_stream::StreamMap

pub struct StreamMap<K, V> {
    entries: Vec<(K, V)>,
}

Combine many streams into one, indexing each source stream with a unique key.

StreamMap is similar to StreamExt::merge in that it combines source streams into a single merged stream that yields values in the order that they arrive from the source streams. However, StreamMap has a lot more flexibility in usage patterns.

StreamMap can:

• Merge an arbitrary number of streams.
• Track which source stream the value was received from.
• Handle inserting and removing streams from the set of managed streams at any point during iteration.

All source streams held by StreamMap are indexed using a key. This key is included with the value when a source stream yields a value. The key is also used to remove the stream from the StreamMap before the stream has completed streaming.

Unpin

Because the StreamMap API moves streams during runtime, both streams and keys must be Unpin. In order to insert a !Unpin stream into a StreamMap, use pin! to pin the stream to the stack or Box::pin to pin the stream in the heap.

Implementation

StreamMap is backed by a Vec<(K, V)>. There is no guarantee that this internal implementation detail will persist in future versions, but it is important to know the runtime implications. In general, StreamMap works best with a “smallish” number of streams as all entries are scanned on insert, remove, and polling. In cases where a large number of streams need to be merged, it may be advisable to use tasks sending values on a shared mpsc channel.

Notes

StreamMap removes finished streams automatically, without alerting the user. In some scenarios, the caller would want to know on closed streams. To do this, use StreamNotifyClose as a wrapper to your stream. It will return None when the stream is closed.

Examples

Merging two streams, then remove them after receiving the first value

use tokio_stream::{StreamExt, StreamMap, Stream};
use tokio::sync::mpsc;
use std::pin::Pin;

#[tokio::main]
async fn main() {
    let (tx1, mut rx1) = mpsc::channel::<usize>(10);
    let (tx2, mut rx2) = mpsc::channel::<usize>(10);

    // Convert the channels to a `Stream`.
    let rx1 = Box::pin(async_stream::stream! {
          while let Some(item) = rx1.recv().await {
              yield item;
          }
    }) as Pin<Box<dyn Stream<Item = usize> + Send>>;

    let rx2 = Box::pin(async_stream::stream! {
          while let Some(item) = rx2.recv().await {
              yield item;
          }
    }) as Pin<Box<dyn Stream<Item = usize> + Send>>;

    tokio::spawn(async move {
        tx1.send(1).await.unwrap();

        // This value will never be received. The send may or may not return
        // `Err` depending on if the remote end closed first or not.
        let _ = tx1.send(2).await;
    });

    tokio::spawn(async move {
        tx2.send(3).await.unwrap();
        let _ = tx2.send(4).await;
    });

    let mut map = StreamMap::new();

    // Insert both streams
    map.insert("one", rx1);
    map.insert("two", rx2);

    // Read twice
    for _ in 0..2 {
        let (key, val) = map.next().await.unwrap();

        if key == "one" {
            assert_eq!(val, 1);
        } else {
            assert_eq!(val, 3);
        }

        // Remove the stream to prevent reading the next value
        map.remove(key);
    }
}

This example models a read-only client to a chat system with channels. The client sends commands to join and leave channels. StreamMap is used to manage active channel subscriptions.

For simplicity, messages are displayed with println!, but they could be sent to the client over a socket.

use tokio_stream::{Stream, StreamExt, StreamMap};

enum Command {
    Join(String),
    Leave(String),
}

fn commands() -> impl Stream<Item = Command> {
    // Streams in user commands by parsing `stdin`.
}

// Join a channel, returns a stream of messages received on the channel.
fn join(channel: &str) -> impl Stream<Item = String> + Unpin {
    // left as an exercise to the reader
}

#[tokio::main]
async fn main() {
    let mut channels = StreamMap::new();

    // Input commands (join / leave channels).
    let cmds = commands();
    tokio::pin!(cmds);

    loop {
        tokio::select! {
            Some(cmd) = cmds.next() => {
                match cmd {
                    Command::Join(chan) => {
                        // Join the channel and add it to the `channels`
                        // stream map
                        let msgs = join(&chan);
                        channels.insert(chan, msgs);
                    }
                    Command::Leave(chan) => {
                        channels.remove(&chan);
                    }
                }
            }
            Some((chan, msg)) = channels.next() => {
                // Received a message, display it on stdout with the channel
                // it originated from.
                println!("{}: {}", chan, msg);
            }
            // Both the `commands` stream and the `channels` stream are
            // complete. There is no more work to do, so leave the loop.
            else => break,
        }
    }
}

Using StreamNotifyClose to handle closed streams with StreamMap.

use tokio_stream::{StreamExt, StreamMap, StreamNotifyClose};

#[tokio::main]
async fn main() {
    let mut map = StreamMap::new();
    let stream = StreamNotifyClose::new(tokio_stream::iter(vec![0, 1]));
    let stream2 = StreamNotifyClose::new(tokio_stream::iter(vec![0, 1]));
    map.insert(0, stream);
    map.insert(1, stream2);
    while let Some((key, val)) = map.next().await {
        match val {
            Some(val) => println!("got {val:?} from stream {key:?}"),
            None => println!("stream {key:?} closed"),
        }
    }
}

Fields

entries: Vec<(K, V)>

Streams stored in the map

Implementations

impl<K, V> StreamMap<K, V>

pub fn iter(&self) -> impl Iterator<Item = &(K, V)>

An iterator visiting all key-value pairs in arbitrary order.

The iterator element type is &'a (K, V).

Examples

use tokio_stream::{StreamMap, pending};

let mut map = StreamMap::new();

map.insert("a", pending::<i32>());
map.insert("b", pending());
map.insert("c", pending());

for (key, stream) in map.iter() {
    println!("({}, {:?})", key, stream);
}

pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut (K, V)>

An iterator visiting all key-value pairs mutably in arbitrary order.

The iterator element type is &'a mut (K, V).

Examples

use tokio_stream::{StreamMap, pending};

let mut map = StreamMap::new();

map.insert("a", pending::<i32>());
map.insert("b", pending());
map.insert("c", pending());

for (key, stream) in map.iter_mut() {
    println!("({}, {:?})", key, stream);
}

pub fn new() -> StreamMap<K, V>

Creates an empty StreamMap.

The stream map is initially created with a capacity of 0, so it will not allocate until it is first inserted into.

Examples

use tokio_stream::{StreamMap, Pending};

let map: StreamMap<&str, Pending<()>> = StreamMap::new();

pub fn with_capacity(capacity: usize) -> StreamMap<K, V>

Creates an empty StreamMap with the specified capacity.

The stream map will be able to hold at least capacity elements without reallocating. If capacity is 0, the stream map will not allocate.

Examples

use tokio_stream::{StreamMap, Pending};

let map: StreamMap<&str, Pending<()>> = StreamMap::with_capacity(10);

pub fn keys(&self) -> impl Iterator<Item = &K>

Returns an iterator visiting all keys in arbitrary order.

The iterator element type is &'a K.

Examples

use tokio_stream::{StreamMap, pending};

let mut map = StreamMap::new();

map.insert("a", pending::<i32>());
map.insert("b", pending());
map.insert("c", pending());

for key in map.keys() {
    println!("{}", key);
}

pub fn values(&self) -> impl Iterator<Item = &V>

An iterator visiting all values in arbitrary order.

The iterator element type is &'a V.

Examples

use tokio_stream::{StreamMap, pending};

let mut map = StreamMap::new();

map.insert("a", pending::<i32>());
map.insert("b", pending());
map.insert("c", pending());

for stream in map.values() {
    println!("{:?}", stream);
}

pub fn values_mut(&mut self) -> impl Iterator<Item = &mut V>

An iterator visiting all values mutably in arbitrary order.

The iterator element type is &'a mut V.

Examples

use tokio_stream::{StreamMap, pending};

let mut map = StreamMap::new();

map.insert("a", pending::<i32>());
map.insert("b", pending());
map.insert("c", pending());

for stream in map.values_mut() {
    println!("{:?}", stream);
}

pub fn capacity(&self) -> usize

Returns the number of streams the map can hold without reallocating.

This number is a lower bound; the StreamMap might be able to hold more, but is guaranteed to be able to hold at least this many.

Examples

use tokio_stream::{StreamMap, Pending};

let map: StreamMap<i32, Pending<()>> = StreamMap::with_capacity(100);
assert!(map.capacity() >= 100);

pub fn len(&self) -> usize

Returns the number of streams in the map.

Examples

use tokio_stream::{StreamMap, pending};

let mut a = StreamMap::new();
assert_eq!(a.len(), 0);
a.insert(1, pending::<i32>());
assert_eq!(a.len(), 1);

pub fn is_empty(&self) -> bool

Returns true if the map contains no elements.

Examples

use tokio_stream::{StreamMap, pending};

let mut a = StreamMap::new();
assert!(a.is_empty());
a.insert(1, pending::<i32>());
assert!(!a.is_empty());

pub fn clear(&mut self)

Clears the map, removing all key-stream pairs. Keeps the allocated memory for reuse.

Examples

use tokio_stream::{StreamMap, pending};

let mut a = StreamMap::new();
a.insert(1, pending::<i32>());
a.clear();
assert!(a.is_empty());

pub fn insert(&mut self, k: K, stream: V) -> Option<V>

where K: Hash + Eq,

Insert a key-stream pair into the map.

If the map did not have this key present, None is returned.

If the map did have this key present, the new stream replaces the old one and the old stream is returned.

Examples

use tokio_stream::{StreamMap, pending};

let mut map = StreamMap::new();

assert!(map.insert(37, pending::<i32>()).is_none());
assert!(!map.is_empty());

map.insert(37, pending());
assert!(map.insert(37, pending()).is_some());

pub fn remove<Q>(&mut self, k: &Q) -> Option<V>

where K: Borrow<Q>, Q: Hash + Eq + ?Sized,

Removes a key from the map, returning the stream at the key if the key was previously in the map.

The key may be any borrowed form of the map’s key type, but Hash and Eq on the borrowed form must match those for the key type.

Examples

use tokio_stream::{StreamMap, pending};

let mut map = StreamMap::new();
map.insert(1, pending::<i32>());
assert!(map.remove(&1).is_some());
assert!(map.remove(&1).is_none());

pub fn contains_key<Q>(&self, k: &Q) -> bool

where K: Borrow<Q>, Q: Hash + Eq + ?Sized,

Returns true if the map contains a stream for the specified key.

The key may be any borrowed form of the map’s key type, but Hash and Eq on the borrowed form must match those for the key type.

Examples

use tokio_stream::{StreamMap, pending};

let mut map = StreamMap::new();
map.insert(1, pending::<i32>());
assert_eq!(map.contains_key(&1), true);
assert_eq!(map.contains_key(&2), false);

impl<K, V> StreamMap<K, V>

where K: Unpin, V: Stream + Unpin,

fn poll_next_entry( &mut self, cx: &mut Context<'_>, ) -> Poll<Option<(usize, V::Item)>>

Polls the next value, includes the vec entry index

impl<K, V> StreamMap<K, V>

where K: Clone + Unpin, V: Stream + Unpin,

pub async fn next_many( &mut self, buffer: &mut Vec<(K, V::Item)>, limit: usize, ) -> usize

Receives multiple items on this StreamMap, extending the provided buffer.

This method returns the number of items that is appended to the buffer.

Note that this method does not guarantee that exactly limit items are received. Rather, if at least one item is available, it returns as many items as it can up to the given limit. This method returns zero only if the StreamMap is empty (or if limit is zero).

Cancel safety

This method is cancel safe. If next_many is used as the event in a tokio::select! statement and some other branch completes first, it is guaranteed that no items were received on any of the underlying streams.

pub fn poll_next_many( &mut self, cx: &mut Context<'_>, buffer: &mut Vec<(K, V::Item)>, limit: usize, ) -> Poll<usize>

Polls to receive multiple items on this StreamMap, extending the provided buffer.

This method returns:

• Poll::Pending if no items are available but the StreamMap is not empty.
• Poll::Ready(count) where count is the number of items successfully received and stored in buffer. This can be less than, or equal to, limit.
• Poll::Ready(0) if limit is set to zero or when the StreamMap is empty.

Note that this method does not guarantee that exactly limit items are received. Rather, if at least one item is available, it returns as many items as it can up to the given limit. This method returns zero only if the StreamMap is empty (or if limit is zero).

Trait Implementations

impl<K: Debug, V: Debug> Debug for StreamMap<K, V>

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl<K, V> Default for StreamMap<K, V>

fn default() -> Self

Returns the “default value” for a type.

impl<K, V> Extend<(K, V)> for StreamMap<K, V>

fn extend<T>(&mut self, iter: T)

where T: IntoIterator<Item = (K, V)>,

Extends a collection with the contents of an iterator.

fn extend_one(&mut self, item: A)

🔬This is a nightly-only experimental API. (extend_one)

Extends a collection with exactly one element.

fn extend_reserve(&mut self, additional: usize)

🔬This is a nightly-only experimental API. (extend_one)

Reserves capacity in a collection for the given number of additional elements.

impl<K, V> FromIterator<(K, V)> for StreamMap<K, V>

where K: Hash + Eq,

fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self

Creates a value from an iterator.

impl<K, V> Stream for StreamMap<K, V>

where K: Clone + Unpin, V: Stream + Unpin,

type Item = (K, <V as Stream>::Item)

Values yielded by the stream.

fn poll_next( self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll<Option<Self::Item>>

Attempt to pull out the next value of this stream, registering the current task for wakeup if the value is not yet available, and returning None if the stream is exhausted.

fn size_hint(&self) -> (usize, Option<usize>)

Returns the bounds on the remaining length of the stream.