stream_cancel/
wrapper.rs

1use crate::{StreamExt, TakeUntilIf, Trigger, Tripwire};
2use futures_core::stream::Stream;
3use pin_project::pin_project;
4use std::future::Future;
5use std::pin::Pin;
6use std::task::{Context, Poll};
7
8/// A `Valve` is associated with a [`Trigger`], and can be used to wrap one or more
9/// asynchronous streams. All streams wrapped by a given `Valve` (or its clones) will be
10/// interrupted when [`Trigger::close`] is called on the valve's associated handle.
11#[pin_project]
12#[derive(Clone, Debug)]
13pub struct Valve(#[pin] Tripwire);
14
15impl Valve {
16    /// Make a new `Valve` and an associated [`Trigger`].
17    pub fn new() -> (Trigger, Self) {
18        let (t, tw) = Tripwire::new();
19        (t, Valve(tw))
20    }
21
22    /// Wrap the given `stream` with this `Valve`.
23    ///
24    /// When [`Trigger::close`] is called on the handle associated with this valve, the given
25    /// stream will immediately yield `None`.
26    pub fn wrap<S>(&self, stream: S) -> Valved<S>
27    where
28        S: Stream,
29    {
30        Valved(stream.take_until_if(self.0.clone()))
31    }
32
33    /// Check if the valve has been closed.
34    ///
35    /// If `Ready`, contains `true` if the stream should be considered closed, and `false`
36    /// if the `Trigger` has been disabled.
37    pub fn poll_closed(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<bool> {
38        self.project().0.poll(cx)
39    }
40}
41
42/// A `Valved` is wrapper around a `Stream` that enables the stream to be turned off remotely to
43/// initiate a graceful shutdown. When a new `Valved` is created with [`Valved::new`], a handle to
44/// that `Valved` is also produced; when [`Trigger::close`] is called on that handle, the
45/// wrapped stream will immediately yield `None` to indicate that it has completed.
46#[derive(Clone, Debug)]
47pub struct Valved<S>(TakeUntilIf<S, Tripwire>);
48
49impl<S> Valved<S> {
50    /// Make the given stream cancellable.
51    ///
52    /// To cancel the stream, call [`Trigger::close`] on the returned handle.
53    pub fn new(stream: S) -> (Trigger, Self)
54    where
55        S: Stream,
56    {
57        let (vh, v) = Valve::new();
58        (vh, v.wrap(stream))
59    }
60
61    /// Consumes this wrapper, returning the underlying stream.
62    pub fn into_inner(self) -> S {
63        self.0.into_inner()
64    }
65}
66
67impl<S> Stream for Valved<S>
68where
69    S: Stream,
70{
71    type Item = S::Item;
72
73    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
74        // safe since we never move nor leak &mut
75        let inner = unsafe { self.map_unchecked_mut(|s| &mut s.0) };
76        inner.poll_next(cx)
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83    use futures_util::stream::empty;
84    #[test]
85    fn valved_stream_may_be_dropped_safely() {
86        let _orphan = {
87            let s = empty::<()>();
88            let (trigger, valve) = Valve::new();
89            let _wrapped = valve.wrap(s);
90            trigger
91        };
92    }
93}
stream_cancel/wrapper.rs

stream_cancel/
wrapper.rs
