rate-limit: Refresh layout (#189)
- Switch to `Box<Error>` - Break up lib.rs into multiple files. - Use `tokio::clock::now` instead of `Instant::now`.
This commit is contained in:
parent
20102a647b
commit
720d31c65f
|
@ -0,0 +1,18 @@
|
|||
use std::error;
|
||||
|
||||
pub(crate) type Error = Box<error::Error + Send + Sync>;
|
||||
|
||||
pub(crate) mod never {
|
||||
use std::{error, fmt};
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum Never {}
|
||||
|
||||
impl fmt::Display for Never {
|
||||
fn fmt(&self, _: &mut fmt::Formatter) -> fmt::Result {
|
||||
unreachable!();
|
||||
}
|
||||
}
|
||||
|
||||
impl error::Error for Never {}
|
||||
}
|
|
@ -0,0 +1,26 @@
|
|||
use crate::Error;
|
||||
use futures::{Future, Poll};
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct ResponseFuture<T> {
|
||||
inner: T,
|
||||
}
|
||||
|
||||
impl<T> ResponseFuture<T> {
|
||||
pub(crate) fn new(inner: T) -> ResponseFuture<T> {
|
||||
ResponseFuture { inner }
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> Future for ResponseFuture<T>
|
||||
where
|
||||
T: Future,
|
||||
Error: From<T::Error>,
|
||||
{
|
||||
type Item = T::Item;
|
||||
type Error = Error;
|
||||
|
||||
fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
|
||||
self.inner.poll().map_err(Into::into)
|
||||
}
|
||||
}
|
|
@ -0,0 +1,32 @@
|
|||
use crate::error::{never::Never, Error};
|
||||
use crate::{Rate, RateLimit};
|
||||
use std::time::Duration;
|
||||
use tower_layer::Layer;
|
||||
use tower_service::Service;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct RateLimitLayer {
|
||||
rate: Rate,
|
||||
}
|
||||
|
||||
impl RateLimitLayer {
|
||||
pub fn new(num: u64, per: Duration) -> Self {
|
||||
let rate = Rate::new(num, per);
|
||||
RateLimitLayer { rate }
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, Request> Layer<S, Request> for RateLimitLayer
|
||||
where
|
||||
S: Service<Request>,
|
||||
Error: From<S::Error>,
|
||||
{
|
||||
type Response = S::Response;
|
||||
type Error = Error;
|
||||
type LayerError = Never;
|
||||
type Service = RateLimit<S>;
|
||||
|
||||
fn layer(&self, service: S) -> Result<Self::Service, Self::LayerError> {
|
||||
Ok(RateLimit::new(service, self.rate))
|
||||
}
|
||||
}
|
|
@ -7,13 +7,21 @@ extern crate tokio_timer;
|
|||
extern crate tower_layer;
|
||||
extern crate tower_service;
|
||||
|
||||
pub mod error;
|
||||
pub mod future;
|
||||
mod layer;
|
||||
mod rate;
|
||||
|
||||
pub use crate::layer::RateLimitLayer;
|
||||
pub use crate::rate::Rate;
|
||||
|
||||
use crate::error::Error;
|
||||
use crate::future::ResponseFuture;
|
||||
use futures::{Future, Poll};
|
||||
use tokio_timer::Delay;
|
||||
use tower_layer::Layer;
|
||||
use tokio_timer::{clock, Delay};
|
||||
use tower_service::Service;
|
||||
|
||||
use std::time::{Duration, Instant};
|
||||
use std::{error, fmt};
|
||||
use std::time::Instant;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct RateLimit<T> {
|
||||
|
@ -22,30 +30,6 @@ pub struct RateLimit<T> {
|
|||
state: State,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct RateLimitLayer {
|
||||
rate: Rate,
|
||||
}
|
||||
|
||||
#[derive(Debug, Copy, Clone)]
|
||||
pub struct Rate {
|
||||
num: u64,
|
||||
per: Duration,
|
||||
}
|
||||
|
||||
/// The request has been rate limited
|
||||
///
|
||||
/// TODO: Consider returning the original request
|
||||
#[derive(Debug)]
|
||||
pub enum Error<T> {
|
||||
RateLimit,
|
||||
Upstream(T),
|
||||
}
|
||||
|
||||
pub struct ResponseFuture<T> {
|
||||
inner: T,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum State {
|
||||
// The service has hit its limit
|
||||
|
@ -53,27 +37,6 @@ enum State {
|
|||
Ready { until: Instant, rem: u64 },
|
||||
}
|
||||
|
||||
impl RateLimitLayer {
|
||||
pub fn new(num: u64, per: Duration) -> Self {
|
||||
let rate = Rate { num, per };
|
||||
RateLimitLayer { rate }
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, Request> Layer<S, Request> for RateLimitLayer
|
||||
where
|
||||
S: Service<Request>,
|
||||
{
|
||||
type Response = S::Response;
|
||||
type Error = Error<S::Error>;
|
||||
type LayerError = ();
|
||||
type Service = RateLimit<S>;
|
||||
|
||||
fn layer(&self, service: S) -> Result<Self::Service, Self::LayerError> {
|
||||
Ok(RateLimit::new(service, self.rate))
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> RateLimit<T> {
|
||||
/// Create a new rate limiter
|
||||
pub fn new<Request>(inner: T, rate: Rate) -> Self
|
||||
|
@ -81,8 +44,8 @@ impl<T> RateLimit<T> {
|
|||
T: Service<Request>,
|
||||
{
|
||||
let state = State::Ready {
|
||||
until: Instant::now(),
|
||||
rem: rate.num,
|
||||
until: clock::now(),
|
||||
rem: rate.num(),
|
||||
};
|
||||
|
||||
RateLimit {
|
||||
|
@ -108,41 +71,26 @@ impl<T> RateLimit<T> {
|
|||
}
|
||||
}
|
||||
|
||||
impl Rate {
|
||||
/// Create a new rate
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// This function panics if `num` or `per` is 0.
|
||||
pub fn new(num: u64, per: Duration) -> Self {
|
||||
assert!(num > 0);
|
||||
assert!(per > Duration::from_millis(0));
|
||||
|
||||
Rate { num, per }
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, Request> Service<Request> for RateLimit<S>
|
||||
where
|
||||
S: Service<Request>,
|
||||
Error: From<S::Error>,
|
||||
{
|
||||
type Response = S::Response;
|
||||
type Error = Error<S::Error>;
|
||||
type Error = Error;
|
||||
type Future = ResponseFuture<S::Future>;
|
||||
|
||||
fn poll_ready(&mut self) -> Poll<(), Self::Error> {
|
||||
match self.state {
|
||||
State::Ready { .. } => return Ok(().into()),
|
||||
State::Limited(ref mut sleep) => {
|
||||
let res = sleep.poll().map_err(|_| Error::RateLimit);
|
||||
|
||||
try_ready!(res);
|
||||
try_ready!(sleep.poll());
|
||||
}
|
||||
}
|
||||
|
||||
self.state = State::Ready {
|
||||
until: Instant::now() + self.rate.per,
|
||||
rem: self.rate.num,
|
||||
until: clock::now() + self.rate.per(),
|
||||
rem: self.rate.num(),
|
||||
};
|
||||
|
||||
Ok(().into())
|
||||
|
@ -151,12 +99,12 @@ where
|
|||
fn call(&mut self, request: Request) -> Self::Future {
|
||||
match self.state {
|
||||
State::Ready { mut until, mut rem } => {
|
||||
let now = Instant::now();
|
||||
let now = clock::now();
|
||||
|
||||
// If the period has elapsed, reset it.
|
||||
if now >= until {
|
||||
until = now + self.rate.per;
|
||||
let rem = self.rate.num;
|
||||
until = now + self.rate.per();
|
||||
let rem = self.rate.num();
|
||||
|
||||
self.state = State::Ready { until, rem }
|
||||
}
|
||||
|
@ -172,55 +120,9 @@ where
|
|||
|
||||
// Call the inner future
|
||||
let inner = self.inner.call(request);
|
||||
ResponseFuture { inner }
|
||||
ResponseFuture::new(inner)
|
||||
}
|
||||
State::Limited(..) => panic!("service not ready; poll_ready must be called first"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> Future for ResponseFuture<T>
|
||||
where
|
||||
T: Future,
|
||||
{
|
||||
type Item = T::Item;
|
||||
type Error = Error<T::Error>;
|
||||
|
||||
fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
|
||||
self.inner.poll().map_err(Error::Upstream)
|
||||
}
|
||||
}
|
||||
|
||||
// ===== impl Error =====
|
||||
|
||||
impl<T> fmt::Display for Error<T>
|
||||
where
|
||||
T: fmt::Display,
|
||||
{
|
||||
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
|
||||
match *self {
|
||||
Error::Upstream(ref why) => fmt::Display::fmt(why, f),
|
||||
Error::RateLimit => f.pad("rate limit exceeded"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> error::Error for Error<T>
|
||||
where
|
||||
T: error::Error,
|
||||
{
|
||||
fn cause(&self) -> Option<&error::Error> {
|
||||
if let Error::Upstream(ref why) = *self {
|
||||
Some(why)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
match *self {
|
||||
Error::Upstream(_) => "upstream service error",
|
||||
Error::RateLimit => "rate limit exceeded",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
@ -0,0 +1,29 @@
|
|||
use std::time::Duration;
|
||||
|
||||
#[derive(Debug, Copy, Clone)]
|
||||
pub struct Rate {
|
||||
num: u64,
|
||||
per: Duration,
|
||||
}
|
||||
|
||||
impl Rate {
|
||||
/// Create a new rate
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// This function panics if `num` or `per` is 0.
|
||||
pub fn new(num: u64, per: Duration) -> Self {
|
||||
assert!(num > 0);
|
||||
assert!(per > Duration::from_millis(0));
|
||||
|
||||
Rate { num, per }
|
||||
}
|
||||
|
||||
pub(crate) fn num(&self) -> u64 {
|
||||
self.num
|
||||
}
|
||||
|
||||
pub(crate) fn per(&self) -> Duration {
|
||||
self.per
|
||||
}
|
||||
}
|
Loading…
Reference in New Issue