mirror of
https://github.com/tokio-rs/tokio.git
synced 2026-08-21 00:00:10 +02:00
stream: provide StreamMap utility (#2185)
`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. 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.
This commit is contained in:
Carl Lerche
@@ -50,6 +50,9 @@ pub use once::{once, Once};
|
||||
mod pending;
|
||||
pub use pending::{pending, Pending};
|
||||
|
||||
mod stream_map;
|
||||
pub use stream_map::StreamMap;
|
||||
|
||||
mod try_next;
|
||||
use try_next::TryNext;
|
||||
|
||||
|
||||
@@ -0,0 +1,503 @@
|
||||
use crate::stream::Stream;
|
||||
|
||||
use std::borrow::Borrow;
|
||||
use std::hash::Hash;
|
||||
use std::pin::Pin;
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
/// 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.
|
||||
///
|
||||
/// [`StreamExt::merge`]: crate::stream::StreamExt::merge
|
||||
/// [`mpsc`]: crate::sync::mpsc
|
||||
/// [`pin!`]: macro@pin
|
||||
/// [`Box::pin`]: std::boxed::Box::pin
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// Merging two streams, then remove them after receiving the first value
|
||||
///
|
||||
/// ```
|
||||
/// use tokio::stream::{StreamExt, StreamMap};
|
||||
/// use tokio::sync::mpsc;
|
||||
///
|
||||
/// #[tokio::main]
|
||||
/// async fn main() {
|
||||
/// let (mut tx1, rx1) = mpsc::channel(10);
|
||||
/// let (mut tx2, rx2) = mpsc::channel(10);
|
||||
///
|
||||
/// 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.
|
||||
///
|
||||
/// ```no_run
|
||||
/// 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`.
|
||||
/// # tokio::stream::pending()
|
||||
/// }
|
||||
///
|
||||
/// // 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::stream::pending()
|
||||
/// }
|
||||
///
|
||||
/// #[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,
|
||||
/// }
|
||||
/// }
|
||||
/// }
|
||||
/// ```
|
||||
#[derive(Debug, Default)]
|
||||
pub struct StreamMap<K, V> {
|
||||
/// Streams stored in the map
|
||||
entries: Vec<(K, V)>,
|
||||
}
|
||||
|
||||
impl<K, V> 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 new() -> StreamMap<K, V> {
|
||||
StreamMap { entries: vec![] }
|
||||
}
|
||||
|
||||
/// 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 with_capacity(capacity: usize) -> StreamMap<K, V> {
|
||||
StreamMap {
|
||||
entries: Vec::with_capacity(capacity),
|
||||
}
|
||||
}
|
||||
|
||||
/// 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 keys(&self) -> impl Iterator<Item = &K> {
|
||||
self.entries.iter().map(|(k, _)| k)
|
||||
}
|
||||
|
||||
/// 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(&self) -> impl Iterator<Item = &V> {
|
||||
self.entries.iter().map(|(_, v)| 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 values_mut(&mut self) -> impl Iterator<Item = &mut V> {
|
||||
self.entries.iter_mut().map(|(_, v)| v)
|
||||
}
|
||||
|
||||
/// 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 capacity(&self) -> usize {
|
||||
self.entries.capacity()
|
||||
}
|
||||
|
||||
/// 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 len(&self) -> usize {
|
||||
self.entries.len()
|
||||
}
|
||||
|
||||
/// Returns `true` if the map contains no elements.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use std::collections::HashMap;
|
||||
///
|
||||
/// let mut a = HashMap::new();
|
||||
/// assert!(a.is_empty());
|
||||
/// a.insert(1, "a");
|
||||
/// assert!(!a.is_empty());
|
||||
/// ```
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.entries.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());
|
||||
/// ```
|
||||
pub fn clear(&mut self) {
|
||||
self.entries.clear();
|
||||
}
|
||||
|
||||
/// 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 insert(&mut self, k: K, stream: V) -> Option<V>
|
||||
where
|
||||
K: Hash + Eq,
|
||||
{
|
||||
let ret = self.remove(&k);
|
||||
self.entries.push((k, stream));
|
||||
|
||||
ret
|
||||
}
|
||||
|
||||
/// 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 remove<Q: ?Sized>(&mut self, k: &Q) -> Option<V>
|
||||
where
|
||||
K: Borrow<Q>,
|
||||
Q: Hash + Eq,
|
||||
{
|
||||
for i in 0..self.entries.len() {
|
||||
if self.entries[i].0.borrow() == k {
|
||||
return Some(self.entries.swap_remove(i).1);
|
||||
}
|
||||
}
|
||||
|
||||
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);
|
||||
/// ```
|
||||
pub fn contains_key<Q: ?Sized>(&self, k: &Q) -> bool
|
||||
where
|
||||
K: Borrow<Q>,
|
||||
Q: Hash + Eq,
|
||||
{
|
||||
for i in 0..self.entries.len() {
|
||||
if self.entries[i].0.borrow() == k {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
impl<K, V> StreamMap<K, V>
|
||||
where
|
||||
K: Unpin,
|
||||
V: Stream + Unpin,
|
||||
{
|
||||
/// Polls the next value, includes the vec entry index
|
||||
fn poll_next_entry(&mut self, cx: &mut Context<'_>) -> Poll<Option<(usize, V::Item)>> {
|
||||
use Poll::*;
|
||||
|
||||
let start = crate::util::thread_rng_n(self.entries.len() as u32) as usize;
|
||||
let mut idx = start;
|
||||
|
||||
for _ in 0..self.entries.len() {
|
||||
let (_, stream) = &mut self.entries[idx];
|
||||
|
||||
match Pin::new(stream).poll_next(cx) {
|
||||
Ready(Some(val)) => return Ready(Some((idx, val))),
|
||||
Ready(None) => {
|
||||
// Remove the entry
|
||||
self.entries.swap_remove(idx);
|
||||
|
||||
// Check if this was the last entry, if so the cursor needs
|
||||
// to wrap
|
||||
if idx == self.entries.len() {
|
||||
idx = 0;
|
||||
} else if idx < start && start <= self.entries.len() {
|
||||
// The stream being swapped into the current index has
|
||||
// already been polled, so skip it.
|
||||
idx = idx.wrapping_add(1) % self.entries.len();
|
||||
}
|
||||
}
|
||||
Pending => {
|
||||
idx = idx.wrapping_add(1) % self.entries.len();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// If the map is empty, then the stream is complete.
|
||||
if self.entries.is_empty() {
|
||||
Ready(None)
|
||||
} else {
|
||||
Pending
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<K, V> Stream for StreamMap<K, V>
|
||||
where
|
||||
K: Clone + Unpin,
|
||||
V: Stream + Unpin,
|
||||
{
|
||||
type Item = (K, V::Item);
|
||||
|
||||
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||
if let Some((idx, val)) = ready!(self.poll_next_entry(cx)) {
|
||||
let key = self.entries[idx].0.clone();
|
||||
Poll::Ready(Some((key, val)))
|
||||
} else {
|
||||
Poll::Ready(None)
|
||||
}
|
||||
}
|
||||
|
||||
fn size_hint(&self) -> (usize, Option<usize>) {
|
||||
let mut ret = (0, Some(0));
|
||||
|
||||
for (_, stream) in &self.entries {
|
||||
let hint = stream.size_hint();
|
||||
|
||||
ret.0 += hint.0;
|
||||
|
||||
match (ret.1, hint.1) {
|
||||
(Some(a), Some(b)) => ret.1 = Some(a + b),
|
||||
(Some(_), None) => ret.1 = None,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
ret
|
||||
}
|
||||
}
|
||||
@@ -3,7 +3,7 @@ cfg_io_driver! {
|
||||
pub(crate) mod slab;
|
||||
}
|
||||
|
||||
#[cfg(any(feature = "rt-threaded", feature = "macros"))]
|
||||
#[cfg(any(feature = "rt-threaded", feature = "macros", feature = "stream"))]
|
||||
mod rand;
|
||||
|
||||
cfg_rt_threaded! {
|
||||
@@ -16,6 +16,6 @@ cfg_rt_threaded! {
|
||||
pub(crate) use try_lock::TryLock;
|
||||
}
|
||||
|
||||
cfg_macros! {
|
||||
pub use rand::thread_rng_n;
|
||||
}
|
||||
#[cfg(any(feature = "macros", feature = "stream"))]
|
||||
#[cfg_attr(not(feature = "macros"), allow(unreachable_pub))]
|
||||
pub use rand::thread_rng_n;
|
||||
|
||||
@@ -51,15 +51,14 @@ impl FastRand {
|
||||
}
|
||||
}
|
||||
|
||||
// Used by the select macro
|
||||
cfg_macros! {
|
||||
// Used by the select macro and `StreamMap`
|
||||
#[cfg(any(feature = "macros", feature = "stream"))]
|
||||
#[doc(hidden)]
|
||||
#[cfg_attr(not(feature = "macros"), allow(unreachable_pub))]
|
||||
pub fn thread_rng_n(n: u32) -> u32 {
|
||||
thread_local! {
|
||||
static THREAD_RNG: FastRand = FastRand::new(crate::loom::rand::seed());
|
||||
}
|
||||
|
||||
// Used by macros
|
||||
#[doc(hidden)]
|
||||
pub fn thread_rng_n(n: u32) -> u32 {
|
||||
THREAD_RNG.with(|rng| rng.fastrand_n(n))
|
||||
}
|
||||
THREAD_RNG.with(|rng| rng.fastrand_n(n))
|
||||
}
|
||||
|
||||
@@ -0,0 +1,374 @@
|
||||
use tokio::stream::{self, pending, Stream, StreamExt, StreamMap};
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_test::{assert_ok, assert_pending, assert_ready, task};
|
||||
|
||||
use std::pin::Pin;
|
||||
|
||||
macro_rules! assert_ready_some {
|
||||
($($t:tt)*) => {
|
||||
match assert_ready!($($t)*) {
|
||||
Some(v) => v,
|
||||
None => panic!("expected `Some`, got `None`"),
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
macro_rules! assert_ready_none {
|
||||
($($t:tt)*) => {
|
||||
match assert_ready!($($t)*) {
|
||||
None => {}
|
||||
Some(v) => panic!("expected `None`, got `Some({:?})`", v),
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn empty() {
|
||||
let mut map = StreamMap::<&str, stream::Pending<()>>::new();
|
||||
|
||||
assert_eq!(map.len(), 0);
|
||||
assert!(map.is_empty());
|
||||
|
||||
assert!(map.next().await.is_none());
|
||||
assert!(map.next().await.is_none());
|
||||
|
||||
assert!(map.remove("foo").is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn single_entry() {
|
||||
let mut map = task::spawn(StreamMap::new());
|
||||
let (tx, rx) = mpsc::unbounded_channel();
|
||||
|
||||
assert_ready_none!(map.poll_next());
|
||||
|
||||
assert!(map.insert("foo", rx).is_none());
|
||||
assert!(map.contains_key("foo"));
|
||||
assert!(!map.contains_key("bar"));
|
||||
|
||||
assert_eq!(map.len(), 1);
|
||||
assert!(!map.is_empty());
|
||||
|
||||
assert_pending!(map.poll_next());
|
||||
|
||||
assert_ok!(tx.send(1));
|
||||
|
||||
assert!(map.is_woken());
|
||||
let (k, v) = assert_ready_some!(map.poll_next());
|
||||
assert_eq!(k, "foo");
|
||||
assert_eq!(v, 1);
|
||||
|
||||
assert_pending!(map.poll_next());
|
||||
|
||||
assert_ok!(tx.send(2));
|
||||
|
||||
assert!(map.is_woken());
|
||||
let (k, v) = assert_ready_some!(map.poll_next());
|
||||
assert_eq!(k, "foo");
|
||||
assert_eq!(v, 2);
|
||||
|
||||
assert_pending!(map.poll_next());
|
||||
drop(tx);
|
||||
assert!(map.is_woken());
|
||||
assert_ready_none!(map.poll_next());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn multiple_entries() {
|
||||
let mut map = task::spawn(StreamMap::new());
|
||||
let (tx1, rx1) = mpsc::unbounded_channel();
|
||||
let (tx2, rx2) = mpsc::unbounded_channel();
|
||||
|
||||
map.insert("foo", rx1);
|
||||
map.insert("bar", rx2);
|
||||
|
||||
assert_pending!(map.poll_next());
|
||||
|
||||
assert_ok!(tx1.send(1));
|
||||
|
||||
assert!(map.is_woken());
|
||||
let (k, v) = assert_ready_some!(map.poll_next());
|
||||
assert_eq!(k, "foo");
|
||||
assert_eq!(v, 1);
|
||||
|
||||
assert_pending!(map.poll_next());
|
||||
|
||||
assert_ok!(tx2.send(2));
|
||||
|
||||
assert!(map.is_woken());
|
||||
let (k, v) = assert_ready_some!(map.poll_next());
|
||||
assert_eq!(k, "bar");
|
||||
assert_eq!(v, 2);
|
||||
|
||||
assert_pending!(map.poll_next());
|
||||
|
||||
assert_ok!(tx1.send(3));
|
||||
assert_ok!(tx2.send(4));
|
||||
|
||||
assert!(map.is_woken());
|
||||
|
||||
// Given the randomization, there is no guarantee what order the values will
|
||||
// be received in.
|
||||
let mut v = (0..2)
|
||||
.map(|_| assert_ready_some!(map.poll_next()))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
assert_pending!(map.poll_next());
|
||||
|
||||
v.sort();
|
||||
assert_eq!(v[0].0, "bar");
|
||||
assert_eq!(v[0].1, 4);
|
||||
assert_eq!(v[1].0, "foo");
|
||||
assert_eq!(v[1].1, 3);
|
||||
|
||||
drop(tx1);
|
||||
assert!(map.is_woken());
|
||||
assert_pending!(map.poll_next());
|
||||
drop(tx2);
|
||||
|
||||
assert_ready_none!(map.poll_next());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn insert_remove() {
|
||||
let mut map = task::spawn(StreamMap::new());
|
||||
let (tx, rx) = mpsc::unbounded_channel();
|
||||
|
||||
assert_ready_none!(map.poll_next());
|
||||
|
||||
assert!(map.insert("foo", rx).is_none());
|
||||
let rx = map.remove("foo").unwrap();
|
||||
|
||||
assert_ok!(tx.send(1));
|
||||
|
||||
assert!(!map.is_woken());
|
||||
assert_ready_none!(map.poll_next());
|
||||
|
||||
assert!(map.insert("bar", rx).is_none());
|
||||
|
||||
let v = assert_ready_some!(map.poll_next());
|
||||
assert_eq!(v.0, "bar");
|
||||
assert_eq!(v.1, 1);
|
||||
|
||||
assert!(map.remove("bar").is_some());
|
||||
assert_ready_none!(map.poll_next());
|
||||
|
||||
assert!(map.is_empty());
|
||||
assert_eq!(0, map.len());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn replace() {
|
||||
let mut map = task::spawn(StreamMap::new());
|
||||
let (tx1, rx1) = mpsc::unbounded_channel();
|
||||
let (tx2, rx2) = mpsc::unbounded_channel();
|
||||
|
||||
assert!(map.insert("foo", rx1).is_none());
|
||||
|
||||
assert_pending!(map.poll_next());
|
||||
|
||||
let _rx1 = map.insert("foo", rx2).unwrap();
|
||||
|
||||
assert_pending!(map.poll_next());
|
||||
|
||||
tx1.send(1).unwrap();
|
||||
assert_pending!(map.poll_next());
|
||||
|
||||
tx2.send(2).unwrap();
|
||||
assert!(map.is_woken());
|
||||
let v = assert_ready_some!(map.poll_next());
|
||||
assert_eq!(v.0, "foo");
|
||||
assert_eq!(v.1, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn size_hint_with_upper() {
|
||||
let mut map = StreamMap::new();
|
||||
|
||||
map.insert("a", stream::iter(vec![1]));
|
||||
map.insert("b", stream::iter(vec![1, 2]));
|
||||
map.insert("c", stream::iter(vec![1, 2, 3]));
|
||||
|
||||
assert_eq!(3, map.len());
|
||||
assert!(!map.is_empty());
|
||||
|
||||
let size_hint = map.size_hint();
|
||||
assert_eq!(size_hint, (6, Some(6)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn size_hint_without_upper() {
|
||||
let mut map = StreamMap::new();
|
||||
|
||||
map.insert("a", pin_box(stream::iter(vec![1])));
|
||||
map.insert("b", pin_box(stream::iter(vec![1, 2])));
|
||||
map.insert("c", pin_box(pending()));
|
||||
|
||||
let size_hint = map.size_hint();
|
||||
assert_eq!(size_hint, (3, None));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn new_capacity_zero() {
|
||||
let map = StreamMap::<&str, stream::Pending<()>>::new();
|
||||
assert_eq!(0, map.capacity());
|
||||
|
||||
let keys = map.keys().collect::<Vec<_>>();
|
||||
assert!(keys.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn with_capacity() {
|
||||
let map = StreamMap::<&str, stream::Pending<()>>::with_capacity(10);
|
||||
assert!(10 <= map.capacity());
|
||||
|
||||
let keys = map.keys().collect::<Vec<_>>();
|
||||
assert!(keys.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn iter_keys() {
|
||||
let mut map = StreamMap::new();
|
||||
|
||||
map.insert("a", pending::<i32>());
|
||||
map.insert("b", pending());
|
||||
map.insert("c", pending());
|
||||
|
||||
let mut keys = map.keys().collect::<Vec<_>>();
|
||||
keys.sort();
|
||||
|
||||
assert_eq!(&keys[..], &[&"a", &"b", &"c"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn iter_values() {
|
||||
let mut map = StreamMap::new();
|
||||
|
||||
map.insert("a", stream::iter(vec![1]));
|
||||
map.insert("b", stream::iter(vec![1, 2]));
|
||||
map.insert("c", stream::iter(vec![1, 2, 3]));
|
||||
|
||||
let mut size_hints = map.values().map(|s| s.size_hint().0).collect::<Vec<_>>();
|
||||
|
||||
size_hints.sort();
|
||||
|
||||
assert_eq!(&size_hints[..], &[1, 2, 3]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn iter_values_mut() {
|
||||
let mut map = StreamMap::new();
|
||||
|
||||
map.insert("a", stream::iter(vec![1]));
|
||||
map.insert("b", stream::iter(vec![1, 2]));
|
||||
map.insert("c", stream::iter(vec![1, 2, 3]));
|
||||
|
||||
let mut size_hints = map
|
||||
.values_mut()
|
||||
.map(|s: &mut _| s.size_hint().0)
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
size_hints.sort();
|
||||
|
||||
assert_eq!(&size_hints[..], &[1, 2, 3]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clear() {
|
||||
let mut map = task::spawn(StreamMap::new());
|
||||
|
||||
map.insert("a", stream::iter(vec![1]));
|
||||
map.insert("b", stream::iter(vec![1, 2]));
|
||||
map.insert("c", stream::iter(vec![1, 2, 3]));
|
||||
|
||||
assert_ready_some!(map.poll_next());
|
||||
|
||||
map.clear();
|
||||
|
||||
assert_ready_none!(map.poll_next());
|
||||
assert!(map.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn contains_key_borrow() {
|
||||
let mut map = StreamMap::new();
|
||||
map.insert("foo".to_string(), pending::<()>());
|
||||
|
||||
assert!(map.contains_key("foo"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn one_ready_many_none() {
|
||||
// Run a few times because of randomness
|
||||
for _ in 0..100 {
|
||||
let mut map = task::spawn(StreamMap::new());
|
||||
|
||||
map.insert(0, pin_box(stream::empty()));
|
||||
map.insert(1, pin_box(stream::empty()));
|
||||
map.insert(2, pin_box(stream::once("hello")));
|
||||
map.insert(3, pin_box(stream::pending()));
|
||||
|
||||
let v = assert_ready_some!(map.poll_next());
|
||||
assert_eq!(v, (2, "hello"));
|
||||
}
|
||||
}
|
||||
|
||||
proptest::proptest! {
|
||||
#[test]
|
||||
fn fuzz_pending_complete_mix(kinds: Vec<bool>) {
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
struct DidPoll<T> {
|
||||
did_poll: bool,
|
||||
inner: T,
|
||||
}
|
||||
|
||||
impl<T: Stream + Unpin> Stream for DidPoll<T> {
|
||||
type Item = T::Item;
|
||||
|
||||
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>)
|
||||
-> Poll<Option<T::Item>>
|
||||
{
|
||||
self.did_poll = true;
|
||||
Pin::new(&mut self.inner).poll_next(cx)
|
||||
}
|
||||
}
|
||||
|
||||
for _ in 0..10 {
|
||||
let mut map = task::spawn(StreamMap::new());
|
||||
let mut expect = 0;
|
||||
|
||||
for (i, &is_empty) in kinds.iter().enumerate() {
|
||||
let inner = if is_empty {
|
||||
pin_box(stream::empty::<()>())
|
||||
} else {
|
||||
expect += 1;
|
||||
pin_box(stream::pending::<()>())
|
||||
};
|
||||
|
||||
let stream = DidPoll {
|
||||
did_poll: false,
|
||||
inner,
|
||||
};
|
||||
|
||||
map.insert(i, stream);
|
||||
}
|
||||
|
||||
if expect == 0 {
|
||||
assert_ready_none!(map.poll_next());
|
||||
} else {
|
||||
assert_pending!(map.poll_next());
|
||||
|
||||
assert_eq!(expect, map.values().count());
|
||||
|
||||
for stream in map.values() {
|
||||
assert!(stream.did_poll);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn pin_box<T: Stream<Item = U> + 'static, U>(s: T) -> Pin<Box<dyn Stream<Item = U>>> {
|
||||
Box::pin(s)
|
||||
}
|
||||
Reference in New Issue
Block a user