From 59e15be2d74f566dea1240a74ae71449fc9a13c5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=A9bastien=20Pierre?= <sebastien.pierre@gmail.com>
Date: Wed, 2 Oct 2024 13:12:10 +1300
Subject: [PATCH] [Update] server: added server state with exception logging

---
 src/py/extra/server.py | 26 +++++++++++++++++++++++++-
 1 file changed, 25 insertions(+), 1 deletion(-)

diff --git a/src/py/extra/server.py b/src/py/extra/server.py
index 1b7aba1..262f6f8 100644
--- a/src/py/extra/server.py
+++ b/src/py/extra/server.py
@@ -1,5 +1,7 @@
 from typing import Callable, NamedTuple, Any, Coroutine, Literal
 from pathlib import Path
+from signal import SIGINT, SIGTERM
+from dataclasses import dataclass
 import socket
 import asyncio
 from .utils.logging import exception, info, warning, event
@@ -17,6 +19,21 @@
 from .config import HOST, PORT
 
 
+@dataclass(slots=True)
+class ServerState:
+    isRunning: bool = True
+
+    def stop(self) -> None:
+        self.isRunning = False
+
+    def onException(
+        self, loop: asyncio.AbstractEventLoop, context: dict[str, Any]
+    ) -> None:
+        e = context.get("exception")
+        if e:
+            exception(e)
+
+
 class ServerOptions(NamedTuple):
     host: str = "0.0.0.0"  # nosec: B104
     port: int = 8000
@@ -336,6 +353,13 @@ async def Serve(
         except RuntimeError:
             loop = asyncio.new_event_loop()
 
+        # Manage server state
+        state = ServerState()
+        # Registers handlers for signals and exception (so that we log them)
+        loop.add_signal_handler(SIGINT, lambda: state.stop())
+        loop.add_signal_handler(SIGTERM, lambda: state.stop())
+        loop.set_exception_handler(state.onException)
+
         info(
             "Extra AIO Server listening",
             icon="🚀",
@@ -344,7 +368,7 @@ async def Serve(
         )
 
         try:
-            while True:
+            while state.isRunning:
                 if options.condition and not options.condition():
                     break
                 try: