1 //! Slow down a stream by enforcing a delay between items.
2
3 use crate::Stream;
4 use tokio::time::{Duration, Instant, Sleep};
5
6 use std::future::Future;
7 use std::marker::Unpin;
8 use std::pin::Pin;
9 use std::task::{self, Poll};
10
11 use pin_project_lite::pin_project;
12
throttle<T>(duration: Duration, stream: T) -> Throttle<T> where T: Stream,13 pub(super) fn throttle<T>(duration: Duration, stream: T) -> Throttle<T>
14 where
15 T: Stream,
16 {
17 Throttle {
18 delay: tokio::time::sleep_until(Instant::now() + duration),
19 duration,
20 has_delayed: true,
21 stream,
22 }
23 }
24
25 pin_project! {
26 /// Stream for the [`throttle`](throttle) function. This object is `!Unpin`. If you need it to
27 /// implement `Unpin` you can pin your throttle like this: `Box::pin(your_throttle)`.
28 #[derive(Debug)]
29 #[must_use = "streams do nothing unless polled"]
30 pub struct Throttle<T> {
31 #[pin]
32 delay: Sleep,
33 duration: Duration,
34
35 // Set to true when `delay` has returned ready, but `stream` hasn't.
36 has_delayed: bool,
37
38 // The stream to throttle
39 #[pin]
40 stream: T,
41 }
42 }
43
44 // XXX: are these safe if `T: !Unpin`?
45 impl<T: Unpin> Throttle<T> {
46 /// Acquires a reference to the underlying stream that this combinator is
47 /// pulling from.
get_ref(&self) -> &T48 pub fn get_ref(&self) -> &T {
49 &self.stream
50 }
51
52 /// Acquires a mutable reference to the underlying stream that this combinator
53 /// is pulling from.
54 ///
55 /// Note that care must be taken to avoid tampering with the state of the stream
56 /// which may otherwise confuse this combinator.
get_mut(&mut self) -> &mut T57 pub fn get_mut(&mut self) -> &mut T {
58 &mut self.stream
59 }
60
61 /// Consumes this combinator, returning the underlying stream.
62 ///
63 /// Note that this may discard intermediate state of this combinator, so care
64 /// should be taken to avoid losing resources when this is called.
into_inner(self) -> T65 pub fn into_inner(self) -> T {
66 self.stream
67 }
68 }
69
70 impl<T: Stream> Stream for Throttle<T> {
71 type Item = T::Item;
72
poll_next(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Option<Self::Item>>73 fn poll_next(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Option<Self::Item>> {
74 let mut me = self.project();
75 let dur = *me.duration;
76
77 if !*me.has_delayed && !is_zero(dur) {
78 ready!(me.delay.as_mut().poll(cx));
79 *me.has_delayed = true;
80 }
81
82 let value = ready!(me.stream.poll_next(cx));
83
84 if value.is_some() {
85 if !is_zero(dur) {
86 me.delay.reset(Instant::now() + dur);
87 }
88
89 *me.has_delayed = false;
90 }
91
92 Poll::Ready(value)
93 }
94 }
95
is_zero(dur: Duration) -> bool96 fn is_zero(dur: Duration) -> bool {
97 dur == Duration::from_millis(0)
98 }
99