1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64

use futures_util::stream::BoxStream;
use futures_util::Stream;
use std::task::{ready, Context, Poll};
use std::time::Duration;

use crate::{PushError, StreamMap, Timeout};

/// Represents a set of [Stream]s.
///
/// Each stream must finish within the specified time and the list never outgrows its capacity.
pub struct StreamSet<O> {
    id: u32,
    inner: StreamMap<u32, O>,
}

impl<O> StreamSet<O> {
    pub fn new(timeout: Duration, capacity: usize) -> Self {
        Self {
            id: 0,
            inner: StreamMap::new(timeout, capacity),
        }
    }
}

impl<O> StreamSet<O>
where
    O: Send + 'static,
{
    /// Push a stream into the list.
    ///
    /// This method adds the given stream to the list.
    /// If the length of the list is equal to the capacity, this method returns a error that contains the passed stream.
    /// In that case, the stream is not added to the set.
    pub fn try_push<F>(&mut self, stream: F) -> Result<(), BoxStream<O>>
    where
        F: Stream<Item = O> + Send + 'static,
    {
        self.id = self.id.wrapping_add(1);

        match self.inner.try_push(self.id, stream) {
            Ok(()) => Ok(()),
            Err(PushError::BeyondCapacity(w)) => Err(w),
            Err(PushError::Replaced(_)) => unreachable!("we never reuse IDs"),
        }
    }

    pub fn len(&self) -> usize {
        self.inner.len()
    }

    pub fn is_empty(&self) -> bool {
        self.inner.is_empty()
    }

    pub fn poll_ready_unpin(&mut self, cx: &mut Context<'_>) -> Poll<()> {
        self.inner.poll_ready_unpin(cx)
    }

    pub fn poll_next_unpin(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<O, Timeout>>> {
        let (_, res) = ready!(self.inner.poll_next_unpin(cx));

        Poll::Ready(res)
    }
}