Use pin_project_lite instead for throttle

This commit is contained in:
Wouter Geraedts 2019-11-11 13:09:35 +01:00
parent 139a34b685
commit ef958f0408
2 changed files with 20 additions and 19 deletions

View file

@ -315,7 +315,6 @@ extension_trait! {
TakeWhile::new(self, predicate) TakeWhile::new(self, predicate)
} }
#[cfg(all(feature = "default", feature = "unstable"))]
#[doc = r#" #[doc = r#"
Limit the amount of items yielded per timeslice in a stream. Limit the amount of items yielded per timeslice in a stream.
@ -342,6 +341,8 @@ extension_trait! {
# }) } # }) }
``` ```
"#] "#]
#[cfg(all(feature = "default", feature = "unstable"))]
#[cfg_attr(feature = "docs", doc(cfg(unstable)))]
fn throttle(self, d: Duration) -> Throttle<Self> fn throttle(self, d: Duration) -> Throttle<Self>
where where
Self: Sized, Self: Sized,

View file

@ -3,26 +3,25 @@ use std::pin::Pin;
use std::time::Duration; use std::time::Duration;
use futures_timer::Delay; use futures_timer::Delay;
use pin_project_lite::pin_project;
use crate::stream::Stream; use crate::stream::Stream;
use crate::task::{Context, Poll}; use crate::task::{Context, Poll};
pin_project! {
/// A stream that only yields one element once every `duration`, and applies backpressure. Does not drop any elements. /// A stream that only yields one element once every `duration`, and applies backpressure. Does not drop any elements.
/// #[doc(hidden)] #[doc(hidden)]
#[allow(missing_debug_implementations)] #[allow(missing_debug_implementations)]
pub struct Throttle<S> { pub struct Throttle<S> {
#[pin]
stream: S, stream: S,
duration: Duration, duration: Duration,
#[pin]
delay: Option<Delay>, delay: Option<Delay>,
} }
}
impl<S: Unpin> Unpin for Throttle<S> {}
impl<S: Stream> Throttle<S> { impl<S: Stream> Throttle<S> {
pin_utils::unsafe_pinned!(stream: S);
pin_utils::unsafe_unpinned!(duration: Duration);
pin_utils::unsafe_pinned!(delay: Option<Delay>);
pub(super) fn new(stream: S, duration: Duration) -> Self { pub(super) fn new(stream: S, duration: Duration) -> Self {
Throttle { Throttle {
stream, stream,
@ -35,23 +34,24 @@ impl<S: Stream> Throttle<S> {
impl<S: Stream> Stream for Throttle<S> { impl<S: Stream> Stream for Throttle<S> {
type Item = S::Item; type Item = S::Item;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<S::Item>> { fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<S::Item>> {
if let Some(d) = self.as_mut().delay().as_pin_mut() { let mut this = self.project();
if let Some(d) = this.delay.as_mut().as_pin_mut() {
if d.poll(cx).is_ready() { if d.poll(cx).is_ready() {
*self.as_mut().delay() = None; this.delay.set(None);
} else { } else {
return Poll::Pending; return Poll::Pending;
} }
} }
match self.as_mut().stream().poll_next(cx) { match this.stream.poll_next(cx) {
Poll::Pending => { Poll::Pending => {
cx.waker().wake_by_ref(); // Continue driving even though emitting Pending cx.waker().wake_by_ref(); // Continue driving even though emitting Pending
Poll::Pending Poll::Pending
} }
Poll::Ready(None) => Poll::Ready(None), Poll::Ready(None) => Poll::Ready(None),
Poll::Ready(Some(v)) => { Poll::Ready(Some(v)) => {
*self.as_mut().delay() = Some(Delay::new(self.duration)); this.delay.set(Some(Delay::new(*this.duration)));
Poll::Ready(Some(v)) Poll::Ready(Some(v))
} }
} }