Fix bug in rate limiter which caused tokens to be underprovisioned

I think the primary cause of the rate-limiting was simply that we didn't
have enough capacity. However, this could underprovision new tokens when
e.g. the `refill_time_per_token` was 1s and a request occurred after 1.5s.
Previously, we would throw away those 0.5s, which would cause the user to
have to wait for another full second before getting a new token, as opposed
to the correct 500ms.
This commit is contained in:
Antonio Scandurra 2024-05-29 11:36:30 +02:00
parent cad3d2d355
commit 34b88d14f6

View File

@ -148,9 +148,12 @@ impl RateBucket {
if elapsed >= self.refill_time_per_token {
let new_tokens =
elapsed.num_milliseconds() / self.refill_time_per_token.num_milliseconds();
self.token_count = (self.token_count + new_tokens as usize).min(self.capacity);
self.last_refill = now;
let unused_refill_time = Duration::milliseconds(
elapsed.num_milliseconds() % self.refill_time_per_token.num_milliseconds(),
);
self.last_refill = now - unused_refill_time;
}
}
}
@ -218,8 +221,19 @@ mod tests {
.await
.unwrap();
// After one second, user 1 can make another request before being rate-limited again.
now += Duration::seconds(1);
// After 1.5s, user 1 can make another request before being rate-limited again.
now += Duration::milliseconds(1500);
rate_limiter
.check_internal::<RateLimitA>(user_1, now)
.await
.unwrap();
rate_limiter
.check_internal::<RateLimitA>(user_1, now)
.await
.unwrap_err();
// After 500ms, user 1 can make another request before being rate-limited again.
now += Duration::milliseconds(500);
rate_limiter
.check_internal::<RateLimitA>(user_1, now)
.await