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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

use futures_core::Stream;
use std::{
    pin::Pin,
    task::{Context, Poll},
};

/// A stream that wraps another stream with a closure that is called once it is dropped.
/// Very useful for libraries that use streams for data transfer and you need to connect
/// when the opposite site drops the connection, thus dropping the stream.
///
/// Example
/// ```
/// use futures::Stream;
/// use drop_stream::DropStream;
///
/// let test_stream = futures::stream::repeat(true);
/// {
///     let wrapped_stream = DropStream::new(test_stream, move || {
///         println!("Stream has been dropped!");
///     });
///
///     let mut wrapped_stream = Box::pin(wrapped_stream);
///
///     let waker = futures::task::noop_waker();
///     let mut context = futures::task::Context::from_waker(&waker);
///     assert_eq!(
///         wrapped_stream.as_mut().poll_next(&mut context),
///         std::task::Poll::Ready(Some(true))
///     );
/// }
/// ```
pub struct DropStream<S: Stream<Item = T> + Unpin, T, U: FnOnce() + Unpin> {
    stream: S,
    // ManuallyDrop used to support FnOnce since ownership of FnOnce needs to be gained in the Drop::drop() method.
    dropper: Option<U>,
}

impl<S: Stream<Item = T> + Unpin, T, U: FnOnce() + Unpin> DropStream<S, T, U> {
    pub fn new(stream: S, dropper: U) -> Self {
        Self {
            stream,
            dropper: Some(dropper),
        }
    }
}

impl<S: Stream<Item = T> + Unpin, T, U: FnOnce() + Unpin> Stream for DropStream<S, T, U> {
    type Item = T;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let stream = Pin::new(&mut self.stream);
        stream.poll_next(cx)
    }
}

impl<S: Stream<Item = T> + Unpin, T, U: FnOnce() + Unpin> Drop for DropStream<S, T, U> {
    fn drop(&mut self) {
        let Some(dropper) = self.dropper.take() else {
            // Only taken in the "drop"-method, and always set in the constructor. Thus it cannot be None here.
            unreachable!()
        };

        dropper()
    }
}

#[cfg(test)]
mod tests {
    use std::task::Poll;

    use crate::DropStream;
    use futures::{stream::repeat, Stream};

    #[test]
    fn dropper_runs_on_drop() {
        let test_stream = repeat(true);

        let mut has_run = false;

        {
            let has_run_ref = &mut has_run;
            let _drop_stream = DropStream::new(test_stream, move || {
                *has_run_ref = true;
            });
        }

        assert!(has_run)
    }

    #[test]
    fn stream_passes_through_result() {
        let test_stream = repeat(true);

        let drop_stream = DropStream::new(test_stream, || {});

        let mut drop_stream = Box::pin(drop_stream);

        let waker = futures::task::noop_waker();
        let mut context = futures::task::Context::from_waker(&waker);
        assert_eq!(
            drop_stream.as_mut().poll_next(&mut context),
            Poll::Ready(Some(true))
        );
    }

    #[test]
    fn dropper_runs_on_drop_after_passing_result() {
        let test_stream = repeat(true);

        let mut has_run = false;

        {
            let has_run_ref = &mut has_run;
            let drop_stream = DropStream::new(test_stream, move || {
                *has_run_ref = true;
            });

            let mut drop_stream = Box::pin(drop_stream);

            let waker = futures::task::noop_waker();
            let mut context = futures::task::Context::from_waker(&waker);
            assert_eq!(
                drop_stream.as_mut().poll_next(&mut context),
                Poll::Ready(Some(true))
            );
        }

        assert!(has_run)
    }
}