1"""Chapter 25: Rate limiting. Tested with Python 3, standard library only.
2
3Running example: the bookstore's partner API allows 100 requests a minute per API key,
4and the login endpoint allows 5 attempts per 15 minutes per account.
5"""
6import math
7from collections import deque
8
9
10# ---------- 25.2.1 token bucket ----------
11
12class TokenBucket:
13 def __init__(self, capacity, refill_per_s):
14 self.capacity, self.rate = capacity, refill_per_s
15 self.tokens, self.last = capacity, 0.0
16
17 def allow(self, now, cost=1):
18 self.tokens = min(self.capacity, self.tokens + (now - self.last) * self.rate)
19 self.last = now
20 if self.tokens >= cost:
21 self.tokens -= cost
22 return True
23 return False
24
25
26# ---------- 25.2.2 leaky bucket (as a queue that drains at a fixed rate) ----------
27
28class LeakyBucket:
29 def __init__(self, capacity, leak_per_s):
30 self.capacity, self.rate = capacity, leak_per_s
31 self.level, self.last = 0.0, 0.0
32
33 def offer(self, now):
34 """Returns the delay before this request is processed, or None if dropped."""
35 self.level = max(0.0, self.level - (now - self.last) * self.rate)
36 self.last = now
37 if self.level + 1 > self.capacity:
38 return None
39 self.level += 1
40 return (self.level - 1) / self.rate # wait behind the ones already queued
41
42
43# ---------- 25.2.3 fixed window ----------
44
45class FixedWindow:
46 def __init__(self, limit, window):
47 self.limit, self.window = limit, window
48 self.counts = {}
49
50 def allow(self, now):
51 w = int(now // self.window)
52 self.counts[w] = self.counts.get(w, 0) + 1
53 return self.counts[w] <= self.limit
54
55
56# ---------- 25.2.4 sliding log ----------
57
58class SlidingLog:
59 def __init__(self, limit, window):
60 self.limit, self.window, self.log = limit, window, deque()
61
62 def allow(self, now):
63 while self.log and self.log[0] <= now - self.window:
64 self.log.popleft()
65 if len(self.log) < self.limit:
66 self.log.append(now)
67 return True
68 return False
69
70
71# ---------- 25.2.5 sliding window counter ----------
72
73class SlidingWindowCounter:
74 def __init__(self, limit, window):
75 self.limit, self.window = limit, window
76 self.counts = {}
77
78 def estimate(self, now):
79 w = int(now // self.window)
80 elapsed = (now % self.window) / self.window
81 return self.counts.get(w - 1, 0) * (1 - elapsed) + self.counts.get(w, 0)
82
83 def allow(self, now):
84 if self.estimate(now) + 1 > self.limit:
85 return False
86 w = int(now // self.window)
87 self.counts[w] = self.counts.get(w, 0) + 1
88 return True
89
90
91def boundary_burst(limiter):
92 """100 requests at 59.9-59.99 s, then 100 at 60.0-60.09 s. How many get through?"""
93 allowed = 0
94 for i in range(100):
95 allowed += limiter.allow(59.9 + i * 0.0009)
96 for i in range(100):
97 allowed += limiter.allow(60.0 + i * 0.0009)
98 return allowed
99
100
101# ---------- 25.4 distributed: a shared counter ----------
102
103class FakeRedis:
104 def __init__(self):
105 self.data = {}
106
107 def get(self, k):
108 return self.data.get(k, 0)
109
110 def set(self, k, v):
111 self.data[k] = v
112
113 def incr(self, k): # atomic in real Redis
114 self.data[k] = self.data.get(k, 0) + 1
115 return self.data[k]
116
117
118def racy_allow_many(r, key, limit, concurrent):
119 """Every server reads the count first, then each writes count+1."""
120 reads = [r.get(key) for _ in range(concurrent)] # all read before anyone writes
121 allowed = 0
122 for c in reads:
123 if c < limit:
124 r.set(key, c + 1)
125 allowed += 1
126 return allowed
127
128
129def atomic_allow_many(r, key, limit, concurrent):
130 return sum(1 for _ in range(concurrent) if r.incr(key) <= limit)
131
132
133def local_limits(per_node_limit, requests_per_node):
134 """Each node enforces limit/N on its own share of one client's traffic."""
135 return sum(min(per_node_limit, n) for n in requests_per_node)
136
137
138# ---------- 25.6 headers ----------
139
140def headers(limit, used, window_reset_in):
141 remaining = max(0, limit - used)
142 h = {"RateLimit-Limit": limit, "RateLimit-Remaining": remaining, "RateLimit-Reset": window_reset_in}
143 if remaining == 0:
144 h["Retry-After"] = window_reset_in
145 return h
146
147
148# ---------- 25.7 bots: login attempts ----------
149
150def stuffing_attack(attempts, ips, per_ip_limit, per_account_limit):
151 """Many IPs, one account. How many guesses get through under each limit?"""
152 per_ip = [0] * ips
153 per_ip_ok = per_acct_ok = 0
154 for i in range(attempts):
155 ip = i % ips
156 per_ip[ip] += 1
157 if per_ip[ip] <= per_ip_limit:
158 per_ip_ok += 1
159 if i < per_account_limit:
160 per_acct_ok += 1
161 return per_ip_ok, per_acct_ok
162
163
164# ---------- checks ----------
165
166def run_checks():
167 # token bucket: burst of 100, then 100 per minute
168 tb = TokenBucket(capacity=100, refill_per_s=100 / 60)
169 assert sum(tb.allow(0) for _ in range(150)) == 100 # burst capped at capacity
170 assert tb.allow(0.5) is False and tb.allow(0.61) is True # 1 token every 0.6 s
171
172 # leaky bucket: smooth output, queue of 10, drains 2 per second
173 lb = LeakyBucket(capacity=10, leak_per_s=2)
174 delays = [lb.offer(0) for _ in range(12)]
175 assert delays[:3] == [0.0, 0.5, 1.0] and delays[-2:] == [None, None]
176
177 # the boundary problem
178 limit, window = 100, 60
179 fixed = boundary_burst(FixedWindow(limit, window))
180 log = boundary_burst(SlidingLog(limit, window))
181 counter = boundary_burst(SlidingWindowCounter(limit, window))
182 bucket = boundary_burst(TokenBucket(100, 100 / 60))
183 assert fixed == 200 # double the limit in 0.2 s
184 assert log == 100
185 assert counter == 100
186 assert bucket == 100
187
188 # sliding window counter is an estimate: halfway through, half the old window counts
189 swc = SlidingWindowCounter(100, 60)
190 swc.counts = {0: 80, 1: 30}
191 assert swc.estimate(90) == 80 * 0.5 + 30 # 70
192
193 # sliding log memory: one timestamp per request
194 keys, per_key = 1_000_000, 100
195 assert keys * per_key * 8 / 1e9 == 0.8 # GB for timestamps alone
196 assert keys * 2 * 8 / 1e6 == 16 # MB for two counters per key
197
198 # distributed: read-then-write races; INCR doesn't
199 r = FakeRedis()
200 r.set("rl:key42:min", 95)
201 assert racy_allow_many(r, "rl:key42:min", 100, concurrent=20) == 20 # 15 too many
202 r = FakeRedis()
203 r.set("rl:key42:min", 95)
204 assert atomic_allow_many(r, "rl:key42:min", 100, concurrent=20) == 5 # exactly the 5 left
205
206 # local per-node limits: 10 nodes x 10, but the client's traffic lands unevenly
207 assert local_limits(10, [30, 25, 20, 10, 5, 4, 3, 2, 1, 0]) == 55 # sends 100, gets only 55 through
208
209 # headers
210 h = headers(100, 100, 23)
211 assert h["RateLimit-Remaining"] == 0 and h["Retry-After"] == 23
212 assert "Retry-After" not in headers(100, 40, 23)
213
214 # credential stuffing: 10,000 guesses from 1,000 IPs at one account
215 per_ip_ok, per_acct_ok = stuffing_attack(10_000, 1_000, per_ip_limit=10, per_account_limit=5)
216 assert per_ip_ok == 10_000 and per_acct_ok == 5 # per-IP limits miss it entirely
217
218 # Level 2 arithmetic
219 assert math.isclose(60 / 100, 0.6) # seconds per token
220 assert 1_000_000 / 60 > 16_000 # Redis ops/s for 1M req/min
221
222 return {"boundary": {"fixed": fixed, "log": log, "counter": counter, "bucket": bucket},
223 "race": (20, 5), "stuffing": (per_ip_ok, per_acct_ok)}
224
225
226if __name__ == "__main__":
227 print("All checks pass.", run_checks())