futures_bounded/
stream_set.rs

1use std::{
2    pin::Pin,
3    task::{ready, Context, Poll},
4};
5
6use crate::{AnyStream, BoxStream, Delay, PushError, StreamMap, Timeout};
7
8/// Represents a set of [Stream]s.
9///
10/// Each stream must finish within the specified time and the list never outgrows its capacity.
11pub struct StreamSet<O> {
12    id: u32,
13    inner: StreamMap<u32, O>,
14}
15
16impl<O> StreamSet<O> {
17    pub fn new(make_delay: impl Fn() -> Delay + Send + Sync + 'static, capacity: usize) -> Self {
18        Self {
19            id: 0,
20            inner: StreamMap::new(make_delay, capacity),
21        }
22    }
23}
24
25impl<O> StreamSet<O>
26where
27    O: Send + 'static,
28{
29    /// Push a stream into the list.
30    ///
31    /// This method adds the given stream to the list.
32    /// If the length of the list is equal to the capacity, this method returns a error that contains the passed stream.
33    /// In that case, the stream is not added to the set.
34    pub fn try_push<F>(&mut self, stream: F) -> Result<(), BoxStream<O>>
35    where
36        F: AnyStream<Item = O>,
37    {
38        self.id = self.id.wrapping_add(1);
39
40        match self.inner.try_push(self.id, stream) {
41            Ok(()) => Ok(()),
42            Err(PushError::BeyondCapacity(w)) => Err(w),
43            Err(PushError::Replaced(_)) => unreachable!("we never reuse IDs"),
44        }
45    }
46
47    pub fn len(&self) -> usize {
48        self.inner.len()
49    }
50
51    pub fn is_empty(&self) -> bool {
52        self.inner.is_empty()
53    }
54
55    pub fn poll_ready_unpin(&mut self, cx: &mut Context<'_>) -> Poll<()> {
56        self.inner.poll_ready_unpin(cx)
57    }
58
59    pub fn poll_next_unpin(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<O, Timeout>>> {
60        let (_, res) = ready!(self.inner.poll_next_unpin(cx));
61
62        Poll::Ready(res)
63    }
64
65    /// Returns an iterator over all streams of type `T` pushed via [`StreamSet::try_push`].
66    /// This iterator returns streams in an arbitrary order, which may change.
67    ///
68    /// If downcasting a stream to `T` fails it will be skipped in the iterator.
69    pub fn iter_of_type<T>(&self) -> impl Iterator<Item = Pin<&T>>
70    where
71        T: 'static,
72    {
73        self.inner.iter_of_type().map(|(_, item)| item)
74    }
75
76    /// Returns an iterator with mutable access over all streams of type `T`
77    /// pushed via [`StreamSet::try_push`].
78    /// This iterator returns streams in an arbitrary order, which may change.
79    ///
80    /// If downcasting a stream to `T` fails it will be skipped in the iterator.
81    pub fn iter_mut_of_type<T>(&mut self) -> impl Iterator<Item = Pin<&mut T>>
82    where
83        T: 'static,
84    {
85        self.inner.iter_mut_of_type().map(|(_, item)| item)
86    }
87}
futures_bounded/stream_set.rs

futures_bounded/
stream_set.rs
