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>
impl<K, V> StreamMap<K, V>
sourcepub fn iter(&self) -> impl Iterator<Item = &(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);
}
sourcepub fn iter_mut(&mut self) -> impl Iterator<Item = &mut (K, V)>
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);
}
sourcepub fn new() -> StreamMap<K, V>
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();
sourcepub fn with_capacity(capacity: usize) -> StreamMap<K, V>
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);
sourcepub fn keys(&self) -> impl Iterator<Item = &K>
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);
}
sourcepub fn values(&self) -> impl Iterator<Item = &V>
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);
}
sourcepub fn values_mut(&mut self) -> impl Iterator<Item = &mut V>
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);
}
sourcepub fn capacity(&self) -> usize
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);
sourcepub fn len(&self) -> usize
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);
sourcepub fn is_empty(&self) -> bool
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());
sourcepub fn clear(&mut self)
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());
sourcepub fn insert(&mut self, k: K, stream: V) -> Option<V>
pub fn insert(&mut self, k: K, stream: V) -> Option<V>
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());
sourcepub fn remove<Q>(&mut self, k: &Q) -> Option<V>
pub fn remove<Q>(&mut self, k: &Q) -> Option<V>
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());
sourcepub fn contains_key<Q>(&self, k: &Q) -> bool
pub fn contains_key<Q>(&self, k: &Q) -> bool
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>
impl<K, V> StreamMap<K, V>
sourcepub async fn next_many(
&mut self,
buffer: &mut Vec<(K, V::Item)>,
limit: usize,
) -> usize
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.
sourcepub fn poll_next_many(
&mut self,
cx: &mut Context<'_>,
buffer: &mut Vec<(K, V::Item)>,
limit: usize,
) -> Poll<usize>
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 theStreamMap
is not empty.Poll::Ready(count)
wherecount
is the number of items successfully received and stored inbuffer
. This can be less than, or equal to,limit
.Poll::Ready(0)
iflimit
is set to zero or when theStreamMap
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, V> Extend<(K, V)> for StreamMap<K, V>
impl<K, V> Extend<(K, V)> for StreamMap<K, V>
source§fn extend<T>(&mut self, iter: T)where
T: IntoIterator<Item = (K, V)>,
fn extend<T>(&mut self, iter: T)where
T: IntoIterator<Item = (K, V)>,
source§fn extend_one(&mut self, item: A)
fn extend_one(&mut self, item: A)
extend_one
)source§fn extend_reserve(&mut self, additional: usize)
fn extend_reserve(&mut self, additional: usize)
extend_one
)source§impl<K, V> FromIterator<(K, V)> for StreamMap<K, V>
impl<K, V> FromIterator<(K, V)> for StreamMap<K, V>
source§impl<K, V> Stream for StreamMap<K, V>
impl<K, V> Stream for StreamMap<K, V>
Auto Trait Implementations§
impl<K, V> Freeze for StreamMap<K, V>
impl<K, V> RefUnwindSafe for StreamMap<K, V>where
K: RefUnwindSafe,
V: RefUnwindSafe,
impl<K, V> Send for StreamMap<K, V>
impl<K, V> Sync for StreamMap<K, V>
impl<K, V> Unpin for StreamMap<K, V>
impl<K, V> UnwindSafe for StreamMap<K, V>where
K: UnwindSafe,
V: UnwindSafe,
Blanket Implementations§
source§impl<T> BorrowMut<T> for Twhere
T: ?Sized,
impl<T> BorrowMut<T> for Twhere
T: ?Sized,
source§fn borrow_mut(&mut self) -> &mut T
fn borrow_mut(&mut self) -> &mut T
source§impl<St> StreamExt for St
impl<St> StreamExt for St
source§fn next(&mut self) -> Next<'_, Self>where
Self: Unpin,
fn next(&mut self) -> Next<'_, Self>where
Self: Unpin,
None
if the
stream is finished. Read moresource§fn try_next<T, E>(&mut self) -> TryNext<'_, Self>
fn try_next<T, E>(&mut self) -> TryNext<'_, Self>
source§fn map<T, F>(self, f: F) -> Map<Self, F>
fn map<T, F>(self, f: F) -> Map<Self, F>
source§fn map_while<T, F>(self, f: F) -> MapWhile<Self, F>
fn map_while<T, F>(self, f: F) -> MapWhile<Self, F>
None
. Read moresource§fn then<F, Fut>(self, f: F) -> Then<Self, Fut, F>
fn then<F, Fut>(self, f: F) -> Then<Self, Fut, F>
source§fn merge<U>(self, other: U) -> Merge<Self, U>
fn merge<U>(self, other: U) -> Merge<Self, U>
source§fn filter<F>(self, f: F) -> Filter<Self, F>
fn filter<F>(self, f: F) -> Filter<Self, F>
source§fn filter_map<T, F>(self, f: F) -> FilterMap<Self, F>
fn filter_map<T, F>(self, f: F) -> FilterMap<Self, F>
source§fn fuse(self) -> Fuse<Self>where
Self: Sized,
fn fuse(self) -> Fuse<Self>where
Self: Sized,
None
. Read moresource§fn take(self, n: usize) -> Take<Self>where
Self: Sized,
fn take(self, n: usize) -> Take<Self>where
Self: Sized,
n
items of the underlying stream. Read moresource§fn take_while<F>(self, f: F) -> TakeWhile<Self, F>
fn take_while<F>(self, f: F) -> TakeWhile<Self, F>
true
. Read moresource§fn skip(self, n: usize) -> Skip<Self>where
Self: Sized,
fn skip(self, n: usize) -> Skip<Self>where
Self: Sized,
n
first items of the
underlying stream. Read moresource§fn skip_while<F>(self, f: F) -> SkipWhile<Self, F>
fn skip_while<F>(self, f: F) -> SkipWhile<Self, F>
true
. Read more