Implement throttle combinator (#736)

Throttle down a stream by enforcing a fixed delay between items.

src/util/stream.rs

@@ -1,4 +1,7 @@
-use tokio_timer::Timeout;
+use tokio_timer::{
+    throttle::Throttle,
+    Timeout,
+};
 
 use futures::Stream;
 
@@ -19,6 +22,14 @@ use std::time::Duration;
 ///
 /// [`timeout`]: #method.timeout
 pub trait StreamExt: Stream {
+    /// Throttle down the stream by enforcing a fixed delay between items.
+    ///
+    /// Errors are also delayed.
+    fn throttle(self, duration: Duration) -> Throttle<Self>
+    where Self: Sized
+    {
+        Throttle::new(self, duration)
+    }
 
     /// Creates a new stream which allows `self` until `timeout`.
     ///

tokio-timer/src/lib.rs

@@ -9,6 +9,8 @@
 //!
 //! * [`Interval`] A stream that yields at fixed time intervals.
 //!
+//! * [`Throttle`]: Throttle down a stream by enforcing a fixed delay between items.
+//!
 //! * [`Timeout`]: Wraps a future or stream, setting an upper bound to the
 //!   amount of time it is allowed to execute. If the future or stream does not
 //!   complete in time, then it is canceled and an error is returned.
@@ -21,6 +23,7 @@
 //! [`Timer`] instance must be running on some thread.
 //!
 //! [`Delay`]: struct.Delay.html
+//! [`Throttle`]: throttle/struct.Throttle.html
 //! [`Timeout`]: struct.Timeout.html
 //! [`Interval`]: struct.Interval.html
 //! [`Timer`]: timer/struct.Timer.html
@@ -34,6 +37,7 @@ extern crate slab;
 
 pub mod clock;
 pub mod delay_queue;
+pub mod throttle;
 pub mod timeout;
 pub mod timer;
 

tokio-timer/src/throttle.rs

@@ -0,0 +1,170 @@
+//! Slow down a stream by enforcing a delay between items.
+
+use {clock, Delay, Error};
+
+use futures::{Async, Future, Poll, Stream};
+use futures::future::Either;
+
+use std::{
+    error::Error as StdError,
+    fmt::{Display, Formatter, Result as FmtResult},
+    time::Duration,
+};
+
+/// Slow down a stream by enforcing a delay between items.
+#[derive(Debug)]
+#[must_use = "streams do nothing unless polled"]
+pub struct Throttle<T> {
+    delay: Option<Delay>,
+    duration: Duration,
+    stream: T,
+}
+
+/// Either the error of the underlying stream, or an error within
+/// tokio's timing machinery.
+#[derive(Debug)]
+pub struct ThrottleError<T>(Either<T, Error>);
+
+impl<T> Throttle<T> {
+    /// Slow down a stream by enforcing a delay between items.
+    pub fn new(stream: T, duration: Duration) -> Self {
+        Self {
+            delay: None,
+            duration: duration,
+            stream: stream,
+        }
+    }
+
+    /// Acquires a reference to the underlying stream that this combinator is
+    /// pulling from.
+    pub fn get_ref(&self) -> &T {
+        &self.stream
+    }
+
+    /// Acquires a mutable reference to the underlying stream that this combinator
+    /// is pulling from.
+    ///
+    /// Note that care must be taken to avoid tampering with the state of the stream
+    /// which may otherwise confuse this combinator.
+    pub fn get_mut(&mut self) -> &mut T {
+        &mut self.stream
+    }
+
+    /// Consumes this combinator, returning the underlying stream.
+    ///
+    /// Note that this may discard intermediate state of this combinator, so care
+    /// should be taken to avoid losing resources when this is called.
+    pub fn into_inner(self) -> T {
+        self.stream
+    }
+}
+
+impl<T: Stream> Stream for Throttle<T> {
+    type Item = T::Item;
+    type Error = ThrottleError<T::Error>;
+
+    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
+        if let Some(ref mut delay) = self.delay {
+            try_ready!({
+                delay.poll()
+                    .map_err(ThrottleError::from_timer_err)
+            });
+        }
+
+        self.delay = None;
+        let value = try_ready!({
+            self.stream.poll()
+                .map_err(ThrottleError::from_stream_err)
+        });
+
+        if value.is_some() {
+            self.delay = Some(Delay::new(clock::now() + self.duration));
+        }
+
+        Ok(Async::Ready(value))
+    }
+}
+
+impl<T> ThrottleError<T> {
+    /// Creates a new `ThrottleError` from the given stream error.
+    pub fn from_stream_err(err: T) -> Self {
+        ThrottleError(Either::A(err))
+    }
+
+    /// Creates a new `ThrottleError` from the given tokio timer error.
+    pub fn from_timer_err(err: Error) -> Self {
+        ThrottleError(Either::B(err))
+    }
+
+    /// Attempts to get the underlying stream error, if it is present.
+    pub fn get_stream_error(&self) -> Option<&T> {
+        match self.0 {
+            Either::A(ref x) => Some(x),
+            _ => None,
+        }
+    }
+
+    /// Attempts to get the underlying timer error, if it is present.
+    pub fn get_timer_error(&self) -> Option<&Error> {
+        match self.0 {
+            Either::B(ref x) => Some(x),
+            _ => None,
+        }
+    }
+
+    /// Attempts to extract the underlying stream error, if it is present.
+    pub fn into_stream_error(self) -> Option<T> {
+        match self.0 {
+            Either::A(x) => Some(x),
+            _ => None,
+        }
+    }
+
+    /// Attempts to extract the underlying timer error, if it is present.
+    pub fn into_timer_error(self) -> Option<Error> {
+        match self.0 {
+            Either::B(x) => Some(x),
+            _ => None,
+        }
+    }
+
+    /// Returns whether the throttle error has occured because of an error
+    /// in the underlying stream.
+    pub fn is_stream_error(&self) -> bool {
+        !self.is_timer_error()
+    }
+
+    /// Returns whether the throttle error has occured because of an error
+    /// in tokio's timer system.
+    pub fn is_timer_error(&self) -> bool {
+        match self.0 {
+            Either::A(_) => false,
+            Either::B(_) => true,
+        }
+    }
+}
+
+impl<T: StdError> Display for ThrottleError<T> {
+    fn fmt(&self, f: &mut Formatter) -> FmtResult {
+        match self.0 {
+            Either::A(ref err) => write!(f, "stream error: {}", err),
+            Either::B(ref err) => write!(f, "timer error: {}", err),
+        }
+    }
+}
+
+impl<T: StdError + 'static> StdError for ThrottleError<T> {
+    fn description(&self) -> &str {
+        match self.0 {
+            Either::A(_) => "stream error",
+            Either::B(_) => "timer error",
+        }
+    }
+
+    fn cause(&self) -> Option<&StdError> {
+        match self.0 {
+            Either::A(ref err) => Some(err),
+            Either::B(ref err) => Some(err),
+        }
+    }
+}

tokio-timer/tests/throttle.rs

@@ -0,0 +1,57 @@
+extern crate futures;
+extern crate tokio;
+extern crate tokio_executor;
+extern crate tokio_timer;
+
+#[macro_use]
+mod support;
+use support::*;
+
+use futures::{
+    prelude::*,
+    sync::mpsc,
+};
+use tokio::util::StreamExt;
+
+#[test]
+fn throttle() {
+    mocked(|timer, _| {
+        let (tx, rx) = mpsc::unbounded();
+        let mut stream = rx.throttle(ms(1))
+            .map_err(|e| panic!("{:?}", e));
+
+        assert_not_ready!(stream);
+
+        for i in 0..3 {
+            tx.unbounded_send(i).unwrap();
+        }
+        for i in 0..3 {
+            assert_ready_eq!(stream, Some(i));
+            assert_not_ready!(stream);
+
+            advance(timer, ms(1));
+        }
+
+        assert_not_ready!(stream);
+    });
+}
+
+#[test]
+fn throttle_dur_0() {
+    mocked(|_, _| {
+        let (tx, rx) = mpsc::unbounded();
+        let mut stream = rx.throttle(ms(0))
+            .map_err(|e| panic!("{:?}", e));
+
+        assert_not_ready!(stream);
+
+        for i in 0..3 {
+            tx.unbounded_send(i).unwrap();
+        }
+        for i in 0..3 {
+            assert_ready_eq!(stream, Some(i));
+        }
+
+        assert_not_ready!(stream);
+    });
+}