eden-api-gateway/registry/host.py

import asyncio
import signal
import socketio
from aiohttp import web

SERVER_0_IP = "192.168.1.113"
FLASK_SERVER_PORT = 4999
HEARTBEAT_INTERVAL = 5
HEARTBEAT_TIMEOUT = 3

sio = socketio.AsyncServer(async_mode='aiohttp')
app = web.Application()
sio.attach(app)

servers = {}

async def available(request):
    return web.json_response(servers)

app.router.add_get("/available", available)

@sio.event
async def connect(sid, environ):
    print("I'm connected!", sid)

@sio.event
async def register(sid, data):
    server_info = data
    name = server_info["name"]

    servers[name] = {"ip": server_info["ip"], "port": server_info["port"], "sid": sid}
    print(servers)

@sio.event
async def disconnect(sid):
    print("I'm disconnected!", sid)
    for name, server in servers.items():
        if server["sid"] == sid:
            del servers[name]
            break

async def heartbeat():
    while True:
        await asyncio.sleep(HEARTBEAT_INTERVAL)
        server_values_copy = list(servers.values())
        for server in server_values_copy:
            sid = server["sid"]
            try:
                print(f"Sending heartbeat to {sid}...")
                heartbeat_future = sio.emit("heartbeat", to=sid)
                await asyncio.wait_for(heartbeat_future, timeout=HEARTBEAT_TIMEOUT)
            except (asyncio.TimeoutError, socketio.exceptions.TimeoutError):
                print(f"Server {sid} failed to respond to heartbeat.")
                await sio.disconnect(sid)

def exit_handler(sig, frame):
    print("Shutting down host...")
    loop = asyncio.get_event_loop()
    heartbeat_task.cancel()
    loop.run_until_complete(loop.shutdown_asyncgens())
    loop.stop()

if __name__ == "__main__":
    signal.signal(signal.SIGINT, exit_handler)
    signal.signal(signal.SIGTERM, exit_handler)

    loop = asyncio.get_event_loop()
    heartbeat_task = loop.create_task(heartbeat())
    aiohttp_app = loop.create_task(web._run_app(app, host=SERVER_0_IP, port=FLASK_SERVER_PORT))
    
    try:
        loop.run_until_complete(asyncio.gather(heartbeat_task, aiohttp_app))
    except asyncio.CancelledError:
        pass
    finally:
        loop.run_until_complete(loop.shutdown_asyncgens())
        loop.stop()