Handle SIGINT and SIGTERM (more) gracefully

argh/node.py

@@ -1,6 +1,7 @@
 # Copyright (c) 2025 Julian Müller (ChaoticByte)
 
 import asyncio
+import signal
 
 from hashlib import md5
 
@@ -58,9 +59,12 @@ class Node:
         identifier = addr.identifier()
         self._connections[identifier] = ws
         log("Accepted connection", identifier)
+        try:
             async for msg in ws:
                 self._relay(msg)
+        except:
             log("Lost connection", identifier)
+        finally:
             del self._connections[identifier]
 
     async def run(self):
@@ -71,9 +75,17 @@ class Node:
             receiver_tasks.append(t)
         # server loop
         async with serve(self._conn_from_node_or_client, self.listen_address.host, self.listen_address.port) as server:
-            await server.serve_forever()
-        # wait for receivers
-        await asyncio.gather(receiver_tasks)
+            loop = asyncio.get_running_loop()
+
+            def _close():
+                log("Cancelling receiver connections")
+                for r in receiver_tasks:
+                    r.cancel()
+                log("Closing server...")
+                server.close()
+
+            loop.add_signal_handler(signal.SIGTERM, _close)
+            await server.wait_closed()
 
 
 def node_from_yml(yml_data: str) -> Node:

node.py

@@ -3,10 +3,12 @@
 # Copyright (c) 2025 Julian Müller (ChaoticByte)
 
 import asyncio
+import signal
 
 from argparse import ArgumentParser
 from pathlib import Path
 
+from argh.log import log
 from argh.node import node_from_yml
 
 
@@ -18,4 +20,11 @@ if __name__ == "__main__":
 
     node = node_from_yml(args.config.read_text())
 
+    def _sigint(*args):
+        # don't need special sigint behaviour
+        # -> send SIGTERM
+        signal.raise_signal(signal.SIGTERM)
+    signal.signal(signal.SIGINT, _sigint)
+
     asyncio.run(node.run())
+    log("Bye")