rate: Fix rate limit not resetting (#439)
This commit is contained in:
parent
cd7dd12315
commit
5e1788f494
|
@ -67,7 +67,10 @@ where
|
||||||
match self.state {
|
match self.state {
|
||||||
State::Ready { .. } => return Poll::Ready(ready!(self.inner.poll_ready(cx))),
|
State::Ready { .. } => return Poll::Ready(ready!(self.inner.poll_ready(cx))),
|
||||||
State::Limited(ref mut sleep) => {
|
State::Limited(ref mut sleep) => {
|
||||||
ready!(Pin::new(sleep).poll(cx));
|
if let Poll::Pending = Pin::new(sleep).poll(cx) {
|
||||||
|
tracing::trace!("rate limit exceeded; sleeping.");
|
||||||
|
return Poll::Pending;
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -87,9 +90,7 @@ where
|
||||||
// If the period has elapsed, reset it.
|
// If the period has elapsed, reset it.
|
||||||
if now >= until {
|
if now >= until {
|
||||||
until = now + self.rate.per();
|
until = now + self.rate.per();
|
||||||
let rem = self.rate.num();
|
rem = self.rate.num();
|
||||||
|
|
||||||
self.state = State::Ready { until, rem }
|
|
||||||
}
|
}
|
||||||
|
|
||||||
if rem > 1 {
|
if rem > 1 {
|
||||||
|
|
|
@ -9,7 +9,6 @@ async fn reaching_capacity() {
|
||||||
time::pause();
|
time::pause();
|
||||||
|
|
||||||
let rate_limit = RateLimitLayer::new(1, Duration::from_millis(100));
|
let rate_limit = RateLimitLayer::new(1, Duration::from_millis(100));
|
||||||
|
|
||||||
let (mut service, mut handle) = mock::spawn_layer(rate_limit);
|
let (mut service, mut handle) = mock::spawn_layer(rate_limit);
|
||||||
|
|
||||||
assert_ready_ok!(service.poll_ready());
|
assert_ready_ok!(service.poll_ready());
|
||||||
|
@ -33,3 +32,38 @@ async fn reaching_capacity() {
|
||||||
|
|
||||||
assert_eq!(response.await.unwrap(), "done");
|
assert_eq!(response.await.unwrap(), "done");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn remaining_gets_reset() {
|
||||||
|
// This test checks for the case where the `until` state gets reset
|
||||||
|
// but the `rem` does not. This was a bug found `cd7dd12315706fc0860a35646b1eb7b60c50a5c1`.
|
||||||
|
//
|
||||||
|
// The main premise here is that we can make one request which should initialize the state
|
||||||
|
// as ready. Then we can advance the clock to put us beyond the current period. When we make
|
||||||
|
// subsequent requests the `rem` for the next window is continued from the previous when
|
||||||
|
// it should be totally reset.
|
||||||
|
|
||||||
|
time::pause();
|
||||||
|
|
||||||
|
let rate_limit = RateLimitLayer::new(3, Duration::from_millis(100));
|
||||||
|
let (mut service, mut handle) = mock::spawn_layer(rate_limit);
|
||||||
|
|
||||||
|
assert_ready_ok!(service.poll_ready());
|
||||||
|
let response = service.call("hello");
|
||||||
|
assert_request_eq!(handle, "hello").send_response("world");
|
||||||
|
assert_eq!(response.await.unwrap(), "world");
|
||||||
|
|
||||||
|
time::advance(Duration::from_millis(100)).await;
|
||||||
|
|
||||||
|
assert_ready_ok!(service.poll_ready());
|
||||||
|
let response = service.call("hello");
|
||||||
|
assert_request_eq!(handle, "hello").send_response("world");
|
||||||
|
assert_eq!(response.await.unwrap(), "world");
|
||||||
|
|
||||||
|
assert_ready_ok!(service.poll_ready());
|
||||||
|
let response = service.call("hello");
|
||||||
|
assert_request_eq!(handle, "hello").send_response("world");
|
||||||
|
assert_eq!(response.await.unwrap(), "world");
|
||||||
|
|
||||||
|
assert_ready_ok!(service.poll_ready());
|
||||||
|
}
|
||||||
|
|
Loading…
Reference in New Issue