1"""Chapter 32 (case study: rate limiter as a service) deep-dive code.
21. A token bucket that runs as one atomic step (what a Redis Lua script does).
32. The read-then-write race, and checking several rules at once without leaking tokens.
43. Limits across regions: splitting a global limit, and rebalancing shares.
5Run: python3 ch32_rate_limiter_service.py"""
6import math
7
8
9# ---- a fake Redis: one dict, and "scripts" that run without interruption ----------
10
11class FakeRedis:
12 def __init__(self):
13 self.data = {}
14 self.ops = 0
15
16 def get(self, k):
17 self.ops += 1
18 return self.data.get(k)
19
20 def set(self, k, v):
21 self.ops += 1
22 self.data[k] = v
23
24 def eval(self, script, keys, args):
25 """Real Redis runs a Lua script atomically: no other command runs in between."""
26 self.ops += 1
27 return script(self.data, keys, args)
28
29
30# ---- 1. token bucket as a script ---------------------------------------------------
31
32def token_bucket_script(data, keys, args):
33 """keys: [bucket key]; args: [capacity, refill_per_s, now, cost]."""
34 capacity, rate, now, cost = args
35 tokens, last = data.get(keys[0], (capacity, now))
36 tokens = min(capacity, tokens + (now - last) * rate)
37 allowed = tokens >= cost
38 if allowed:
39 tokens -= cost
40 data[keys[0]] = (tokens, now)
41 retry_after = 0 if allowed else math.ceil((cost - tokens) / rate)
42 return allowed, int(tokens), retry_after
43
44
45def check(redis, rule, key, now, cost=1):
46 return redis.eval(token_bucket_script, [f"rl:{{{rule['id']}:{key}}}"],
47 [rule["limit"], rule["limit"] / rule["window_s"], now, cost])
48
49
50# ---- 2a. the race ------------------------------------------------------------------
51
52def racy_checks(redis, key, limit, concurrent):
53 """Every limiter instance GETs, decides, then SETs: all decide on the same old value."""
54 seen = [redis.get(key) or 0 for _ in range(concurrent)]
55 allowed = 0
56 for count in seen:
57 if count < limit:
58 redis.set(key, count + 1)
59 allowed += 1
60 return allowed
61
62
63def atomic_checks(redis, key, limit, concurrent):
64 def incr_if_under(data, keys, args):
65 c = data.get(keys[0], 0)
66 if c < args[0]:
67 data[keys[0]] = c + 1
68 return True
69 return False
70 return sum(redis.eval(incr_if_under, [key], [limit]) for _ in range(concurrent))
71
72
73# ---- 2b. several rules for one request ---------------------------------------------
74
75def multi_rule_naive(redis, rules_keys, now):
76 """Checks rules one by one: an earlier rule's token is spent even if a later rule says no."""
77 for rule, key in rules_keys:
78 allowed, _, _ = check(redis, rule, key, now)
79 if not allowed:
80 return False
81 return True
82
83
84def multi_rule_atomic(redis, rules_keys, now):
85 """One script: check every bucket first, spend only if all allow."""
86 def script(data, keys, args):
87 states = []
88 for k, (cap, rate) in zip(keys, args):
89 tokens, last = data.get(k, (cap, now))
90 states.append((k, min(cap, tokens + (now - last) * rate)))
91 if all(t >= 1 for _, t in states):
92 for k, t in states:
93 data[k] = (t - 1, now)
94 return True
95 for k, t in states:
96 data[k] = (t, now)
97 return False
98 keys = [f"rl:{{{r['id']}:{k}}}" for r, k in rules_keys]
99 args = [(r["limit"], r["limit"] / r["window_s"]) for r, _ in rules_keys]
100 return redis.eval(script, keys, args)
101
102
103# ---- 3. across regions -------------------------------------------------------------
104
105def split_evenly(limit, regions):
106 return {r: limit // len(regions) for r in regions}
107
108
109def served(shares, demand):
110 return sum(min(shares[r], demand[r]) for r in demand)
111
112
113def rebalance(limit, recent_demand, floor=0.1):
114 """Give each region a share in proportion to recent demand, with a small floor."""
115 regions = list(recent_demand)
116 base = int(limit * floor / len(regions))
117 rest = limit - base * len(regions)
118 total = sum(recent_demand.values()) or 1
119 return {r: base + int(rest * recent_demand[r] / total) for r in regions}
120
121
122# ---- checks ------------------------------------------------------------------------
123
124def run_checks():
125 partner = {"id": "partner-api", "limit": 100, "window_s": 60}
126 login = {"id": "login-account", "limit": 5, "window_s": 900}
127 login_ip = {"id": "login-ip", "limit": 50, "window_s": 900}
128
129 # 1. token bucket: burst of 100, then 100 a minute; retry-after reported
130 r = FakeRedis()
131 results = [check(r, partner, "key42", now=0) for _ in range(101)]
132 assert sum(a for a, _, _ in results) == 100
133 assert results[-1] == (False, 0, 1) # next token in ~0.6 s → retry in 1 s
134 assert check(r, partner, "key42", now=0.6)[0] is True
135 assert r.ops == 102 # one round trip per check
136
137 # 2a. race: 20 checks arrive with 95 used
138 r = FakeRedis(); r.data["c"] = 95
139 assert racy_checks(r, "c", 100, 20) == 20 # 15 too many
140 r = FakeRedis(); r.data["c"] = 95
141 assert atomic_checks(r, "c", 100, 20) == 5
142
143 # 2b. two rules: per-account (5) and per-IP (50). Account is exhausted.
144 r = FakeRedis()
145 for _ in range(5):
146 check(r, login, "asha", now=0)
147 pair = [(login_ip, "1.2.3.4"), (login, "asha")]
148 for _ in range(10):
149 multi_rule_naive(r, pair, now=1)
150 ip_tokens_naive = r.data["rl:{login-ip:1.2.3.4}"][0]
151 assert round(ip_tokens_naive) == 40 # 10 IP tokens spent on denied requests
152 r2 = FakeRedis()
153 for _ in range(5):
154 check(r2, login, "asha", now=0)
155 for _ in range(10):
156 assert multi_rule_atomic(r2, pair, now=1) is False
157 assert round(r2.data["rl:{login-ip:1.2.3.4}"][0]) == 50 # nothing leaked
158
159 # 3. regions: global limit 900 a minute for one big partner, across 3 regions
160 regions = ["mumbai", "frankfurt", "virginia"]
161 shares = split_evenly(900, regions)
162 demand = {"mumbai": 600, "frankfurt": 200, "virginia": 100}
163 assert served(shares, demand) == 600 # 300 + 200 + 100: a third wasted
164 better = rebalance(900, demand)
165 assert better == {"mumbai": 570, "frankfurt": 210, "virginia": 120}
166 assert served(better, demand) == 870 # 97% of demand, vs 67% split evenly
167
168 # sizing (Step 2 arithmetic)
169 assert 200_000 * 1 == 200_000 # checks a second = Redis ops a second
170 assert math.ceil(200_000 / 50_000) == 4 # shards at an assumed 50k ops/s headroom each
171 assert 10_000_000 * 100 / 1e9 == 1.0 # GB for 10M buckets at ~100 B each
172
173 return {"race": (20, 5), "ip_tokens_naive_vs_atomic": (40, 50),
174 "regions_even_vs_rebalanced": (600, served(better, demand)), "shares": better}
175
176
177if __name__ == "__main__":
178 print("All checks pass.", run_checks())