Skip to main content

max / makenotwork

11.0 KB · 336 lines History Blame Raw
1 //! Per-(user, room) send budget.
2 //!
3 //! Chat needs its own limiter rather than borrowing either host's. Multithreaded
4 //! caps posting at 15 per 60 seconds, which is a sensible forum-post rate and an
5 //! unusable chat rate. The per-IP `tower_governor` layers do not fit either: they
6 //! key on the connection, and chat's connection is long-lived and shared by every
7 //! message on it.
8 //!
9 //! A token bucket rather than a fixed window, because the shape of chat traffic is
10 //! bursty by nature. Someone answering a question types four short lines in five
11 //! seconds and then goes quiet for a minute; a fixed window either refuses the
12 //! fourth line or has to be set so loose it stops being a limit.
13 //!
14 //! # Eviction is not optional
15 //!
16 //! One bucket per (user, room) pair, held in a map, is unbounded state keyed by
17 //! something an attacker influences. Under a hard 512M cgroup cap that is a memory
18 //! exhaustion vector, not untidiness, so [`RateLimiter::evict_idle`] exists and
19 //! the host is expected to call it on a timer. Multithreaded's per-IP governor
20 //! already sweeps its own buckets every 300 seconds; this wants the same
21 //! treatment.
22
23 use std::time::{Duration, Instant};
24
25 use dashmap::DashMap;
26
27 use crate::ids::{RoomId, UserId};
28
29 /// Send budget.
30 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
31 pub struct RateLimits {
32 /// Messages allowed back to back from a standing start.
33 pub burst: u32,
34 /// Long-run rate, in messages per minute.
35 pub sustain_per_min: u32,
36 }
37
38 impl Default for RateLimits {
39 fn default() -> Self {
40 // Ten in a row, then one every two seconds. Fast enough that a person
41 // talking normally never meets it, slow enough that a script is throttled
42 // to something a room can absorb.
43 Self {
44 burst: 10,
45 sustain_per_min: 30,
46 }
47 }
48 }
49
50 impl RateLimits {
51 fn refill_per_sec(self) -> f64 {
52 f64::from(self.sustain_per_min) / 60.0
53 }
54 }
55
56 #[derive(Debug, Clone, Copy)]
57 struct Bucket {
58 tokens: f64,
59 last: Instant,
60 }
61
62 /// Outcome of a send check.
63 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
64 pub enum RateDecision {
65 Allow,
66 /// Refused. `retry_after` is when a token next becomes available, so the
67 /// caller can tell the client rather than leaving it to guess.
68 Deny {
69 retry_after: Duration,
70 },
71 }
72
73 impl RateDecision {
74 pub fn is_allowed(&self) -> bool {
75 matches!(self, Self::Allow)
76 }
77 }
78
79 /// Token buckets keyed by (user, room).
80 ///
81 /// Keyed by the pair rather than by user alone so that being talkative in one
82 /// room does not silence someone in another, and so a room's budget cannot be
83 /// consumed from outside it.
84 pub struct RateLimiter {
85 buckets: DashMap<(UserId, RoomId), Bucket>,
86 limits: RateLimits,
87 }
88
89 impl RateLimiter {
90 pub fn new(limits: RateLimits) -> Self {
91 Self {
92 buckets: DashMap::new(),
93 limits,
94 }
95 }
96
97 /// Take a token for one message.
98 ///
99 /// `now` is passed in rather than read here so the behavior is testable
100 /// without sleeping, and so a caller batching work can use one consistent
101 /// instant.
102 pub fn check(&self, user: UserId, room: RoomId, now: Instant) -> RateDecision {
103 let capacity = f64::from(self.limits.burst);
104 let refill = self.limits.refill_per_sec();
105
106 let mut bucket = self.buckets.entry((user, room)).or_insert(Bucket {
107 tokens: capacity,
108 last: now,
109 });
110
111 // `saturating_duration_since` rather than subtraction: a caller passing a
112 // `now` older than the last one (batching, or a clock read from a
113 // different task) would otherwise panic.
114 let elapsed = now.saturating_duration_since(bucket.last).as_secs_f64();
115 bucket.tokens = (bucket.tokens + elapsed * refill).min(capacity);
116 bucket.last = now;
117
118 if bucket.tokens >= 1.0 {
119 bucket.tokens -= 1.0;
120 RateDecision::Allow
121 } else {
122 let deficit = 1.0 - bucket.tokens;
123 RateDecision::Deny {
124 retry_after: Duration::from_secs_f64(deficit / refill),
125 }
126 }
127 }
128
129 /// Drop buckets untouched for longer than `idle_for`.
130 ///
131 /// A bucket at full tokens carries no information: recreating it produces the
132 /// same answer as keeping it. So eviction cannot make the limiter more
133 /// permissive as long as `idle_for` is at least the time it takes a bucket to
134 /// refill, which is what [`RateLimiter::recommended_idle_timeout`] returns.
135 ///
136 /// Returns how many were dropped.
137 pub fn evict_idle(&self, now: Instant, idle_for: Duration) -> usize {
138 let before = self.buckets.len();
139 self.buckets
140 .retain(|_, bucket| now.saturating_duration_since(bucket.last) < idle_for);
141 before - self.buckets.len()
142 }
143
144 /// The shortest idle timeout that cannot make eviction a bypass: long enough
145 /// for an empty bucket to have refilled completely.
146 pub fn recommended_idle_timeout(&self) -> Duration {
147 Duration::from_secs_f64(f64::from(self.limits.burst) / self.limits.refill_per_sec())
148 }
149
150 /// Live bucket count. Diagnostics, and the number to watch if eviction is
151 /// ever suspected of not running.
152 pub fn bucket_count(&self) -> usize {
153 self.buckets.len()
154 }
155 }
156
157 #[cfg(test)]
158 mod tests {
159 use super::*;
160 use uuid::Uuid;
161
162 fn user() -> UserId {
163 UserId(Uuid::new_v4())
164 }
165
166 fn room() -> RoomId {
167 RoomId(Uuid::new_v4())
168 }
169
170 fn limiter(burst: u32, sustain_per_min: u32) -> RateLimiter {
171 RateLimiter::new(RateLimits {
172 burst,
173 sustain_per_min,
174 })
175 }
176
177 #[test]
178 fn burst_is_spendable_immediately_then_refused() {
179 let rl = limiter(3, 60);
180 let (u, r, t0) = (user(), room(), Instant::now());
181
182 for i in 0..3 {
183 assert!(rl.check(u, r, t0).is_allowed(), "burst message {i}");
184 }
185 assert!(!rl.check(u, r, t0).is_allowed(), "burst exhausted");
186 }
187
188 #[test]
189 fn refuses_with_a_usable_retry_after() {
190 // 60/min is one per second, so an empty bucket is one second from a token.
191 let rl = limiter(1, 60);
192 let (u, r, t0) = (user(), room(), Instant::now());
193
194 assert!(rl.check(u, r, t0).is_allowed());
195 match rl.check(u, r, t0) {
196 RateDecision::Deny { retry_after } => {
197 let ms = retry_after.as_millis();
198 assert!((900..=1100).contains(&ms), "retry_after was {ms}ms");
199 }
200 RateDecision::Allow => panic!("should have been refused"),
201 }
202 }
203
204 #[test]
205 fn tokens_come_back_over_time() {
206 let rl = limiter(2, 60);
207 let (u, r, t0) = (user(), room(), Instant::now());
208
209 assert!(rl.check(u, r, t0).is_allowed());
210 assert!(rl.check(u, r, t0).is_allowed());
211 assert!(!rl.check(u, r, t0).is_allowed());
212
213 assert!(rl.check(u, r, t0 + Duration::from_secs(1)).is_allowed());
214 }
215
216 #[test]
217 fn refill_does_not_exceed_the_burst_ceiling() {
218 // Idling for an hour must not bank an hour of messages.
219 let rl = limiter(3, 60);
220 let (u, r, t0) = (user(), room(), Instant::now());
221
222 assert!(rl.check(u, r, t0).is_allowed());
223 let long_later = t0 + Duration::from_hours(1);
224
225 for i in 0..3 {
226 assert!(rl.check(u, r, long_later).is_allowed(), "message {i}");
227 }
228 assert!(
229 !rl.check(u, r, long_later).is_allowed(),
230 "banked more than the burst"
231 );
232 }
233
234 #[test]
235 fn the_sustained_rate_is_actually_the_sustained_rate() {
236 // Spend the burst, then send at exactly the refill interval forever.
237 let rl = limiter(5, 60);
238 let (u, r) = (user(), room());
239 let mut t = Instant::now();
240
241 for _ in 0..5 {
242 assert!(rl.check(u, r, t).is_allowed());
243 }
244 for i in 0..100 {
245 t += Duration::from_secs(1);
246 assert!(rl.check(u, r, t).is_allowed(), "sustained message {i}");
247 }
248 // And faster than that is still refused.
249 assert!(!rl.check(u, r, t).is_allowed());
250 }
251
252 #[test]
253 fn budgets_are_per_user_and_per_room() {
254 let rl = limiter(1, 60);
255 let (u1, u2) = (user(), user());
256 let (r1, r2) = (room(), room());
257 let t0 = Instant::now();
258
259 assert!(rl.check(u1, r1, t0).is_allowed());
260 assert!(!rl.check(u1, r1, t0).is_allowed(), "same pair is spent");
261 assert!(rl.check(u2, r1, t0).is_allowed(), "other user unaffected");
262 assert!(rl.check(u1, r2, t0).is_allowed(), "other room unaffected");
263 }
264
265 #[test]
266 fn a_backwards_now_does_not_panic() {
267 let rl = limiter(2, 60);
268 let (u, r) = (user(), room());
269 let t = Instant::now() + Duration::from_secs(10);
270
271 assert!(rl.check(u, r, t).is_allowed());
272 // Earlier than the bucket's last touch: elapsed saturates to zero rather
273 // than underflowing.
274 assert!(
275 rl.check(u, r, t.checked_sub(Duration::from_secs(5)).unwrap())
276 .is_allowed()
277 );
278 }
279
280 #[test]
281 fn eviction_drops_only_idle_buckets() {
282 let rl = limiter(5, 60);
283 let (busy, idle) = (user(), user());
284 let r = room();
285 let t0 = Instant::now();
286
287 rl.check(busy, r, t0);
288 rl.check(idle, r, t0);
289 assert_eq!(rl.bucket_count(), 2);
290
291 let later = t0 + Duration::from_mins(1);
292 rl.check(busy, r, later);
293
294 assert_eq!(rl.evict_idle(later, Duration::from_secs(30)), 1);
295 assert_eq!(rl.bucket_count(), 1);
296 }
297
298 #[test]
299 fn eviction_at_the_recommended_timeout_is_not_a_bypass() {
300 // The concern: spend the burst, wait out eviction, and get a fresh bucket
301 // sooner than refilling would have allowed. At the recommended timeout the
302 // bucket has fully refilled anyway, so the two paths agree.
303 let rl = limiter(5, 60);
304 let (u, r, t0) = (user(), room(), Instant::now());
305
306 for _ in 0..5 {
307 assert!(rl.check(u, r, t0).is_allowed());
308 }
309 assert!(!rl.check(u, r, t0).is_allowed());
310
311 let timeout = rl.recommended_idle_timeout();
312 assert_eq!(timeout, Duration::from_secs(5), "5 tokens at 1/sec");
313
314 let after = t0 + timeout;
315 assert_eq!(rl.evict_idle(after, timeout), 1);
316
317 // A fresh bucket gives exactly the burst, which is what waiting the same
318 // period without eviction would also have given.
319 for i in 0..5 {
320 assert!(rl.check(u, r, after).is_allowed(), "post-eviction {i}");
321 }
322 assert!(!rl.check(u, r, after).is_allowed(), "no extra budget");
323 }
324
325 #[test]
326 fn default_limits_are_not_the_forum_post_limit() {
327 // The whole reason this module exists: 15/min would make chat unusable.
328 let d = RateLimits::default();
329 assert!(
330 d.sustain_per_min > 15,
331 "chat sustain must exceed the forum post rate"
332 );
333 assert!(d.burst >= 5, "chat must absorb a short burst of typing");
334 }
335 }
336