stream_multiplexer 0.16.1

Combines many streams into a few.
Documentation 
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
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96

use crate::*;

use async_channel as channel;
use futures_util::future::FutureExt;
use futures_util::stream::{StreamExt, StreamFuture, TryStream};

type ReadStream<T> = StreamDropper<T>;

#[derive(Debug)]
pub(crate) struct Channel<T> {
    id: ChannelId,
    add_rx: channel::Receiver<ReadStream<T>>,
    stream_futures: FuturesUnordered<StreamFuture<ReadStream<T>>>,
}

impl<T> Channel<T>
where
    T: TryStream + Unpin,
{
    pub(crate) fn new(id: ChannelId) -> (channel::Sender<ReadStream<T>>, Self) {
        let (add_tx, add_rx) = channel::unbounded::<ReadStream<T>>();
        let stream_futures = FuturesUnordered::new();
        (
            add_tx,
            Self {
                add_rx,
                id,
                stream_futures,
            },
        )
    }

    pub(crate) async fn next(&mut self) -> (StreamId, Option<T::Item>) {
        log::trace!("next() for {}", self.id);
        // select on the stream-add channel and the FuturesUnordered
        loop {
            if self.stream_futures.is_empty() {
                log::trace!("channel {} is empty, awaiting streams.", self.id);
                let added_stream_res = self.add_rx.recv().await;
                // return the erorr if there was one
                self.process_stream_add(added_stream_res);
            } else {
                log::trace!("channel {} is awaiting streams or packets", self.id);
                futures_util::select_biased! {
                    added_stream_res = self.add_rx.recv().fuse() => {
                        log::trace!("channel {} adding new stream", self.id);
                        self.process_stream_add(added_stream_res);
                    }
                    stream_output = self.stream_futures.next().fuse() => {
                        log::trace!("channel {} processing stream data", self.id);
                        return self.process_stream_output(stream_output.unwrap()).await;
                    }
                };
            }
        }
    }

    fn process_stream_add(
        &mut self,
        added_stream: Result<ReadStream<T>, async_channel::RecvError>,
    ) {
        // Unwrapping result as it should not be possbile to have the other end of the channel
        // closed, we hold both ends.
        let added_stream = added_stream.expect("Add channel was closed for adding.");

        let stream_id = added_stream.id;
        if added_stream.stream.is_some() {
            // Add to the FuturesUnordered
            self.stream_futures.push(added_stream.into_future());
        } else {
            log::warn!(
                "Trying to enqueue dropped stream {} into channel {}",
                stream_id,
                self.id
            );
        }
    }

    async fn process_stream_output(
        &self,
        (output, stream): (
            Option<T::Item>, // TryStream
            ReadStream<T>,
        ),
    ) -> (StreamId, Option<T::Item>) {
        let stream_id = stream.id;

        // re-enqueue the stream if it's viable
        if output.is_some() && stream.stream.is_some() {
            self.stream_futures.push(stream.into_future());
        }

        // return the bytes from the stream
        (stream_id, output)
    }
}