Skip to content

Commit a239350

Browse files
committed
Implemented StreamExt::throttle
1 parent ca80ca9 commit a239350

File tree

2 files changed

+66
-0
lines changed

2 files changed

+66
-0
lines changed

src/stream/stream/mod.rs

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ mod skip_while;
5050
mod step_by;
5151
mod take;
5252
mod take_while;
53+
mod throttle;
5354
mod try_fold;
5455
mod try_for_each;
5556
mod zip;
@@ -86,10 +87,12 @@ pub use skip_while::SkipWhile;
8687
pub use step_by::StepBy;
8788
pub use take::Take;
8889
pub use take_while::TakeWhile;
90+
pub use throttle::Throttle;
8991
pub use zip::Zip;
9092

9193
use std::cmp::Ordering;
9294
use std::marker::PhantomData;
95+
use std::time::Duration;
9396

9497
use cfg_if::cfg_if;
9598

@@ -288,6 +291,13 @@ extension_trait! {
288291
TakeWhile::new(self, predicate)
289292
}
290293

294+
fn throttle(self, d: Duration) -> Throttle<Self>
295+
where
296+
Self: Sized,
297+
{
298+
Throttle::new(self, d)
299+
}
300+
291301
#[doc = r#"
292302
Creates a stream that yields each `step`th element.
293303

src/stream/stream/throttle.rs

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
use std::future::Future;
2+
use std::pin::Pin;
3+
use std::time::Duration;
4+
5+
use futures_timer::Delay;
6+
7+
use crate::stream::Stream;
8+
use crate::task::{Context, Poll};
9+
10+
/// A stream that only yields one element once every `duration`, and drops all others.
11+
/// #[doc(hidden)]
12+
#[allow(missing_debug_implementations)]
13+
pub struct Throttle<S> {
14+
stream: S,
15+
duration: Duration,
16+
delay: Option<Delay>,
17+
}
18+
19+
impl<S: Unpin> Unpin for Throttle<S> {}
20+
21+
impl<S: Stream> Throttle<S> {
22+
pin_utils::unsafe_pinned!(stream: S);
23+
pin_utils::unsafe_unpinned!(duration: Duration);
24+
pin_utils::unsafe_pinned!(delay: Option<Delay>);
25+
26+
pub(super) fn new(stream: S, duration: Duration) -> Self {
27+
Throttle {
28+
stream,
29+
duration,
30+
delay: None,
31+
}
32+
}
33+
}
34+
35+
impl<S: Stream> Stream for Throttle<S> {
36+
type Item = S::Item;
37+
38+
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<S::Item>> {
39+
match self.as_mut().stream().poll_next(cx) {
40+
Poll::Ready(v) => match self.as_mut().delay().as_pin_mut() {
41+
None => {
42+
*self.as_mut().delay() = Some(Delay::new(self.duration));
43+
Poll::Ready(v)
44+
}
45+
Some(d) => match d.poll(cx) {
46+
Poll::Ready(_) => {
47+
*self.as_mut().delay() = Some(Delay::new(self.duration));
48+
Poll::Ready(v)
49+
}
50+
Poll::Pending => Poll::Pending,
51+
},
52+
},
53+
Poll::Pending => Poll::Pending,
54+
}
55+
}
56+
}

0 commit comments

Comments
 (0)