1"""Chapter 7 practice: load-balancing algorithms and health checks (Level 3 Build)
2and the 'IP hash with mod N' bug (Level 4 Debug). Standard library only.
3Run: python3 ch07_load_balancer.py"""
4import bisect
5import hashlib
6import random
7
8
9# ---- Algorithms ---------------------------------------------------------------------------
10
11class RoundRobin:
12 def __init__(self, servers):
13 self.servers, self.i = list(servers), 0
14
15 def pick(self, **_):
16 s = self.servers[self.i % len(self.servers)]
17 self.i += 1
18 return s
19
20
21class WeightedRoundRobin:
22 """Smooth weighted round robin: a server with weight 2 gets twice the requests,
23 spread out (A A B, not A A ... B B) rather than in bursts."""
24
25 def __init__(self, weights):
26 self.weights = dict(weights)
27 self.current = {s: 0 for s in weights}
28
29 def pick(self, **_):
30 total = sum(self.weights.values())
31 for s, w in self.weights.items():
32 self.current[s] += w
33 best = max(self.current, key=self.current.get)
34 self.current[best] -= total
35 return best
36
37
38class LeastConnections:
39 def __init__(self, servers):
40 self.active = {s: 0 for s in servers}
41
42 def pick(self, **_):
43 return min(self.active, key=self.active.get)
44
45
46def stable_hash(text):
47 return int(hashlib.md5(text.encode()).hexdigest(), 16)
48
49
50class IPHash:
51 def __init__(self, servers):
52 self.servers = list(servers)
53
54 def pick(self, ip=None, **_):
55 return self.servers[stable_hash(ip) % len(self.servers)] # mod N (see Level 4)
56
57
58# ---- Health checks ------------------------------------------------------------------------
59
60class HealthChecker:
61 """Marks a server down after `unhealthy` failed checks in a row,
62 and back up after `healthy` passing checks in a row."""
63
64 def __init__(self, servers, unhealthy=3, healthy=2):
65 self.up = {s: True for s in servers}
66 self.fails = {s: 0 for s in servers}
67 self.passes = {s: 0 for s in servers}
68 self.unhealthy, self.healthy = unhealthy, healthy
69
70 def record(self, server, ok):
71 if ok:
72 self.passes[server] += 1
73 self.fails[server] = 0
74 if not self.up[server] and self.passes[server] >= self.healthy:
75 self.up[server] = True
76 else:
77 self.fails[server] += 1
78 self.passes[server] = 0
79 if self.up[server] and self.fails[server] >= self.unhealthy:
80 self.up[server] = False
81
82 def healthy_servers(self):
83 return [s for s, ok in self.up.items() if ok]
84
85
86# ---- Simulation: round robin vs least connections with uneven request lengths ------------------
87
88def simulate(balancer, n_servers=4, n_requests=20_000, seed=3):
89 """Requests arrive every 1 time unit. 90% take 2 units, 10% take 60 units (slow searches).
90 Returns the highest number of requests any server had in progress at once."""
91 rng = random.Random(seed)
92 finishing = [] # (finish_time, server)
93 active = {f"s{i}": 0 for i in range(n_servers)}
94 if isinstance(balancer, LeastConnections):
95 balancer.active = active
96 peak = 0
97 for t in range(n_requests):
98 while finishing and finishing[0][0] <= t:
99 _, s = finishing.pop(0)
100 active[s] -= 1
101 s = balancer.pick()
102 active[s] += 1
103 peak = max(peak, active[s])
104 duration = 60 if rng.random() < 0.10 else 2
105 bisect.insort(finishing, (t + duration, s))
106 return peak
107
108
109# ---- Level 4: IP hash with mod N --------------------------------------------------------------
110
111def moved_fraction(n_before, n_after, users=20_000):
112 before = IPHash([f"s{i}" for i in range(n_before)])
113 after = IPHash([f"s{i}" for i in range(n_after)])
114 ips = [f"10.0.{u // 256}.{u % 256}" for u in range(users)]
115 return sum(before.pick(ip=ip) != after.pick(ip=ip) for ip in ips) / users
116
117
118class HashRing:
119 """The fix (Chapter 28): consistent hashing with virtual nodes."""
120
121 def __init__(self, servers, vnodes=200):
122 self.ring = sorted((stable_hash(f"{s}#{v}"), s) for s in servers for v in range(vnodes))
123 self.keys = [k for k, _ in self.ring]
124
125 def pick(self, ip):
126 i = bisect.bisect(self.keys, stable_hash(ip)) % len(self.ring)
127 return self.ring[i][1]
128
129
130def ring_moved_fraction(n_before, n_after, users=20_000):
131 before = HashRing([f"s{i}" for i in range(n_before)])
132 after = HashRing([f"s{i}" for i in range(n_after)])
133 ips = [f"10.0.{u // 256}.{u % 256}" for u in range(users)]
134 return sum(before.pick(ip) != after.pick(ip) for ip in ips) / users
135
136
137def run_checks():
138 rr = RoundRobin(["a", "b", "c"])
139 assert [rr.pick() for _ in range(6)] == ["a", "b", "c", "a", "b", "c"]
140
141 w = WeightedRoundRobin({"big": 2, "small": 1})
142 picks = [w.pick() for _ in range(300)]
143 assert picks.count("big") == 200 and picks.count("small") == 100
144 assert picks[:3] in (["big", "small", "big"], ["big", "big", "small"])
145
146 ip = IPHash(["a", "b", "c"])
147 assert ip.pick(ip="203.0.113.7") == ip.pick(ip="203.0.113.7") # same user, same server
148
149 hc = HealthChecker(["a", "b"])
150 for ok in (False, False):
151 hc.record("a", ok)
152 assert hc.healthy_servers() == ["a", "b"] # 2 failures: still up
153 hc.record("a", False)
154 assert hc.healthy_servers() == ["b"] # 3rd failure: down
155 hc.record("a", True)
156 assert hc.healthy_servers() == ["b"] # 1 pass: not yet
157 hc.record("a", True)
158 assert hc.healthy_servers() == ["a", "b"] # 2 passes: back
159
160 peak_rr = simulate(RoundRobin([f"s{i}" for i in range(4)]))
161 peak_lc = simulate(LeastConnections([f"s{i}" for i in range(4)]))
162 assert peak_lc < peak_rr, (peak_rr, peak_lc)
163
164 # Level 4: removing 1 of 4 servers with mod N moves most users; a hash ring moves about 1 in 4.
165 bad = moved_fraction(4, 3)
166 good = ring_moved_fraction(4, 3)
167 assert bad > 0.70 and good < 0.30, (bad, good)
168
169 print(f"All checks pass. Busiest server, round robin: {peak_rr} in progress; "
170 f"least connections: {peak_lc}.")
171 print(f"Removing 1 of 4 servers: mod-N IP hash moved {bad:.0%} of users, hash ring {good:.0%}.")
172
173
174if __name__ == "__main__":
175 run_checks()