futures_bounded/
stream_map.rs

1use std::any::Any;
2use std::mem;
3use std::pin::Pin;
4use std::task::{Context, Poll, Waker};
5use std::time::Duration;
6
7use futures_util::stream::SelectAll;
8use futures_util::{stream, FutureExt, Stream, StreamExt};
9
10use crate::{AnyStream, BoxStream, Delay, PushError, Timeout};
11
12/// Represents a map of [`Stream`]s.
13///
14/// Each stream must finish within the specified time and the map never outgrows its capacity.
15pub struct StreamMap<ID, O> {
16    make_delay: Box<dyn Fn() -> Delay + Send + Sync>,
17    capacity: usize,
18    inner: SelectAll<TaggedStream<ID, TimeoutStream<BoxStream<O>>>>,
19    empty_waker: Option<Waker>,
20    full_waker: Option<Waker>,
21}
22
23impl<ID, O> StreamMap<ID, O>
24where
25    ID: Clone + Unpin,
26{
27    pub fn new(make_delay: impl Fn() -> Delay + Send + Sync + 'static, capacity: usize) -> Self {
28        Self {
29            make_delay: Box::new(make_delay),
30            capacity,
31            inner: Default::default(),
32            empty_waker: None,
33            full_waker: None,
34        }
35    }
36}
37
38impl<ID, O> StreamMap<ID, O>
39where
40    ID: Clone + PartialEq + Send + Unpin + 'static,
41    O: Send + 'static,
42{
43    /// Push a stream into the map.
44    pub fn try_push<F>(&mut self, id: ID, stream: F) -> Result<(), PushError<BoxStream<O>>>
45    where
46        F: AnyStream<Item = O>,
47    {
48        if self.inner.len() >= self.capacity {
49            return Err(PushError::BeyondCapacity(Box::pin(stream)));
50        }
51
52        if let Some(waker) = self.empty_waker.take() {
53            waker.wake();
54        }
55
56        let old = self.remove(id.clone());
57        self.inner.push(TaggedStream::new(
58            id,
59            TimeoutStream {
60                inner: Box::pin(stream),
61                timeout: (self.make_delay)(),
62            },
63        ));
64
65        match old {
66            None => Ok(()),
67            Some(old) => Err(PushError::Replaced(old)),
68        }
69    }
70
71    pub fn remove(&mut self, id: ID) -> Option<BoxStream<O>> {
72        let tagged = self.inner.iter_mut().find(|s| s.key == id)?;
73
74        let inner = mem::replace(&mut tagged.inner.inner, Box::pin(stream::pending()));
75        tagged.exhausted = true; // Setting this will emit `None` on the next poll and ensure `SelectAll` cleans up the resources.
76
77        Some(inner)
78    }
79
80    pub fn len(&self) -> usize {
81        self.inner.len()
82    }
83
84    pub fn is_empty(&self) -> bool {
85        self.inner.is_empty()
86    }
87
88    #[allow(unknown_lints, clippy::needless_pass_by_ref_mut)] // &mut Context is idiomatic.
89    pub fn poll_ready_unpin(&mut self, cx: &mut Context<'_>) -> Poll<()> {
90        if self.inner.len() < self.capacity {
91            return Poll::Ready(());
92        }
93
94        self.full_waker = Some(cx.waker().clone());
95
96        Poll::Pending
97    }
98
99    pub fn poll_next_unpin(
100        &mut self,
101        cx: &mut Context<'_>,
102    ) -> Poll<(ID, Option<Result<O, Timeout>>)> {
103        match futures_util::ready!(self.inner.poll_next_unpin(cx)) {
104            None => {
105                self.empty_waker = Some(cx.waker().clone());
106                Poll::Pending
107            }
108            Some((id, Some(Ok(output)))) => Poll::Ready((id, Some(Ok(output)))),
109            Some((id, Some(Err(dur)))) => {
110                self.remove(id.clone()); // Remove stream, otherwise we keep reporting the timeout.
111
112                Poll::Ready((id, Some(Err(Timeout::new(dur)))))
113            }
114            Some((id, None)) => Poll::Ready((id, None)),
115        }
116    }
117
118    /// Returns an iterator over all streams of type `T` pushed via [`StreamMap::try_push`].
119    /// This iterator returns streams in an arbitrary order, which may change.
120    ///
121    /// If downcasting a stream to `T` fails it will be skipped in the iterator.
122    pub fn iter_of_type<T>(&self) -> impl Iterator<Item = (&ID, Pin<&T>)>
123    where
124        T: 'static,
125    {
126        self.inner.iter().filter_map(|a| {
127            let pointer = a.inner.inner.as_ref().get_ref();
128
129            let any = pointer as &(dyn Any + Send);
130            let inner = any.downcast_ref::<T>()?;
131
132            // Safety: The pointer is already pinned.
133            let pinned = unsafe { Pin::new_unchecked(inner) };
134
135            Some((&a.key, pinned))
136        })
137    }
138
139    /// Returns an iterator with mutable access over all streams of type `T`
140    /// pushed via [`StreamMap::try_push`].
141    /// This iterator returns streams in an arbitrary order, which may change.
142    ///
143    /// If downcasting a stream to `T` fails it will be skipped in the iterator.
144    pub fn iter_mut_of_type<T>(&mut self) -> impl Iterator<Item = (&ID, Pin<&mut T>)>
145    where
146        T: 'static,
147    {
148        self.inner.iter_mut().filter_map(|a| {
149            let pin = a.inner.inner.as_mut();
150
151            // Safety: We are only temporarily manipulating the pointer and pinning it again further down.
152            let pointer = unsafe { Pin::into_inner_unchecked(pin) };
153            let any = pointer as &mut (dyn Any + Send);
154            let inner = any.downcast_mut::<T>()?;
155
156            // Safety: The pointer is already pinned.
157            let pinned = unsafe { Pin::new_unchecked(inner) };
158
159            Some((&a.key, pinned))
160        })
161    }
162}
163
164struct TimeoutStream<S> {
165    inner: S,
166    timeout: Delay,
167}
168
169impl<F> Stream for TimeoutStream<F>
170where
171    F: Stream + Unpin,
172{
173    type Item = Result<F::Item, Duration>;
174
175    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
176        if let Poll::Ready(dur) = self.timeout.poll_unpin(cx) {
177            return Poll::Ready(Some(Err(dur)));
178        }
179
180        self.inner.poll_next_unpin(cx).map(|a| a.map(Ok))
181    }
182}
183
184struct TaggedStream<K, S> {
185    key: K,
186    inner: S,
187
188    exhausted: bool,
189}
190
191impl<K, S> TaggedStream<K, S> {
192    fn new(key: K, inner: S) -> Self {
193        Self {
194            key,
195            inner,
196            exhausted: false,
197        }
198    }
199}
200
201impl<K, S> Stream for TaggedStream<K, S>
202where
203    K: Clone + Unpin,
204    S: Stream + Unpin,
205{
206    type Item = (K, Option<S::Item>);
207
208    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
209        if self.exhausted {
210            return Poll::Ready(None);
211        }
212
213        match futures_util::ready!(self.inner.poll_next_unpin(cx)) {
214            Some(item) => Poll::Ready(Some((self.key.clone(), Some(item)))),
215            None => {
216                self.exhausted = true;
217
218                Poll::Ready(Some((self.key.clone(), None)))
219            }
220        }
221    }
222}
223
224#[cfg(all(test, feature = "futures-timer"))]
225mod tests {
226    use futures::channel::mpsc;
227    use futures_util::stream::{once, pending};
228    use futures_util::SinkExt;
229    use std::future::{poll_fn, ready, Future};
230    use std::pin::Pin;
231    use std::time::Instant;
232
233    use super::*;
234
235    #[test]
236    fn cannot_push_more_than_capacity_tasks() {
237        let mut streams = StreamMap::new(|| Delay::futures_timer(Duration::from_secs(10)), 1);
238
239        assert!(streams.try_push("ID_1", once(ready(()))).is_ok());
240        matches!(
241            streams.try_push("ID_2", once(ready(()))),
242            Err(PushError::BeyondCapacity(_))
243        );
244    }
245
246    #[test]
247    fn cannot_push_the_same_id_few_times() {
248        let mut streams = StreamMap::new(|| Delay::futures_timer(Duration::from_secs(10)), 5);
249
250        assert!(streams.try_push("ID", once(ready(()))).is_ok());
251        matches!(
252            streams.try_push("ID", once(ready(()))),
253            Err(PushError::Replaced(_))
254        );
255    }
256
257    #[tokio::test]
258    async fn streams_timeout() {
259        let mut streams = StreamMap::new(|| Delay::futures_timer(Duration::from_millis(100)), 1);
260
261        let _ = streams.try_push("ID", pending::<()>());
262        futures_timer::Delay::new(Duration::from_millis(150)).await;
263        let (_, result) = poll_fn(|cx| streams.poll_next_unpin(cx)).await;
264
265        assert!(result.unwrap().is_err())
266    }
267
268    #[tokio::test]
269    async fn timed_out_stream_gets_removed() {
270        let mut streams = StreamMap::new(|| Delay::futures_timer(Duration::from_millis(100)), 1);
271
272        let _ = streams.try_push("ID", pending::<()>());
273        futures_timer::Delay::new(Duration::from_millis(150)).await;
274        poll_fn(|cx| streams.poll_next_unpin(cx)).await;
275
276        let poll = streams.poll_next_unpin(&mut Context::from_waker(
277            futures_util::task::noop_waker_ref(),
278        ));
279        assert!(poll.is_pending())
280    }
281
282    #[test]
283    fn removing_stream() {
284        let mut streams = StreamMap::new(|| Delay::futures_timer(Duration::from_millis(100)), 1);
285
286        let _ = streams.try_push("ID", stream::once(ready(())));
287
288        {
289            let cancelled_stream = streams.remove("ID");
290            assert!(cancelled_stream.is_some());
291        }
292
293        let poll = streams.poll_next_unpin(&mut Context::from_waker(
294            futures_util::task::noop_waker_ref(),
295        ));
296
297        assert!(poll.is_pending());
298        assert_eq!(
299            streams.len(),
300            0,
301            "resources of cancelled streams are cleaned up properly"
302        );
303    }
304
305    #[tokio::test]
306    async fn replaced_stream_is_still_registered() {
307        let mut streams = StreamMap::new(|| Delay::futures_timer(Duration::from_millis(100)), 3);
308
309        let (mut tx1, rx1) = mpsc::channel(5);
310        let (mut tx2, rx2) = mpsc::channel(5);
311
312        let _ = streams.try_push("ID1", rx1);
313        let _ = streams.try_push("ID2", rx2);
314
315        let _ = tx2.send(2).await;
316        let _ = tx1.send(1).await;
317        let _ = tx2.send(3).await;
318        let (id, res) = poll_fn(|cx| streams.poll_next_unpin(cx)).await;
319        assert_eq!(id, "ID1");
320        assert_eq!(res.unwrap().unwrap(), 1);
321        let (id, res) = poll_fn(|cx| streams.poll_next_unpin(cx)).await;
322        assert_eq!(id, "ID2");
323        assert_eq!(res.unwrap().unwrap(), 2);
324        let (id, res) = poll_fn(|cx| streams.poll_next_unpin(cx)).await;
325        assert_eq!(id, "ID2");
326        assert_eq!(res.unwrap().unwrap(), 3);
327
328        let (mut new_tx1, new_rx1) = mpsc::channel(5);
329        let replaced = streams.try_push("ID1", new_rx1);
330        assert!(matches!(replaced.unwrap_err(), PushError::Replaced(_)));
331
332        let _ = new_tx1.send(4).await;
333        let (id, res) = poll_fn(|cx| streams.poll_next_unpin(cx)).await;
334
335        assert_eq!(id, "ID1");
336        assert_eq!(res.unwrap().unwrap(), 4);
337    }
338
339    // Each stream emits 1 item with delay, `Task` only has a capacity of 1, meaning they must be processed in sequence.
340    // We stop after NUM_STREAMS tasks, meaning the overall execution must at least take DELAY * NUM_STREAMS.
341    #[tokio::test]
342    async fn backpressure() {
343        const DELAY: Duration = Duration::from_millis(100);
344        const NUM_STREAMS: u32 = 10;
345
346        let start = Instant::now();
347        Task::new(DELAY, NUM_STREAMS, 1).await;
348        let duration = start.elapsed();
349
350        assert!(duration >= DELAY * NUM_STREAMS);
351    }
352
353    #[test]
354    fn can_iter_typed_streams() {
355        const N: usize = 10;
356        let mut streams = StreamMap::new(|| Delay::futures_timer(Duration::from_millis(100)), N);
357        let mut sender = Vec::with_capacity(N);
358        for i in 0..N {
359            let (tx, rx) = mpsc::channel::<()>(1);
360            streams.try_push(format!("ID{i}"), rx).unwrap();
361            sender.push(tx);
362        }
363        assert_eq!(streams.iter_of_type::<mpsc::Receiver<()>>().count(), N);
364        for (i, (id, _)) in streams.iter_of_type::<mpsc::Receiver<()>>().enumerate() {
365            let expect_id = format!("ID{}", N - i - 1); // Reverse order.
366            assert_eq!(id, &expect_id);
367        }
368        assert!(!sender.iter().any(|tx| tx.is_closed()));
369
370        for (_, mut rx) in streams.iter_mut_of_type::<mpsc::Receiver<()>>() {
371            rx.close();
372        }
373        assert!(sender.iter().all(|tx| tx.is_closed()));
374
375        // Deliberately try a non-matching type
376        assert_eq!(streams.iter_mut_of_type::<()>().count(), 0);
377    }
378
379    struct Task {
380        item_delay: Duration,
381        num_streams: usize,
382        num_processed: usize,
383        inner: StreamMap<u8, ()>,
384    }
385
386    impl Task {
387        fn new(item_delay: Duration, num_streams: u32, capacity: usize) -> Self {
388            Self {
389                item_delay,
390                num_streams: num_streams as usize,
391                num_processed: 0,
392                inner: StreamMap::new(|| Delay::futures_timer(Duration::from_secs(60)), capacity),
393            }
394        }
395    }
396
397    impl Future for Task {
398        type Output = ();
399
400        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
401            let this = self.get_mut();
402
403            while this.num_processed < this.num_streams {
404                match this.inner.poll_next_unpin(cx) {
405                    Poll::Ready((_, Some(result))) => {
406                        if result.is_err() {
407                            panic!("Timeout is great than item delay")
408                        }
409
410                        this.num_processed += 1;
411                        continue;
412                    }
413                    Poll::Ready((_, None)) => {
414                        continue;
415                    }
416                    _ => {}
417                }
418
419                if let Poll::Ready(()) = this.inner.poll_ready_unpin(cx) {
420                    // We push the constant ID to prove that user can use the same ID if the stream was finished
421                    let maybe_future = this
422                        .inner
423                        .try_push(1u8, once(futures_timer::Delay::new(this.item_delay)));
424                    assert!(maybe_future.is_ok(), "we polled for readiness");
425
426                    continue;
427                }
428
429                return Poll::Pending;
430            }
431
432            Poll::Ready(())
433        }
434    }
435}
futures_bounded/stream_map.rs

futures_bounded/
stream_map.rs
