tor_circmgr/mgr/
streams.rs

1//! Stream utilities to help implement
2//! [`AbstractCircMgr`](`super::AbstractCircMgr.`)
3
4use futures::stream::{Fuse, FusedStream, Stream, StreamExt};
5use futures::task::{Context, Poll};
6use pin_project::pin_project;
7use std::pin::Pin;
8
9/// Enumeration to indicate which of two streams provided a result.
10#[derive(Debug, Copy, Clone, Eq, PartialEq)]
11pub(super) enum Source {
12    /// Indicates a result coming from the left (preferred) stream.
13    Left,
14    /// Indicates a result coming from the right (secondary) stream.
15    Right,
16}
17
18/// A stream returned by [`select_biased`]
19///
20/// See that function for more documentation.
21#[pin_project]
22pub(super) struct SelectBiased<S, T> {
23    /// Preferred underlying stream.
24    ///
25    /// When results are available from both streams, we always yield them
26    /// from this one.  When this stream is exhausted, the `SelectBiased`
27    /// is exhausted too.
28    #[pin]
29    left: Fuse<S>,
30    /// Secondary underlying stream.
31    #[pin]
32    right: Fuse<T>,
33}
34
35/// Combine two instances of [`Stream`] into one.
36///
37/// This function is similar to [`futures::stream::select`], but differs
38/// in that it treats the two underlying streams asymmetrically.  Specifically:
39///
40///  * Each result is labeled with [`Source::Left`] or
41///    [`Source::Right`], depending on which of the two streams it came
42///    from.
43///  * If both the "left" and the "right" stream are ready, we always
44///    prefer the left stream.
45///  * We stop iterating over this stream when there are no more
46///    results on the left stream, regardless whether the right stream
47///    is exhausted or not.
48///
49/// # Future plans
50///
51/// This might need a better name, especially if we use it anywhere
52/// else.
53///
54/// If we do expose this function, we might want to split up the ways in
55/// which it differs from `select`.
56pub(super) fn select_biased<S, T>(left: S, right: T) -> SelectBiased<S, T>
57where
58    S: Stream,
59    T: Stream<Item = S::Item>,
60{
61    SelectBiased {
62        left: left.fuse(),
63        right: right.fuse(),
64    }
65}
66
67impl<S, T> FusedStream for SelectBiased<S, T>
68where
69    S: Stream,
70    T: Stream<Item = S::Item>,
71{
72    fn is_terminated(&self) -> bool {
73        // We're done if the left stream is done, whether the right stream
74        // is done or not.
75        self.left.is_terminated()
76    }
77}
78
79impl<S, T> Stream for SelectBiased<S, T>
80where
81    S: Stream,
82    T: Stream<Item = S::Item>,
83{
84    type Item = (Source, S::Item);
85
86    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
87        let this = self.project();
88        // We always check the left stream first.
89        match this.left.poll_next(cx) {
90            Poll::Ready(Some(val)) => {
91                // The left stream has an item: yield it.
92                return Poll::Ready(Some((Source::Left, val)));
93            }
94            Poll::Ready(None) => {
95                // The left stream is exhausted: don't even check the right.
96                return Poll::Ready(None);
97            }
98            Poll::Pending => {}
99        }
100
101        // The left stream is pending: see whether the right stream has
102        // anything to say.
103        match this.right.poll_next(cx) {
104            Poll::Ready(Some(val)) => {
105                // The right stream has an item: yield it.
106                Poll::Ready(Some((Source::Right, val)))
107            }
108            _ => {
109                // The right stream is exhausted or pending: in either case,
110                // we need to wait.
111                Poll::Pending
112            }
113        }
114    }
115}
116
117#[cfg(test)]
118mod test {
119    // @@ begin test lint list maintained by maint/add_warning @@
120    #![allow(clippy::bool_assert_comparison)]
121    #![allow(clippy::clone_on_copy)]
122    #![allow(clippy::dbg_macro)]
123    #![allow(clippy::mixed_attributes_style)]
124    #![allow(clippy::print_stderr)]
125    #![allow(clippy::print_stdout)]
126    #![allow(clippy::single_char_pattern)]
127    #![allow(clippy::unwrap_used)]
128    #![allow(clippy::unchecked_time_subtraction)]
129    #![allow(clippy::useless_vec)]
130    #![allow(clippy::needless_pass_by_value)]
131    #![allow(clippy::string_slice)] // See arti#2571
132    //! <!-- @@ end test lint list maintained by maint/add_warning @@ -->
133    use super::*;
134    use futures_await_test::async_test;
135
136    // Tests where only elements from the left stream should be yielded.
137    #[async_test]
138    async fn left_only() {
139        use Source::Left as L;
140        use futures::stream::iter;
141        // If there's nothing in the right stream, we just yield the left.
142        let left = vec![1_usize, 2, 3];
143        let right = vec![];
144
145        let s = select_biased(iter(left), iter(right));
146        let result: Vec<_> = s.collect().await;
147        assert_eq!(result, vec![(L, 1_usize), (L, 2), (L, 3)]);
148
149        // If the left runs out (which this will), we don't yield anything
150        // from the right.
151        let left = vec![1_usize, 2, 3];
152        let right = vec![4, 5, 6];
153        let s = select_biased(iter(left), iter(right));
154        let result: Vec<_> = s.collect().await;
155        assert_eq!(result, vec![(L, 1_usize), (L, 2), (L, 3)]);
156
157        // The same thing happens if the left stream is completely empty!
158        let left = vec![];
159        let right = vec![4_usize, 5, 6];
160        let s = select_biased(iter(left), iter(right));
161        let result: Vec<_> = s.collect().await;
162        assert_eq!(result, vec![]);
163    }
164
165    // Tests where only elements from the right stream should be yielded.
166    #[async_test]
167    async fn right_only() {
168        use Source::Right as R;
169        use futures::stream::{iter, pending};
170
171        // Try a forever-pending stream for the left hand side.
172        let left = pending();
173        let right = vec![4_usize, 5, 6];
174        let mut s = select_biased(left, iter(right));
175        assert_eq!(s.next().await, Some((R, 4)));
176        assert_eq!(s.next().await, Some((R, 5)));
177        assert_eq!(s.next().await, Some((R, 6)));
178    }
179
180    // Tests where we can find elements from both streams.
181    #[async_test]
182    async fn multiplex() {
183        use Source::{Left as L, Right as R};
184        use futures::SinkExt;
185
186        let (mut snd_l, rcv_l) = futures::channel::mpsc::channel(5);
187        let (mut snd_r, rcv_r) = futures::channel::mpsc::channel(5);
188        let mut s = select_biased(rcv_l, rcv_r);
189
190        snd_l.send(1_usize).await.unwrap();
191        snd_r.send(4_usize).await.unwrap();
192        snd_l.send(2_usize).await.unwrap();
193
194        assert_eq!(s.next().await, Some((L, 1)));
195        assert_eq!(s.next().await, Some((L, 2)));
196        assert_eq!(s.next().await, Some((R, 4)));
197
198        snd_r.send(5_usize).await.unwrap();
199        snd_l.send(3_usize).await.unwrap();
200
201        assert!(!s.is_terminated());
202        drop(snd_r);
203
204        assert_eq!(s.next().await, Some((L, 3)));
205        assert_eq!(s.next().await, Some((R, 5)));
206
207        drop(snd_l);
208        assert_eq!(s.next().await, None);
209
210        assert!(s.is_terminated());
211    }
212}
tor_circmgr/mgr/streams.rs

tor_circmgr/mgr/
streams.rs
