From 1bd4d9eac7390537a4b891255112d92ece38a79c Mon Sep 17 00:00:00 2001
From: Vladyslav Doloman <git@vaku.org.ua>
Date: Wed, 8 Oct 2025 00:27:10 +0300
Subject: [PATCH] Fix: clean up run.py newlines and indentation; add optional
 RUN_SECONDS timeout helper for server tasks

---
 run.py | 30 ++++++++++++++++++++++++++----
 1 file changed, 26 insertions(+), 4 deletions(-)

diff --git a/run.py b/run.py
index b9ba799..0435c8c 100644
--- a/run.py
+++ b/run.py
@@ -1,4 +1,4 @@
-import asyncio
+import asyncio
 import os
 import logging
 
@@ -10,7 +10,8 @@ async def run_in_memory():
     from server.transport import InMemoryTransport
     cfg = ServerConfig()
     server = GameServer(transport=InMemoryTransport(lambda d, p: server.on_datagram(d, p)), config=cfg)
-    await asyncio.gather(server.transport.run(), server.tick_loop())
+    tasks = [asyncio.create_task(server.transport.run()), asyncio.create_task(server.tick_loop())]
+    await _run_tasks_with_optional_timeout(tasks)
 
 
 async def run_quic():
@@ -21,7 +22,8 @@ async def run_quic():
     cert = os.environ["QUIC_CERT"]
     key = os.environ["QUIC_KEY"]
     server = GameServer(transport=QuicWebTransportServer(host, port, cert, key, lambda d, p: server.on_datagram(d, p)), config=cfg)
-    await asyncio.gather(server.transport.run(), server.tick_loop())
+    tasks = [asyncio.create_task(server.transport.run()), asyncio.create_task(server.tick_loop())]
+    await _run_tasks_with_optional_timeout(tasks)
 
 
 async def run_webtransport():
@@ -44,7 +46,8 @@ async def run_webtransport():
     server = GameServer(transport=WebTransportServer(host, port, cert, key, lambda d, p: server.on_datagram(d, p)), config=cfg)
     print(f"WebTransport server: https://{host}:{port}/ (HTTP/3)")
     try:
-        await asyncio.gather(server.transport.run(), server.tick_loop())
+        tasks = [asyncio.create_task(server.transport.run()), asyncio.create_task(server.tick_loop())]
+        await _run_tasks_with_optional_timeout(tasks)
     finally:
         if httpd is not None:
             httpd.shutdown()
@@ -67,3 +70,22 @@ if __name__ == "__main__":
             asyncio.run(run_in_memory())
     except KeyboardInterrupt:
         pass
+
+
+
+
+
+
+async def _run_tasks_with_optional_timeout(tasks):
+    """Await tasks, optionally honoring RUN_SECONDS env var to cancel after a timeout."""
+    timeout_s = os.environ.get("RUN_SECONDS")
+    if not timeout_s:
+        await asyncio.gather(*tasks)
+        return
+    try:
+        await asyncio.wait_for(asyncio.gather(*tasks), timeout=float(timeout_s))
+    except asyncio.TimeoutError:
+        logging.info("Timeout reached (RUN_SECONDS=%s); stopping server tasks...", timeout_s)
+        for t in tasks:
+            t.cancel()
+        await asyncio.gather(*tasks, return_exceptions=True)