src/lib/async-utils/src/stream/future_map.rs - fuchsia - Git at Google

 // Copyright 2020 The Fuchsia Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.

 use {
     futures::{
         stream::{FusedStream, Stream},
         Future,
     },
     std::{
         collections::HashMap,
         hash::Hash,
         pin::Pin,
         task::{Context, Poll},
     },
 };

 /// A collection of Future indexed by key, allowing removal by Key. When polled, a FutureMap yields
 /// from whichever member future is ready first.
 /// The Future type `Fut` can be `?Unpin`, as all futures are stored as pins inside the map. The Key
 /// type `K` must be `Unpin`; it is unlikely that an `!Unpin` type would ever be needed as a Key.
 /// FutureMap yields items of type Fut::Output.
 pub struct FutureMap<K, Fut> {
     inner: HashMap<K, Pin<Box<Fut>>>,
     is_terminated: bool,
 }

 impl<K, Fut> Default for FutureMap<K, Fut> {
     fn default() -> Self {
         Self { inner: Default::default(), is_terminated: false }
     }
 }

 impl<K: Unpin, Fut> Unpin for FutureMap<K, Fut> {}

 impl<K: Eq + Hash + Unpin, Fut: Future> FutureMap<K, Fut> {
     /// Returns an empty `FutureMap`.
     pub fn new() -> Self {
         Self::default()
     }

     /// Insert a future identified by `key` to the map.
     ///
     /// This method will not call `poll` on the submitted stream. The caller must ensure
     /// that `poll_next` is called in order to receive wake-up notifications for the given
     /// stream.
     pub fn insert(&mut self, key: K, future: Fut) -> Option<Pin<Box<Fut>>> {
         let Self { inner, is_terminated } = self;
         *is_terminated = false;
         inner.insert(key, Box::new(future).into())
     }

     /// Returns `true` if the `FutureMap` contains `key`.
     pub fn contains_key(&self, key: &K) -> bool {
         self.inner.contains_key(key)
     }

     /// Remove the future identified by `key`, returning it if it exists.
     pub fn remove(&mut self, key: &K) -> Option<Pin<Box<Fut>>> {
         self.inner.remove(key)
     }

     /// Provide mutable access to the inner hashmap.
     /// This is safe as if the future were being polled, we would not be able to access a mutable
     /// reference to self to pass to this method.
     pub fn inner(&mut self) -> &mut HashMap<K, Pin<Box<Fut>>> {
         &mut self.inner
     }
 }

 impl<K: Clone + Eq + Hash + Unpin, Fut: Future> Stream for FutureMap<K, Fut> {
     type Item = Fut::Output;

     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         // We can pull the inner value out as FutureMap is `Unpin`
         let Self { inner, is_terminated } = Pin::into_inner(self);

         if inner.is_empty() {
             *is_terminated = true;
             Poll::Ready(None)
         } else {
             match inner.iter_mut().find_map(|(key, future)| match Pin::new(future).poll(cx) {
                 Poll::Ready(req) => Some((key.clone(), req)),
                 Poll::Pending => None,
             }) {
                 Some((key, req)) => {
                     assert!(inner.remove(&key).is_some());
                     Poll::Ready(Some(req))
                 }
                 None => Poll::Pending,
             }
         }
     }
 }

 impl<K: Clone + Eq + Hash + Unpin, Fut: Future> FusedStream for FutureMap<K, Fut> {
     fn is_terminated(&self) -> bool {
         let Self { inner: _, is_terminated } = self;
         *is_terminated
     }
 }

 #[cfg(test)]
 mod test {
     //! We validate the behavior of the FutureMap stream by enumerating all possible external
     //! events, and then generating permutations of valid sequences of those events. These model
     //! the possible executions sequences the stream could go through in program execution. We
     //! then assert that:
     //!   a) At all points during execution, all invariants are held
     //!   b) The final result is as expected
     //!
     //! In this case, the invariants are:
     //!   * If the map is empty, it is pending
     //!   * If all futures are pending, the map is pending
     //!   * otherwise the map is ready
     //!
     //! The result is:
     //!   * All test messages have been injected
     //!   * All test messages have been yielded
     //!   * All test futures have terminated
     //!   * No event is yielded with a given key after the future for that key has terminated
     //!
     //! Together these show:
     //!   * Progress is always eventually made - the Stream cannot be stalled
     //!   * All inserted elements will eventually be yielded
     //!   * Elements are never duplicated
     use {
         super::*,
         crate::stream::WithTag,
         futures::{channel::oneshot, StreamExt},
         proptest::prelude::*,
         std::{collections::HashSet, fmt::Debug},
     };

     /// Possible actions to take in evaluating the stream
     enum Event<K> {
         /// Insert a new future
         InsertFuture(K, oneshot::Receiver<Result<u64, ()>>),
         /// Send a value, completing a future.
         CompleteFuture(K, oneshot::Sender<Result<u64, ()>>),
         /// Schedule the executor. The executor will only run the task if awoken, otherwise it will
         /// do nothing
         Execute,
     }

     impl<K: Debug> Debug for Event<K> {
         fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
             match self {
                 Event::InsertFuture(k, _) => write!(f, "InsertFuture({:?})", k),
                 Event::CompleteFuture(k, _) => write!(f, "SendRequest({:?})", k),
                 Event::Execute => write!(f, "Execute"),
             }
         }
     }

     fn stream_events<K: Clone + Eq + Hash>(key: K) -> Vec<Event<K>> {
         let (sender, receiver) = oneshot::channel::<Result<u64, ()>>();
         vec![Event::InsertFuture(key.clone(), receiver), Event::CompleteFuture(key, sender)]
     }

     /// Determine how many events are sent on open channels (a channel is open if it has not been
     /// closed, even if it has not yet been inserted into the FutureMap)
     fn expected_yield<K: Eq + Hash>(events: &Vec<Event<K>>) -> usize {
         events
             .iter()
             .fold((HashSet::new(), 0), |(mut terminated, closed), event| match event {
                 Event::CompleteFuture(k, _) => {
                     assert!(
                         !terminated.contains(k),
                         "There should be no more than one future per key"
                     );
                     let _: bool = terminated.insert(k);
                     (terminated, closed + 1)
                 }
                 _ => (terminated, closed),
             })
             .1
     }

     /// Strategy that produces random permutations of a set of events, corresponding to inserting,
     /// sending and completing up to n different streams in random order, also interspersed with
     /// running the executor
     fn execution_sequences(n: u64) -> impl Strategy<Value = Vec<Event<u64>>> {
         fn generate_events(n: u64) -> Vec<Event<u64>> {
             let mut events = (0..n).flat_map(|n| stream_events(n)).collect::<Vec<_>>();
             events.extend(std::iter::repeat_with(|| Event::Execute).take((n * 3) as usize));
             events
         }

         // We want to produce random permutations of these events
         (0..n).prop_map(generate_events).prop_shuffle()
     }

     proptest! {
         #[test]
         fn test_invariants(mut execution in execution_sequences(4)) {
             let expected = expected_yield(&execution);

             // Add enough execution events to ensure we will complete, no matter the order
             execution.extend(std::iter::repeat_with(|| Event::Execute).take((expected * 3) as usize));

             let (waker, count) = futures_test::task::new_count_waker();
             let mut futures = FutureMap::new();
             let expected = expected as u64;
             let mut next_wake = 0;
             let mut yielded = 0;
             let mut inserted = 0;
             let mut events = vec![];
             for event in execution {
                 match event {
                     Event::InsertFuture(key, future) => {
                         assert_matches::assert_matches!(futures.insert(key, future.tagged(key)), None);
                         // FutureMap does *not* wake on inserting new futures, matching the
                         // behavior of streams::SelectAll. The client *must* arrange for it to be
                         // polled again after a future is inserted; we model that here by forcing a
                         // wake up
                         next_wake = count.get();
                     }
                     Event::CompleteFuture(_, sender) => {
                         prop_assert_eq!(sender.send(Ok(1)), Ok(()));
                         inserted = inserted + 1;
                     }
                     Event::Execute if count.get() >= next_wake => {
                         match Pin::new(&mut futures.next()).poll(&mut Context::from_waker(&waker)) {
                             Poll::Ready(Some((k, v))) => {
                                 events.push((k, v));
                                 yielded = yielded + 1;
                                 // Ensure that we wake up next time;
                                 next_wake = count.get();
                                 // Invariant: future(k) must be in the map
                                 prop_assert!(!futures.contains_key(&k))
                             }
                             Poll::Ready(None) => {
                                 // // the Stream impl for FutureMap never completes
                                 // unreachable!()
                                 prop_assert!(futures.inner.is_empty());
                             }
                             Poll::Pending => {
                                 next_wake = count.get() + 1;
                             }
                         };
                     }
                     Event::Execute => (),
                 }
             }
             prop_assert_eq!(inserted, expected, "All expected requests inserted");
             prop_assert_eq!((next_wake, count.get(), yielded), (next_wake, count.get(), expected), "All expected requests yielded");
             let all_keys = 0..expected;
             for k in all_keys {
                 prop_assert!(!futures.contains_key(&k), "All futures should now have been removed");
             }
         }
     }
 }
	// Copyright 2020 The Fuchsia Authors. All rights reserved.
	// Use of this source code is governed by a BSD-style license that can be
	// found in the LICENSE file.

	use {
	futures::{
	stream::{FusedStream, Stream},
	Future,
	},
	std::{
	collections::HashMap,
	hash::Hash,
	pin::Pin,
	task::{Context, Poll},
	},
	};

	/// A collection of Future indexed by key, allowing removal by Key. When polled, a FutureMap yields
	/// from whichever member future is ready first.
	/// The Future type `Fut` can be `?Unpin`, as all futures are stored as pins inside the map. The Key
	/// type `K` must be `Unpin`; it is unlikely that an `!Unpin` type would ever be needed as a Key.
	/// FutureMap yields items of type Fut::Output.
	pub struct FutureMap<K, Fut> {
	inner: HashMap<K, Pin<Box<Fut>>>,
	is_terminated: bool,
	}

	impl<K, Fut> Default for FutureMap<K, Fut> {
	fn default() -> Self {
	Self { inner: Default::default(), is_terminated: false }
	}
	}

	impl<K: Unpin, Fut> Unpin for FutureMap<K, Fut> {}

	impl<K: Eq + Hash + Unpin, Fut: Future> FutureMap<K, Fut> {
	/// Returns an empty `FutureMap`.
	pub fn new() -> Self {
	Self::default()
	}

	/// Insert a future identified by `key` to the map.
	///
	/// This method will not call `poll` on the submitted stream. The caller must ensure
	/// that `poll_next` is called in order to receive wake-up notifications for the given
	/// stream.
	pub fn insert(&mut self, key: K, future: Fut) -> Option<Pin<Box<Fut>>> {
	let Self { inner, is_terminated } = self;
	*is_terminated = false;
	inner.insert(key, Box::new(future).into())
	}

	/// Returns `true` if the `FutureMap` contains `key`.
	pub fn contains_key(&self, key: &K) -> bool {
	self.inner.contains_key(key)
	}

	/// Remove the future identified by `key`, returning it if it exists.
	pub fn remove(&mut self, key: &K) -> Option<Pin<Box<Fut>>> {
	self.inner.remove(key)
	}

	/// Provide mutable access to the inner hashmap.
	/// This is safe as if the future were being polled, we would not be able to access a mutable
	/// reference to self to pass to this method.
	pub fn inner(&mut self) -> &mut HashMap<K, Pin<Box<Fut>>> {
	&mut self.inner
	}
	}

	impl<K: Clone + Eq + Hash + Unpin, Fut: Future> Stream for FutureMap<K, Fut> {
	type Item = Fut::Output;

	fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
	// We can pull the inner value out as FutureMap is `Unpin`
	let Self { inner, is_terminated } = Pin::into_inner(self);

	if inner.is_empty() {
	*is_terminated = true;
	Poll::Ready(None)
	} else {
	match inner.iter_mut().find_map(\|(key, future)\| match Pin::new(future).poll(cx) {
	Poll::Ready(req) => Some((key.clone(), req)),
	Poll::Pending => None,
	}) {
	Some((key, req)) => {
	assert!(inner.remove(&key).is_some());
	Poll::Ready(Some(req))
	}
	None => Poll::Pending,
	}
	}
	}
	}

	impl<K: Clone + Eq + Hash + Unpin, Fut: Future> FusedStream for FutureMap<K, Fut> {
	fn is_terminated(&self) -> bool {
	let Self { inner: _, is_terminated } = self;
	*is_terminated
	}
	}

	#[cfg(test)]
	mod test {
	//! We validate the behavior of the FutureMap stream by enumerating all possible external
	//! events, and then generating permutations of valid sequences of those events. These model
	//! the possible executions sequences the stream could go through in program execution. We
	//! then assert that:
	//! a) At all points during execution, all invariants are held
	//! b) The final result is as expected
	//!
	//! In this case, the invariants are:
	//! * If the map is empty, it is pending
	//! * If all futures are pending, the map is pending
	//! * otherwise the map is ready
	//!
	//! The result is:
	//! * All test messages have been injected
	//! * All test messages have been yielded
	//! * All test futures have terminated
	//! * No event is yielded with a given key after the future for that key has terminated
	//!
	//! Together these show:
	//! * Progress is always eventually made - the Stream cannot be stalled
	//! * All inserted elements will eventually be yielded
	//! * Elements are never duplicated
	use {
	super::*,
	crate::stream::WithTag,
	futures::{channel::oneshot, StreamExt},
	proptest::prelude::*,
	std::{collections::HashSet, fmt::Debug},
	};

	/// Possible actions to take in evaluating the stream
	enum Event<K> {
	/// Insert a new future
	InsertFuture(K, oneshot::Receiver<Result<u64, ()>>),
	/// Send a value, completing a future.
	CompleteFuture(K, oneshot::Sender<Result<u64, ()>>),
	/// Schedule the executor. The executor will only run the task if awoken, otherwise it will
	/// do nothing
	Execute,
	}

	impl<K: Debug> Debug for Event<K> {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
	match self {
	Event::InsertFuture(k, _) => write!(f, "InsertFuture({:?})", k),
	Event::CompleteFuture(k, _) => write!(f, "SendRequest({:?})", k),
	Event::Execute => write!(f, "Execute"),
	}
	}
	}

	fn stream_events<K: Clone + Eq + Hash>(key: K) -> Vec<Event<K>> {
	let (sender, receiver) = oneshot::channel::<Result<u64, ()>>();
	vec![Event::InsertFuture(key.clone(), receiver), Event::CompleteFuture(key, sender)]
	}

	/// Determine how many events are sent on open channels (a channel is open if it has not been
	/// closed, even if it has not yet been inserted into the FutureMap)
	fn expected_yield<K: Eq + Hash>(events: &Vec<Event<K>>) -> usize {
	events
	.iter()
	.fold((HashSet::new(), 0), \|(mut terminated, closed), event\| match event {
	Event::CompleteFuture(k, _) => {
	assert!(
	!terminated.contains(k),
	"There should be no more than one future per key"
	);
	let _: bool = terminated.insert(k);
	(terminated, closed + 1)
	}
	_ => (terminated, closed),
	})
	.1
	}

	/// Strategy that produces random permutations of a set of events, corresponding to inserting,
	/// sending and completing up to n different streams in random order, also interspersed with
	/// running the executor
	fn execution_sequences(n: u64) -> impl Strategy<Value = Vec<Event<u64>>> {
	fn generate_events(n: u64) -> Vec<Event<u64>> {
	let mut events = (0..n).flat_map(\|n\| stream_events(n)).collect::<Vec<_>>();
	events.extend(std::iter::repeat_with(\|\| Event::Execute).take((n * 3) as usize));
	events
	}

	// We want to produce random permutations of these events
	(0..n).prop_map(generate_events).prop_shuffle()
	}

	proptest! {
	#[test]
	fn test_invariants(mut execution in execution_sequences(4)) {
	let expected = expected_yield(&execution);

	// Add enough execution events to ensure we will complete, no matter the order
	execution.extend(std::iter::repeat_with(\|\| Event::Execute).take((expected * 3) as usize));

	let (waker, count) = futures_test::task::new_count_waker();
	let mut futures = FutureMap::new();
	let expected = expected as u64;
	let mut next_wake = 0;
	let mut yielded = 0;
	let mut inserted = 0;
	let mut events = vec![];
	for event in execution {
	match event {
	Event::InsertFuture(key, future) => {
	assert_matches::assert_matches!(futures.insert(key, future.tagged(key)), None);
	// FutureMap does not wake on inserting new futures, matching the
	// behavior of streams::SelectAll. The client must arrange for it to be
	// polled again after a future is inserted; we model that here by forcing a
	// wake up
	next_wake = count.get();
	}
	Event::CompleteFuture(_, sender) => {
	prop_assert_eq!(sender.send(Ok(1)), Ok(()));
	inserted = inserted + 1;
	}
	Event::Execute if count.get() >= next_wake => {
	match Pin::new(&mut futures.next()).poll(&mut Context::from_waker(&waker)) {
	Poll::Ready(Some((k, v))) => {
	events.push((k, v));
	yielded = yielded + 1;
	// Ensure that we wake up next time;
	next_wake = count.get();
	// Invariant: future(k) must be in the map
	prop_assert!(!futures.contains_key(&k))
	}
	Poll::Ready(None) => {
	// // the Stream impl for FutureMap never completes
	// unreachable!()
	prop_assert!(futures.inner.is_empty());
	}
	Poll::Pending => {
	next_wake = count.get() + 1;
	}
	};
	}
	Event::Execute => (),
	}
	}
	prop_assert_eq!(inserted, expected, "All expected requests inserted");
	prop_assert_eq!((next_wake, count.get(), yielded), (next_wake, count.get(), expected), "All expected requests yielded");
	let all_keys = 0..expected;
	for k in all_keys {
	prop_assert!(!futures.contains_key(&k), "All futures should now have been removed");
	}
	}
	}
	}