futures_bounded/
futures_tuple_set.rs

1use std::collections::HashMap;
2use std::pin::Pin;
3use std::task::{ready, Context, Poll};
4
5use crate::{AnyFuture, BoxFuture, Delay, FuturesMap, PushError, Timeout};
6
7/// Represents a list of tuples of a [Future] and an associated piece of data.
8///
9/// Each future must finish within the specified time and the list never outgrows its capacity.
10pub struct FuturesTupleSet<O, D> {
11    id: u32,
12    inner: FuturesMap<u32, O>,
13    data: HashMap<u32, D>,
14}
15
16impl<O, D> FuturesTupleSet<O, D> {
17    pub fn new(make_delay: impl Fn() -> Delay + Send + Sync + 'static, capacity: usize) -> Self {
18        Self {
19            id: 0,
20            inner: FuturesMap::new(make_delay, capacity),
21            data: HashMap::new(),
22        }
23    }
24}
25
26impl<O, D> FuturesTupleSet<O, D>
27where
28    O: 'static,
29{
30    /// Push a future into the list.
31    ///
32    /// This method adds the given future to the list.
33    /// If the length of the list is equal to the capacity, this method returns a error that contains the passed future.
34    /// In that case, the future is not added to the set.
35    pub fn try_push<F>(&mut self, future: F, data: D) -> Result<(), (BoxFuture<O>, D)>
36    where
37        F: AnyFuture<Output = O>,
38    {
39        self.id = self.id.wrapping_add(1);
40
41        match self.inner.try_push(self.id, future) {
42            Ok(()) => {}
43            Err(PushError::BeyondCapacity(w)) => return Err((w, data)),
44            Err(PushError::Replaced(_)) => unreachable!("we never reuse IDs"),
45        }
46        self.data.insert(self.id, data);
47
48        Ok(())
49    }
50
51    pub fn len(&self) -> usize {
52        self.inner.len()
53    }
54
55    pub fn is_empty(&self) -> bool {
56        self.inner.is_empty()
57    }
58
59    pub fn poll_ready_unpin(&mut self, cx: &mut Context<'_>) -> Poll<()> {
60        self.inner.poll_ready_unpin(cx)
61    }
62
63    pub fn poll_unpin(&mut self, cx: &mut Context<'_>) -> Poll<(Result<O, Timeout>, D)> {
64        let (id, res) = ready!(self.inner.poll_unpin(cx));
65        let data = self.data.remove(&id).expect("must have data for future");
66
67        Poll::Ready((res, data))
68    }
69
70    /// Returns an iterator over all futures of type `T` pushed via [`FuturesTupleSet::try_push`].
71    /// This iterator returns futures in an arbitrary order, which may change.
72    ///
73    /// If downcasting a future to `T` fails it will be skipped in the iterator.
74    pub fn iter_of_type<T>(&self) -> impl Iterator<Item = (Pin<&T>, &D)>
75    where
76        T: 'static,
77    {
78        self.inner
79            .iter_of_type()
80            .map(|(id, item)| (item, self.data.get(id).expect("must have data for future")))
81    }
82
83    /// Returns an iterator with mutable access over all futures of type `T`
84    /// pushed via [`FuturesTupleSet::try_push`].
85    /// This iterator returns futures in an arbitrary order, which may change.
86    ///
87    /// If downcasting a future to `T` fails it will be skipped in the iterator.
88    pub fn iter_mut_of_type<T>(&mut self) -> impl Iterator<Item = (Pin<&mut T>, &mut D)>
89    where
90        T: 'static,
91    {
92        // We need to convince the borrow checker that we are only returning one mutable reference
93        // to each data item. We can do this efficiently by creating a hashmap of mutable
94        // references, then removing each item from the hashmap as we iterate over the futures.
95        let mut mut_refs_data = self
96            .data
97            .iter_mut()
98            .map(|(id, data)| (*id, data))
99            .collect::<HashMap<u32, &mut D>>();
100
101        self.inner.iter_mut_of_type().map(move |(id, item)| {
102            let data = mut_refs_data
103                .remove(id)
104                .expect("each future has a unique id, must have data for future");
105            (item, data)
106        })
107    }
108}
109
110#[cfg(all(test, feature = "futures-timer"))]
111mod tests {
112    use super::*;
113    use futures::channel::oneshot;
114    use futures_util::future::poll_fn;
115    use futures_util::FutureExt;
116    use std::future::ready;
117    use std::time::Duration;
118
119    #[test]
120    fn tracks_associated_data_of_future() {
121        let mut set = FuturesTupleSet::new(|| Delay::futures_timer(Duration::from_secs(10)), 10);
122
123        let _ = set.try_push(ready(1), 1);
124        let _ = set.try_push(ready(2), 2);
125
126        let (res1, data1) = poll_fn(|cx| set.poll_unpin(cx)).now_or_never().unwrap();
127        let (res2, data2) = poll_fn(|cx| set.poll_unpin(cx)).now_or_never().unwrap();
128
129        assert_eq!(res1.unwrap(), data1);
130        assert_eq!(res2.unwrap(), data2);
131    }
132
133    #[test]
134    fn can_iter_typed_futures() {
135        const N: usize = 10;
136        let mut futures =
137            FuturesTupleSet::new(|| Delay::futures_timer(Duration::from_millis(100)), N);
138        let mut sender = Vec::with_capacity(N);
139        for i in 0..N {
140            let (tx, rx) = oneshot::channel::<()>();
141            futures
142                .try_push(rx, format!("Data{i}"))
143                .map_err(|_| PushError::BeyondCapacity(()))
144                .unwrap();
145            sender.push(tx);
146        }
147        assert_eq!(futures.iter_of_type::<oneshot::Receiver<()>>().count(), N);
148        for (i, (_, data)) in futures.iter_of_type::<oneshot::Receiver<()>>().enumerate() {
149            let expect_data = format!("Data{}", N - i - 1); // Reverse order.
150            assert_eq!(data, &expect_data);
151        }
152        assert!(!sender.iter().any(|tx| tx.is_canceled()));
153
154        for (mut rx, _) in futures.iter_mut_of_type::<oneshot::Receiver<()>>() {
155            rx.close();
156        }
157        assert!(sender.iter().all(|tx| tx.is_canceled()));
158
159        // Deliberately try a non-matching type
160        assert_eq!(futures.iter_mut_of_type::<()>().count(), 0);
161    }
162}
futures_bounded/futures_tuple_set.rs

futures_bounded/
futures_tuple_set.rs
