futures_time/stream/
buffer.rs

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
131
132
133
134
135
136
137
138
139
140
141
use std::mem;
use std::pin::Pin;

use pin_project_lite::pin_project;

use core::task::{Context, Poll};
use futures_core::stream::Stream;

pin_project! {
    /// Buffer items and flushes them at each interval.
    ///
    /// This `struct` is created by the [`buffer`] method on [`StreamExt`]. See its
    /// documentation for more.
    ///
    /// [`buffer`]: crate::stream::StreamExt::buffer
    /// [`StreamExt`]: crate::stream::StreamExt
    #[derive(Debug)]
    #[must_use = "streams do nothing unless polled or .awaited"]
    pub struct Buffer<S: Stream, I> {
        #[pin]
        stream: S,
        #[pin]
        interval: I,
        slot: Vec<S::Item>,
        state: State,
    }
}

impl<S: Stream, I> Buffer<S, I> {
    pub(crate) fn new(stream: S, interval: I) -> Self {
        Self {
            stream,
            interval,
            slot: vec![],
            state: State::Streaming,
        }
    }
}

#[derive(Debug)]
enum State {
    /// The underlying stream is yielding items.
    Streaming,
    /// The underlying stream is done yielding items.
    StreamDone,
    /// All timers have completed and all data has been yielded.
    TimerDone,
    /// The closing `Ready(None)` has been yielded.
    AllDone,
}

impl<S: Stream, I: Stream> Stream for Buffer<S, I> {
    type Item = Vec<S::Item>;

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

        match this.state {
            // The underlying stream is yielding items.
            State::Streaming => {
                // Poll the underlying stream until we get to `Poll::Pending`.
                loop {
                    match this.stream.as_mut().poll_next(cx) {
                        Poll::Ready(Some(value)) => this.slot.push(value),
                        Poll::Ready(None) => {
                            *this.state = State::StreamDone;
                            break;
                        }
                        Poll::Pending => break,
                    }
                }

                // After the stream, always poll the interval timer.
                this.interval.as_mut().poll_next(cx).map(move |_| {
                    if let State::StreamDone = this.state {
                        *this.state = State::TimerDone;
                        cx.waker().wake_by_ref();
                    }
                    Some(mem::take(&mut *this.slot))
                })
            }

            // The underlying stream is done yielding items.
            State::StreamDone => this.interval.as_mut().poll_next(cx).map(|_| {
                cx.waker().wake_by_ref();
                *this.state = State::TimerDone;
                Some(mem::take(&mut *this.slot))
            }),

            // All timers have completed and all data has been yielded.
            State::TimerDone => {
                *this.state = State::AllDone;
                Poll::Ready(None)
            }

            // The closing `Ready(None)` has been yielded.
            State::AllDone => panic!("stream polled after completion"),
        }
    }
}

#[cfg(test)]
mod test {
    use crate::prelude::*;
    use crate::time::Duration;
    use futures_lite::prelude::*;

    #[test]
    fn buffer_all_values() {
        async_io::block_on(async {
            let interval = Duration::from_millis(5);
            let buffer = Duration::from_millis(20);

            let mut counter = 0;
            crate::stream::interval(interval)
                .take(10)
                .buffer(buffer)
                .for_each(|buf| counter += buf.len())
                .await;

            assert_eq!(counter, 10);
        })
    }

    #[test]
    fn no_debounces_hit() {
        async_io::block_on(async {
            let interval = Duration::from_millis(20);
            let buffer = Duration::from_millis(10);

            let mut counter = 0;
            crate::stream::interval(interval)
                .take(10)
                .buffer(buffer)
                .for_each(|buf| counter += buf.len())
                .await;

            assert_eq!(counter, 10);
        })
    }
}