Struct tokio_stream::StreamMap

source · [−]
pub struct StreamMap<K, V> { /* private fields */ }
Expand description

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"),
        }
    }
}

Implementations

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);
}

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);
}

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();

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);

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);
}

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);
}

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);
}

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);

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);

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());

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());

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());

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());

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);

Trait Implementations

Formats the value using the given formatter. Read more
Returns the “default value” for a type. Read more
Extends a collection with the contents of an iterator. Read more
🔬This is a nightly-only experimental API. (extend_one)
Extends a collection with exactly one element.
🔬This is a nightly-only experimental API. (extend_one)
Reserves capacity in a collection for the given number of additional elements. Read more
Creates a value from an iterator. Read more
Values yielded by the stream.
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. Read more
Returns the bounds on the remaining length of the stream. Read more

Auto Trait Implementations

Blanket Implementations

Gets the TypeId of self. Read more
Immutably borrows from an owned value. Read more
Mutably borrows from an owned value. Read more

Returns the argument unchanged.

Calls U::from(self).

That is, this conversion is whatever the implementation of From<T> for U chooses to do.

The type returned in the event of a conversion error.
Performs the conversion.
The type returned in the event of a conversion error.
Performs the conversion.