rate: Fix rate limit not resetting (#439)
This commit is contained in:
parent
cd7dd12315
commit
5e1788f494
|
@ -67,7 +67,10 @@ where
|
|||
match self.state {
|
||||
State::Ready { .. } => return Poll::Ready(ready!(self.inner.poll_ready(cx))),
|
||||
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 now >= until {
|
||||
until = now + self.rate.per();
|
||||
let rem = self.rate.num();
|
||||
|
||||
self.state = State::Ready { until, rem }
|
||||
rem = self.rate.num();
|
||||
}
|
||||
|
||||
if rem > 1 {
|
||||
|
|
|
@ -9,7 +9,6 @@ async fn reaching_capacity() {
|
|||
time::pause();
|
||||
|
||||
let rate_limit = RateLimitLayer::new(1, Duration::from_millis(100));
|
||||
|
||||
let (mut service, mut handle) = mock::spawn_layer(rate_limit);
|
||||
|
||||
assert_ready_ok!(service.poll_ready());
|
||||
|
@ -33,3 +32,38 @@ async fn reaching_capacity() {
|
|||
|
||||
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