Remove AIO HTTP

hyperforge/pyproject.toml

@@ -44,6 +44,7 @@ dependencies = [
     "sentry_sdk",
     "fastapi",
     "starlette>1.0.0",
+    "uvicorn>=0.30.0",
 ]
 
 

hyperforge/src/hyperforge/server/session.py

@@ -10,7 +10,7 @@
 import nucliadb_telemetry.context
 import nucliadb_telemetry.metrics
 import prometheus_client
-from aiohttp.web import Server
+from hyperforge.server.web import WebServer
 from lru import LRU
 from nucliadb_telemetry import errors
 from nucliadb_telemetry.utils import get_telemetry
@@ -62,7 +62,7 @@ def tracer():
 
 
 class SessionManager:
-    server: Optional[Server] = None
+    server: Optional[WebServer] = None
     tasks: List[Task]
     hooks: Optional[Dict[str, List[Callable]]] = None
 

hyperforge/src/hyperforge/server/web.py

@@ -1,35 +1,66 @@
+import asyncio
 import os
 
 import prometheus_client  # type: ignore
-from aiohttp import web
+import uvicorn
+from starlette.applications import Starlette
+from starlette.requests import Request
+from starlette.responses import PlainTextResponse, Response
+from starlette.routing import Route
 
 from hyperforge.server import logger
 
 PORT = os.environ.get("HEALTH_CHECK_PORT", "8000")
 
 
-async def http_handler(request: web.Request):
-    if request.path == "/metrics":
-        output = prometheus_client.exposition.generate_latest()
-        return web.Response(text=output.decode("utf8"))
-    elif request.path in ("/health/alive", "/health/ready"):
-        # implement health check here
-        return web.Response(text="OK")
-    else:
-        return web.Response(text="OK", status=404)
+async def metrics(request: Request) -> Response:
+    output = prometheus_client.exposition.generate_latest()
+    return PlainTextResponse(output.decode("utf8"))
 
 
-async def start_web_server() -> web.Server:
-    server = web.Server(http_handler)  # type: ignore
-    runner = web.ServerRunner(server)
-    await runner.setup()
-    site = web.TCPSite(runner, "0.0.0.0", int(PORT))
-    await site.start()
+async def health(request: Request) -> Response:
+    return PlainTextResponse("OK")
 
+
+async def not_found(request: Request, exc: Exception) -> Response:
+    return PlainTextResponse("OK", status_code=404)
+
+
+app = Starlette(
+    routes=[
+        Route("/metrics", metrics),
+        Route("/health/alive", health),
+        Route("/health/ready", health),
+    ],
+    exception_handlers={404: not_found},
+)
+
+
+class WebServer:
+    """Wraps a uvicorn server running as an asyncio task in the current event loop."""
+
+    def __init__(self, server: uvicorn.Server) -> None:
+        self._server = server
+        self._task: asyncio.Task | None = None
+
+    def start(self) -> None:
+        self._task = asyncio.get_event_loop().create_task(self._server.serve())
+
+    async def shutdown(self) -> None:
+        self._server.should_exit = True
+        if self._task is not None:
+            await self._task
+            self._task = None
+
+
+async def start_web_server() -> WebServer:
+    config = uvicorn.Config(app, host="0.0.0.0", port=int(PORT), log_level="warning")
+    server = uvicorn.Server(config)
+    web = WebServer(server)
+    web.start()
     logger.info(f"======= Serving on http://0.0.0.0:{PORT}/ ======")
-    return server
+    return web
 
 
-async def start_health_check():
-    server = await start_web_server()
-    return server
+async def start_health_check() -> WebServer:
+    return await start_web_server()