Pythonload_balancer.py
Copy
1
import threading
2
import time
3
import hashlib
4
import bisect
5
import random
6
from dataclasses import dataclass, field
7
from collections import Counter
8
 
9
@dataclass
10
class Backend:
11
    host: str
12
    port: int
13
    weight: int = 1
14
    active_connections: int = 0
15
    healthy: bool = True
16
    request_count: int = 0
17
    
18
    @property
19
    def address(self) -> str:
20
        return f"{self.host}:{self.port}"
21
 
22
# ── Round-Robin ──────────────────────────────────────────────────────
23
class RoundRobinBalancer:
24
    def __init__(self, backends: list[Backend]):
25
        self._backends = backends
26
        self._index = 0
27
        self._lock = threading.Lock()
28
    
29
    def next(self) -> Backend | None:
30
        with self._lock:
31
            healthy = [b for b in self._backends if b.healthy]
32
            if not healthy:
33
                return None
34
            backend = healthy[self._index % len(healthy)]
35
            self._index = (self._index + 1) % len(healthy)
36
            backend.request_count += 1
37
            return backend
38
 
39
# ── Weighted Round-Robin ─────────────────────────────────────────────
40
class WeightedRoundRobinBalancer:
41
    def __init__(self, backends: list[Backend]):
42
        self._backends = backends
43
        self._lock = threading.Lock()
44
        self._current_weights = {b.address: 0 for b in backends}
45
    
46
    def next(self) -> Backend | None:
47
        """Smooth Weighted Round-Robin (Nginx algorithm)."""
48
        with self._lock:
49
            healthy = [b for b in self._backends if b.healthy]
50
            if not healthy:
51
                return None
52
            
53
            total_weight = sum(b.weight for b in healthy)
54
            
55
            for b in healthy:
56
                self._current_weights[b.address] = self._current_weights.get(b.address, 0) + b.weight
57
            
58
            best = max(healthy, key=lambda b: self._current_weights.get(b.address, 0))
59
            self._current_weights[best.address] -= total_weight
60
            best.request_count += 1
61
            return best
62
 
63
# ── Least Connections ────────────────────────────────────────────────
64
class LeastConnectionsBalancer:
65
    def __init__(self, backends: list[Backend]):
66
        self._backends = backends
67
        self._lock = threading.Lock()
68
    
69
    def next(self) -> Backend | None:
70
        with self._lock:
71
            healthy = [b for b in self._backends if b.healthy]
72
            if not healthy:
73
                return None
74
            backend = min(healthy, key=lambda b: b.active_connections)
75
            backend.active_connections += 1
76
            backend.request_count += 1
77
            return backend
78
    
79
    def release(self, backend: Backend):
80
        with self._lock:
81
            backend.active_connections = max(0, backend.active_connections - 1)
82
 
83
# ── Consistent Hashing ────────────────────────────────────────────────
84
class ConsistentHashRing:
85
    def __init__(self, backends: list[Backend], virtual_nodes: int = 100):
86
        self._ring: dict[int, Backend] = {}
87
        self._sorted_keys: list[int] = []
88
        self._lock = threading.Lock()
89
        
90
        for backend in backends:
91
            self.add_backend(backend, virtual_nodes)
92
    
93
    def _hash(self, key: str) -> int:
94
        return int(hashlib.md5(key.encode()).hexdigest(), 16)
95
    
96
    def add_backend(self, backend: Backend, virtual_nodes: int = 100):
97
        with self._lock:
98
            for i in range(virtual_nodes):
99
                key = self._hash(f"{backend.address}:{i}")
100
                self._ring[key] = backend
101
                bisect.insort(self._sorted_keys, key)
102
    
103
    def get(self, key: str) -> Backend | None:
104
        """Route a request key to the correct backend."""
105
        with self._lock:
106
            if not self._ring:
107
                return None
108
            hash_key = self._hash(key)
109
            idx = bisect.bisect(self._sorted_keys, hash_key) % len(self._sorted_keys)
110
            backend = self._ring[self._sorted_keys[idx]]
111
            backend.request_count += 1
112
            return backend
113
 
114
# ── Simulation ────────────────────────────────────────────────────────
115
backends_rr = [Backend("10.0.0.1", 8080), Backend("10.0.0.2", 8080), Backend("10.0.0.3", 8080)]
116
backends_w = [Backend("10.0.0.1", 8080, weight=1), Backend("10.0.0.2", 8080, weight=2), Backend("10.0.0.3", 8080, weight=3)]
117
backends_lc = [Backend("10.0.0.1", 8080), Backend("10.0.0.2", 8080), Backend("10.0.0.3", 8080)]
118
backends_ch = [Backend("10.0.0.1", 8080), Backend("10.0.0.2", 8080), Backend("10.0.0.3", 8080)]
119
 
120
rr = RoundRobinBalancer(backends_rr)
121
wrr = WeightedRoundRobinBalancer(backends_w)
122
lc = LeastConnectionsBalancer(backends_lc)
123
ch = ConsistentHashRing(backends_ch)
124
 
125
N = 600
126
for i in range(N):
127
    rr.next()
128
    wrr.next()
129
    lc.next()
130
    ch.get(f"user-{random.randint(1, 100)}")
131
 
132
print("=== Load Distribution (N=600) ===")
133
print(f"Round-Robin:          {[b.request_count for b in backends_rr]}")
134
print(f"Weighted(1:2:3):       {[b.request_count for b in backends_w]}")
135
print(f"Least-Connections:     {[b.request_count for b in backends_lc]}")
136
print(f"Consistent Hashing:    {[b.request_count for b in backends_ch]}")
137
 
138
# Consistent hashing: same key always routes to same backend
139
print("\n=== Consistent Hashing Affinity ===")
140
for user in ["user-42", "user-7", "user-99"]:
141
    backend = ch.get(user)
142
    print(f"  {user} → always {backend.address}")
143