rate_limit.nx source
↩ module page · 136 lines · 5305 B
1// rate_limit.nx -- token-bucket rate limiter.
2//
3// Standard algorithm (RFC 2698 / IEEE; also what HAProxy, nginx,
4// AWS API Gateway, and Stripe use):
5//
6// - Bucket has a max capacity (burst size) of B tokens
7// - Tokens refill at rate R per second up to B
8// - Each request consumes 1 token (or more for weighted ops)
9// - Request allowed iff current tokens >= needed; otherwise
10// reject (429 Too Many Requests) or wait
11//
12// Time is passed in as unix_microseconds -- caller reads a
13// monotonic clock and hands it to us. This keeps the module
14// pure + testable and decoupled from the OS time source.
15// For real use pair with ntp.nx + sys_clock_monotonic.
16//
17// Use cases: HTTP API throttling (per IP / per token), login
18// attempt rate limiting (anti-brute-force), outbound API call
19// pacing, mail-send pacing, git clone pacing.
20//
21// Invariants:
22// RL1 tokens clamped to capacity on every refill (bucket
23// never overflows).
24// RL2 Time going backwards does nothing (malicious / NTP
25// adjustment) -- we record max(now, last_refill).
26// RL3 Rate is expressed in tokens per second; sub-second
27// granularity handled via microsecond timestamps.
28
29import "syscalls.nx"
30
31struct RateLimiter {
32 capacity: i64, // bucket max (burst size)
33 refill_per_sec: i64, // tokens added per second
34 tokens_scaled: i64, // current tokens * 1_000_000 (so we can
35 // add fractional tokens without FP)
36 last_refill_us: i64, // last top-up timestamp in microseconds
37}
38
39const RL_SCALE: i64 = 1000000
40
41// Initialise a limiter. Starts full (so new clients get a burst
42// before they have to wait for refill).
43func rate_limit_init(rl: *RateLimiter, capacity: i64,
44 refill_per_sec: i64, now_us: i64) -> i64 {
45 rl.capacity = capacity
46 rl.refill_per_sec = refill_per_sec
47 rl.tokens_scaled = capacity * RL_SCALE
48 rl.last_refill_us = now_us
49 return 0
50}
51
52// Recompute tokens given current time. Safe to call
53// idempotently; extracted for use by try_take + peek.
54func rate_limit_refill(rl: *RateLimiter, now_us: i64) -> i64 {
55 if now_us <= rl.last_refill_us { return 0 }
56 let elapsed_us: i64 = now_us - rl.last_refill_us
57 // Tokens to add scaled: (elapsed_us * refill_per_sec) so the
58 // scale stays tokens*1_000_000.
59 let add_scaled: i64 = elapsed_us * rl.refill_per_sec
60 var new_scaled: i64 = rl.tokens_scaled + add_scaled
61 let max_scaled: i64 = rl.capacity * RL_SCALE
62 if new_scaled > max_scaled { new_scaled = max_scaled }
63 rl.tokens_scaled = new_scaled
64 rl.last_refill_us = now_us
65 return 0
66}
67
68// Try to take n tokens. Returns 1 if granted, 0 if denied.
69func rate_limit_try_take(rl: *RateLimiter, n: i64, now_us: i64) -> i64 {
70 rate_limit_refill(rl, now_us)
71 let need_scaled: i64 = n * RL_SCALE
72 if rl.tokens_scaled < need_scaled { return 0 }
73 rl.tokens_scaled = rl.tokens_scaled - need_scaled
74 return 1
75}
76
77// Peek: how many integer tokens available right now.
78func rate_limit_available(rl: *RateLimiter, now_us: i64) -> i64 {
79 rate_limit_refill(rl, now_us)
80 return rl.tokens_scaled / RL_SCALE
81}
82
83// Time until next single token (or 0 if available now), in us.
84// Useful for 429-Retry-After header values.
85func rate_limit_next_token_us(rl: *RateLimiter, now_us: i64) -> i64 {
86 rate_limit_refill(rl, now_us)
87 if rl.tokens_scaled >= RL_SCALE { return 0 }
88 let deficit_scaled: i64 = RL_SCALE - rl.tokens_scaled
89 if rl.refill_per_sec == 0 { return -1 } // never
90 // deficit_scaled tokens to generate at refill_per_sec per
91 // second. Each us of elapsed_us contributes refill_per_sec
92 // to new_scaled. us_needed = ceil(deficit_scaled / refill_per_sec).
93 let us_needed: i64 = (deficit_scaled + rl.refill_per_sec - 1) / rl.refill_per_sec
94 return us_needed
95}
96
97// Compile-only smoke.
98func main() -> i64 {
99 let rl_raw: *u8 = sys_mmap(64)
100 let rl: *RateLimiter = rl_raw as *RateLimiter
101
102 // 10-token burst, 2 tokens per second, start at t=0us.
103 rate_limit_init(rl, 10, 2, 0)
104 if rate_limit_available(rl, 0) != 10 { return 1 }
105
106 // Consume 10 tokens rapidly -- all should succeed.
107 var i: i64 = 0
108 while i < 10 {
109 if rate_limit_try_take(rl, 1, 0) != 1 { return 2 }
110 i = i + 1
111 }
112 if rate_limit_available(rl, 0) != 0 { return 3 }
113
114 // 11th -- no tokens.
115 if rate_limit_try_take(rl, 1, 0) != 0 { return 4 }
116
117 // At t=500ms (500_000 us), we've refilled 1 token.
118 if rate_limit_available(rl, 500000) != 1 { return 5 }
119 if rate_limit_try_take(rl, 1, 500000) != 1 { return 6 }
120
121 // Can't take another immediately.
122 if rate_limit_try_take(rl, 1, 500000) != 0 { return 7 }
123
124 // Time needed for next token: 500_000us.
125 let wait_us: i64 = rate_limit_next_token_us(rl, 500000)
126 if wait_us != 500000 { return 8 }
127
128 // 5 seconds later, bucket refills up to capacity.
129 if rate_limit_available(rl, 5500000) != 10 { return 9 }
130
131 // Time going backwards: harmless.
132 if rate_limit_try_take(rl, 5, 5500000) != 1 { return 10 }
133 if rate_limit_available(rl, 100) != 5 { return 11 }
134
135 return 0
136}