130 lines
4.2 KiB
Python
130 lines
4.2 KiB
Python
import asyncio
|
|
import sys
|
|
|
|
|
|
class LoadBalancer:
|
|
def __init__(self, backend_ports):
|
|
self.backend_ports = backend_ports
|
|
self.backend_processes = []
|
|
self.client_counts = [0] * len(backend_ports)
|
|
self.lock = asyncio.Lock()
|
|
|
|
async def start_backend_servers(self, port, workers):
|
|
for x in range(workers):
|
|
port += 1
|
|
process = await asyncio.create_subprocess_exec(
|
|
sys.executable,
|
|
sys.argv[0],
|
|
"backend",
|
|
str(port),
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
)
|
|
port += 1
|
|
self.backend_processes.append(process)
|
|
print(
|
|
f"Started backend server on port {(port-1)/port} with PID {process.pid}"
|
|
)
|
|
|
|
async def handle_client(self, reader, writer):
|
|
async with self.lock:
|
|
min_clients = min(self.client_counts)
|
|
server_index = self.client_counts.index(min_clients)
|
|
self.client_counts[server_index] += 1
|
|
backend = ("127.0.0.1", self.backend_ports[server_index])
|
|
try:
|
|
backend_reader, backend_writer = await asyncio.open_connection(*backend)
|
|
|
|
async def forward(r, w):
|
|
try:
|
|
while True:
|
|
data = await r.read(1024)
|
|
if not data:
|
|
break
|
|
w.write(data)
|
|
await w.drain()
|
|
except asyncio.CancelledError:
|
|
pass
|
|
finally:
|
|
w.close()
|
|
|
|
task1 = asyncio.create_task(forward(reader, backend_writer))
|
|
task2 = asyncio.create_task(forward(backend_reader, writer))
|
|
await asyncio.gather(task1, task2)
|
|
except Exception as e:
|
|
print(f"Error: {e}")
|
|
finally:
|
|
writer.close()
|
|
async with self.lock:
|
|
self.client_counts[server_index] -= 1
|
|
|
|
async def monitor(self):
|
|
while True:
|
|
await asyncio.sleep(5)
|
|
print("Connected clients per server:")
|
|
for i, count in enumerate(self.client_counts):
|
|
print(f"Server {self.backend_ports[i]}: {count} clients")
|
|
|
|
async def start(self, host="0.0.0.0", port=8081, workers=5):
|
|
await self.start_backend_servers(port, workers)
|
|
server = await asyncio.start_server(self.handle_client, host, port)
|
|
asyncio.create_task(self.monitor())
|
|
|
|
# Handle shutdown gracefully
|
|
try:
|
|
async with server:
|
|
await server.serve_forever()
|
|
except asyncio.CancelledError:
|
|
pass
|
|
finally:
|
|
# Terminate backend processes
|
|
for process in self.backend_processes:
|
|
process.terminate()
|
|
await asyncio.gather(*(p.wait() for p in self.backend_processes))
|
|
print("Backend processes terminated.")
|
|
|
|
|
|
async def backend_echo_server(port):
|
|
async def handle_echo(reader, writer):
|
|
try:
|
|
while True:
|
|
data = await reader.read(1024)
|
|
if not data:
|
|
break
|
|
writer.write(data)
|
|
await writer.drain()
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
writer.close()
|
|
|
|
server = await asyncio.start_server(handle_echo, "127.0.0.1", port)
|
|
print(f"Backend echo server running on port {port}")
|
|
await server.serve_forever()
|
|
|
|
|
|
async def main():
|
|
backend_ports = [8001, 8003, 8005, 8006]
|
|
# Launch backend echo servers
|
|
# Wait a moment for servers to start
|
|
lb = LoadBalancer(backend_ports)
|
|
await lb.start()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) > 1:
|
|
if sys.argv[1] == "backend":
|
|
port = int(sys.argv[2])
|
|
from snek.app import Application
|
|
|
|
snek = Application(port=port)
|
|
web.run_app(snek, port=port, host="127.0.0.1")
|
|
elif sys.argv[1] == "sync":
|
|
|
|
web.run_app(snek, port=port, host="127.0.0.1")
|
|
else:
|
|
try:
|
|
asyncio.run(main())
|
|
except KeyboardInterrupt:
|
|
print("Shutting down...")
|