2018-02-05 08:38:04 +00:00
|
|
|
use std::fmt::{Debug, Display, Formatter, Result as FmtResult};
|
|
|
|
use std::time::{Duration, Instant};
|
|
|
|
use std::ops::Deref;
|
|
|
|
use std::sync::Mutex;
|
|
|
|
|
2018-02-05 09:23:40 +00:00
|
|
|
use failure::{Error, Fail};
|
2018-02-05 08:38:04 +00:00
|
|
|
use futures::{Future, Poll};
|
|
|
|
use futures::future::{Shared, SharedError, SharedItem};
|
|
|
|
use tokio_service::Service;
|
|
|
|
|
2018-02-05 09:23:40 +00:00
|
|
|
pub struct Throttle<S>
|
|
|
|
where S: Service<Request=(), Error=Error>
|
|
|
|
{
|
2018-02-05 08:38:04 +00:00
|
|
|
inner: S,
|
|
|
|
duration: Duration,
|
|
|
|
current: Mutex<Option<(Instant, Shared<S::Future>)>>
|
|
|
|
}
|
|
|
|
|
|
|
|
impl<S> Debug for Throttle<S>
|
2018-02-05 09:23:40 +00:00
|
|
|
where S: Service<Request=(), Error=Error> + Debug
|
2018-02-05 08:38:04 +00:00
|
|
|
{
|
|
|
|
fn fmt(&self, fmt: &mut Formatter) -> FmtResult {
|
|
|
|
fmt.debug_struct("Throttle")
|
|
|
|
.field("inner", &self.inner)
|
|
|
|
.field("duration", &self.duration)
|
|
|
|
.finish()
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2018-02-05 09:23:40 +00:00
|
|
|
impl<S> Throttle<S>
|
|
|
|
where S: Service<Request=(), Error=Error>
|
|
|
|
{
|
2018-02-05 08:38:04 +00:00
|
|
|
pub fn new(service: S, duration: Duration) -> Throttle<S> {
|
|
|
|
Throttle {
|
|
|
|
inner: service,
|
|
|
|
duration,
|
|
|
|
current: Mutex::new(None)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2018-02-05 09:23:40 +00:00
|
|
|
impl<S> Service for Throttle<S>
|
|
|
|
where S: Service<Request=(), Error=Error>
|
|
|
|
{
|
2018-02-05 08:38:04 +00:00
|
|
|
type Request = ();
|
|
|
|
type Response = ThrottledItem<S::Response>;
|
2018-02-05 09:23:40 +00:00
|
|
|
type Error = ThrottledError;
|
2018-02-05 08:38:04 +00:00
|
|
|
type Future = Throttled<S::Future>;
|
|
|
|
|
|
|
|
fn call(&self, _: ()) -> Self::Future {
|
|
|
|
let now = Instant::now();
|
|
|
|
let mut current = self.current.lock().expect("lock poisoned");
|
|
|
|
if let Some((valid_until, ref shared_future)) = *current {
|
|
|
|
if valid_until > now {
|
|
|
|
if let Some(Ok(_)) = shared_future.peek() {
|
|
|
|
return Throttled(shared_future.clone());
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
let shared_future = self.inner.call(()).shared();
|
|
|
|
*current = Some((now + self.duration, shared_future.clone()));
|
|
|
|
Throttled(shared_future)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
pub struct Throttled<F: Future>(Shared<F>);
|
|
|
|
|
|
|
|
impl<F> Debug for Throttled<F>
|
|
|
|
where F: Future + Debug,
|
|
|
|
F::Item: Debug,
|
|
|
|
F::Error: Debug
|
|
|
|
{
|
|
|
|
fn fmt(&self, fmt: &mut Formatter) -> FmtResult {
|
|
|
|
self.0.fmt(fmt)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2018-02-05 09:23:40 +00:00
|
|
|
impl<F: Future<Error=Error>> Future for Throttled<F> {
|
2018-02-05 08:38:04 +00:00
|
|
|
type Item = ThrottledItem<F::Item>;
|
2018-02-05 09:23:40 +00:00
|
|
|
type Error = ThrottledError;
|
2018-02-05 08:38:04 +00:00
|
|
|
|
|
|
|
fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
|
|
|
|
self.0.poll()
|
|
|
|
.map_err(ThrottledError)
|
|
|
|
.map(|async| async.map(ThrottledItem))
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
#[derive(Debug)]
|
|
|
|
pub struct ThrottledItem<T>(SharedItem<T>);
|
|
|
|
|
|
|
|
impl<T> Deref for ThrottledItem<T> {
|
|
|
|
type Target = T;
|
|
|
|
|
|
|
|
fn deref(&self) -> &T {
|
|
|
|
&self.0.deref()
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
#[derive(Debug)]
|
2018-02-05 09:23:40 +00:00
|
|
|
pub struct ThrottledError(SharedError<Error>);
|
2018-02-05 08:38:04 +00:00
|
|
|
|
2018-02-05 09:23:40 +00:00
|
|
|
impl Fail for ThrottledError {
|
|
|
|
fn cause(&self) -> Option<&Fail> {
|
|
|
|
Some(self.0.cause())
|
2018-02-05 08:38:04 +00:00
|
|
|
}
|
|
|
|
|
2018-02-05 09:23:40 +00:00
|
|
|
fn backtrace(&self) -> Option<&::failure::Backtrace> {
|
|
|
|
Some(self.0.backtrace())
|
2018-02-05 08:38:04 +00:00
|
|
|
}
|
|
|
|
|
2018-02-05 09:23:40 +00:00
|
|
|
fn causes(&self) -> ::failure::Causes {
|
|
|
|
self.0.causes()
|
2018-02-05 08:38:04 +00:00
|
|
|
}
|
2018-02-05 09:23:40 +00:00
|
|
|
}
|
2018-02-05 08:38:04 +00:00
|
|
|
|
2018-02-05 09:23:40 +00:00
|
|
|
impl Display for ThrottledError {
|
|
|
|
fn fmt(&self, f: &mut Formatter) -> FmtResult {
|
|
|
|
Display::fmt(&self.0, f)
|
2018-02-05 08:38:04 +00:00
|
|
|
}
|
|
|
|
}
|