ratelimit: do not overflow leaky bucket

This way the time to wait until next message can
be sent is always limited.
This commit is contained in:
link2xt
2022-07-10 14:08:54 +00:00
parent 294d8862e4
commit 9926804f1b
2 changed files with 7 additions and 11 deletions

View File

@@ -4,6 +4,7 @@
### Changes ### Changes
- handle drafts from mailto links in scanned QR #3492 - handle drafts from mailto links in scanned QR #3492
- do not overflow ratelimiter leaky bucket #3496
### Fixes ### Fixes
- don't squash text parts of NDN into attachments #3497 - don't squash text parts of NDN into attachments #3497

View File

@@ -51,7 +51,7 @@ impl Ratelimit {
/// Returns true if it is allowed to send a message. /// Returns true if it is allowed to send a message.
fn can_send_at(&self, now: SystemTime) -> bool { fn can_send_at(&self, now: SystemTime) -> bool {
self.current_value_at(now) <= self.quota self.current_value_at(now) + 1.0 <= self.quota
} }
/// Returns true if can send another message now. /// Returns true if can send another message now.
@@ -62,7 +62,7 @@ impl Ratelimit {
} }
fn send_at(&mut self, now: SystemTime) { fn send_at(&mut self, now: SystemTime) {
self.current_value = self.current_value_at(now) + 1.0; self.current_value = f64::min(self.quota, self.current_value_at(now) + 1.0);
self.last_update = now; self.last_update = now;
} }
@@ -77,10 +77,10 @@ impl Ratelimit {
fn until_can_send_at(&self, now: SystemTime) -> Duration { fn until_can_send_at(&self, now: SystemTime) -> Duration {
let current_value = self.current_value_at(now); let current_value = self.current_value_at(now);
if current_value <= self.quota { if current_value + 1.0 <= self.quota {
Duration::ZERO Duration::ZERO
} else { } else {
let requirement = current_value - self.quota; let requirement = current_value + 1.0 - self.quota;
let rate = self.quota / self.window.as_secs_f64(); let rate = self.quota / self.window.as_secs_f64();
Duration::from_secs_f64(requirement / rate) Duration::from_secs_f64(requirement / rate)
} }
@@ -109,8 +109,6 @@ mod tests {
ratelimit.send_at(now); ratelimit.send_at(now);
assert!(ratelimit.can_send_at(now)); assert!(ratelimit.can_send_at(now));
ratelimit.send_at(now); ratelimit.send_at(now);
assert!(ratelimit.can_send_at(now));
ratelimit.send_at(now);
// Can't send more messages now. // Can't send more messages now.
assert!(!ratelimit.can_send_at(now)); assert!(!ratelimit.can_send_at(now));
@@ -125,11 +123,8 @@ mod tests {
// Send one more message anyway, over quota. // Send one more message anyway, over quota.
ratelimit.send_at(now); ratelimit.send_at(now);
// Waiting 20 seconds is not enough. // Always can send another message after 20 seconds,
let now = now + Duration::from_secs(20); // leaky bucket never overflows.
assert!(!ratelimit.can_send_at(now));
// Can send another message after 40 seconds.
let now = now + Duration::from_secs(20); let now = now + Duration::from_secs(20);
assert!(ratelimit.can_send_at(now)); assert!(ratelimit.can_send_at(now));