diff --git a/src/util/stream.rs b/src/util/stream.rs index ef268483c..8204ae0d5 100644 --- a/src/util/stream.rs +++ b/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 + where Self: Sized + { + Throttle::new(self, duration) + } /// Creates a new stream which allows `self` until `timeout`. /// diff --git a/tokio-timer/src/lib.rs b/tokio-timer/src/lib.rs index 6c1e164e8..c85f2db79 100644 --- a/tokio-timer/src/lib.rs +++ b/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; diff --git a/tokio-timer/src/throttle.rs b/tokio-timer/src/throttle.rs new file mode 100644 index 000000000..c1da76bbf --- /dev/null +++ b/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 { + delay: Option, + duration: Duration, + stream: T, +} + +/// Either the error of the underlying stream, or an error within +/// tokio's timing machinery. +#[derive(Debug)] +pub struct ThrottleError(Either); + +impl Throttle { + /// 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 Stream for Throttle { + type Item = T::Item; + type Error = ThrottleError; + + fn poll(&mut self) -> Poll, 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 ThrottleError { + /// 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 { + 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 { + 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 Display for ThrottleError { + 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 StdError for ThrottleError { + 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), + } + } +} diff --git a/tokio-timer/tests/throttle.rs b/tokio-timer/tests/throttle.rs new file mode 100644 index 000000000..ace947dd1 --- /dev/null +++ b/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); + }); +}