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§

source§

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

source

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

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

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

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

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

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

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

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

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

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

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

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

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

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

impl<K, V> StreamMap<K, V>
where K: Clone + Unpin, V: Stream + Unpin,

source

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.

source

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§

source§

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

source§

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

Formats the value using the given formatter. Read more
source§

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

source§

fn default() -> Self

Returns the “default value” for a type. Read more
source§

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

source§

fn extend<T>(&mut self, iter: T)
where T: IntoIterator<Item = (K, V)>,

Extends a collection with the contents of an iterator. Read more
source§

fn extend_one(&mut self, item: A)

🔬This is a nightly-only experimental API. (extend_one)
Extends a collection with exactly one element.
source§

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. Read more
source§

impl<K, V> FromIterator<(K, V)> for StreamMap<K, V>
where K: Hash + Eq,

source§

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

Creates a value from an iterator. Read more
source§

impl<K, V> Stream for StreamMap<K, V>
where K: Clone + Unpin, V: Stream + Unpin,

source§

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

Values yielded by the stream.
source§

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. Read more
source§

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

Returns the bounds on the remaining length of the stream. Read more

Auto Trait Implementations§

§

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

§

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

§

impl<K, V> Send for StreamMap<K, V>
where K: Send, V: Send,

§

impl<K, V> Sync for StreamMap<K, V>
where K: Sync, V: Sync,

§

impl<K, V> Unpin for StreamMap<K, V>
where K: Unpin, V: Unpin,

§

impl<K, V> UnwindSafe for StreamMap<K, V>
where K: UnwindSafe, V: UnwindSafe,

Blanket Implementations§

source§

impl<T> Any for T
where T: 'static + ?Sized,

source§

fn type_id(&self) -> TypeId

Gets the TypeId of self. Read more
source§

impl<T> Borrow<T> for T
where T: ?Sized,

source§

fn borrow(&self) -> &T

Immutably borrows from an owned value. Read more
source§

impl<T> BorrowMut<T> for T
where T: ?Sized,

source§

fn borrow_mut(&mut self) -> &mut T

Mutably borrows from an owned value. Read more
source§

impl<T> From<T> for T

source§

fn from(t: T) -> T

Returns the argument unchanged.

source§

impl<T, U> Into<U> for T
where U: From<T>,

source§

fn into(self) -> U

Calls U::from(self).

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

source§

impl<St> StreamExt for St
where St: Stream + ?Sized,

source§

fn next(&mut self) -> Next<'_, Self>
where Self: Unpin,

Consumes and returns the next value in the stream or None if the stream is finished. Read more
source§

fn try_next<T, E>(&mut self) -> TryNext<'_, Self>
where Self: Stream<Item = Result<T, E>> + Unpin,

Consumes and returns the next item in the stream. If an error is encountered before the next item, the error is returned instead. Read more
source§

fn map<T, F>(self, f: F) -> Map<Self, F>
where F: FnMut(Self::Item) -> T, Self: Sized,

Maps this stream’s items to a different type, returning a new stream of the resulting type. Read more
source§

fn map_while<T, F>(self, f: F) -> MapWhile<Self, F>
where F: FnMut(Self::Item) -> Option<T>, Self: Sized,

Map this stream’s items to a different type for as long as determined by the provided closure. A stream of the target type will be returned, which will yield elements until the closure returns None. Read more
source§

fn then<F, Fut>(self, f: F) -> Then<Self, Fut, F>
where F: FnMut(Self::Item) -> Fut, Fut: Future, Self: Sized,

Maps this stream’s items asynchronously to a different type, returning a new stream of the resulting type. Read more
source§

fn merge<U>(self, other: U) -> Merge<Self, U>
where U: Stream<Item = Self::Item>, Self: Sized,

Combine two streams into one by interleaving the output of both as it is produced. Read more
source§

fn filter<F>(self, f: F) -> Filter<Self, F>
where F: FnMut(&Self::Item) -> bool, Self: Sized,

Filters the values produced by this stream according to the provided predicate. Read more
source§

fn filter_map<T, F>(self, f: F) -> FilterMap<Self, F>
where F: FnMut(Self::Item) -> Option<T>, Self: Sized,

Filters the values produced by this stream while simultaneously mapping them to a different type according to the provided closure. Read more
source§

fn fuse(self) -> Fuse<Self>
where Self: Sized,

Creates a stream which ends after the first None. Read more
source§

fn take(self, n: usize) -> Take<Self>
where Self: Sized,

Creates a new stream of at most n items of the underlying stream. Read more
source§

fn take_while<F>(self, f: F) -> TakeWhile<Self, F>
where F: FnMut(&Self::Item) -> bool, Self: Sized,

Take elements from this stream while the provided predicate resolves to true. Read more
source§

fn skip(self, n: usize) -> Skip<Self>
where Self: Sized,

Creates a new stream that will skip the n first items of the underlying stream. Read more
source§

fn skip_while<F>(self, f: F) -> SkipWhile<Self, F>
where F: FnMut(&Self::Item) -> bool, Self: Sized,

Skip elements from the underlying stream while the provided predicate resolves to true. Read more
source§

fn all<F>(&mut self, f: F) -> AllFuture<'_, Self, F>
where Self: Unpin, F: FnMut(Self::Item) -> bool,

Tests if every element of the stream matches a predicate. Read more
source§

fn any<F>(&mut self, f: F) -> AnyFuture<'_, Self, F>
where Self: Unpin, F: FnMut(Self::Item) -> bool,

Tests if any element of the stream matches a predicate. Read more
source§

fn chain<U>(self, other: U) -> Chain<Self, U>
where U: Stream<Item = Self::Item>, Self: Sized,

Combine two streams into one by first returning all values from the first stream then all values from the second stream. Read more
source§

fn fold<B, F>(self, init: B, f: F) -> FoldFuture<Self, B, F>
where Self: Sized, F: FnMut(B, Self::Item) -> B,

A combinator that applies a function to every element in a stream producing a single, final value. Read more
source§

fn collect<T>(self) -> Collect<Self, T>
where T: FromStream<Self::Item>, Self: Sized,

Drain stream pushing all emitted values into a collection. Read more
source§

fn timeout(self, duration: Duration) -> Timeout<Self>
where Self: Sized,

Applies a per-item timeout to the passed stream. Read more
source§

fn timeout_repeating(self, interval: Interval) -> TimeoutRepeating<Self>
where Self: Sized,

Applies a per-item timeout to the passed stream. Read more
source§

fn throttle(self, duration: Duration) -> Throttle<Self>
where Self: Sized,

Slows down a stream by enforcing a delay between items. Read more
source§

fn chunks_timeout( self, max_size: usize, duration: Duration, ) -> ChunksTimeout<Self>
where Self: Sized,

Batches the items in the given stream using a maximum duration and size for each batch. Read more
source§

fn peekable(self) -> Peekable<Self>
where Self: Sized,

Turns the stream into a peekable stream, whose next element can be peeked at without being consumed. Read more
source§

impl<T, U> TryFrom<U> for T
where U: Into<T>,

source§

type Error = Infallible

The type returned in the event of a conversion error.
source§

fn try_from(value: U) -> Result<T, <T as TryFrom<U>>::Error>

Performs the conversion.
source§

impl<T, U> TryInto<U> for T
where U: TryFrom<T>,

source§

type Error = <U as TryFrom<T>>::Error

The type returned in the event of a conversion error.
source§

fn try_into(self) -> Result<U, <U as TryFrom<T>>::Error>

Performs the conversion.