vmbus_channel/
stream.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! `Stream` combinators shared across the vmbus channel/client stack.
5
6use core::pin::Pin;
7use core::task::Context;
8use core::task::Poll;
9use futures::StreamExt;
10
11/// A [`futures::Stream`] adapter that tags each item produced by the
12/// inner stream with a caller-supplied `T`. Yields `(tag, Some(item))`
13/// for each inner item, then a single `(tag, None)` sentinel when the
14/// inner stream ends, and `None` thereafter.
15#[derive(Debug)]
16pub struct TaggedStream<T, S>(Option<T>, S);
17
18impl<T: Clone, S: futures::Stream + Unpin> TaggedStream<T, S> {
19    /// Wraps `s` with the tag `t`.
20    pub fn new(t: T, s: S) -> Self {
21        Self(Some(t), s)
22    }
23
24    /// Returns the tag, or `None` if the inner stream has terminated
25    /// and the terminal `(tag, None)` sentinel has already been
26    /// yielded.
27    pub fn value(&self) -> Option<&T> {
28        self.0.as_ref()
29    }
30}
31
32impl<T: Clone, S: futures::Stream + Unpin> futures::Stream for TaggedStream<T, S>
33where
34    Self: Unpin,
35{
36    type Item = (T, Option<S::Item>);
37
38    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
39        let this = self.get_mut();
40        if let Some(t) = this.0.clone() {
41            let v = core::task::ready!(this.1.poll_next_unpin(cx));
42            if v.is_none() {
43                // Return `None` next time poll_next is called.
44                this.0 = None;
45            }
46            Poll::Ready(Some((t, v)))
47        } else {
48            Poll::Ready(None)
49        }
50    }
51}
vmbus_channel/stream.rs

vmbus_channel/
stream.rs
