diff --git a/server/udp_server.py b/server/udp_server.py index a397997..b69a1ff 100644 --- a/server/udp_server.py +++ b/server/udp_server.py @@ -52,7 +52,9 @@ def connection_made(self, transport: asyncio.DatagramTransport): def datagram_received(self, data: bytes, addr: tuple): """Handle incoming UDP datagram.""" # Schedule processing as a task to avoid blocking the protocol - asyncio.create_task(self.server.handle_packet(data, addr)) + task = asyncio.create_task(self.server.handle_packet(data, addr)) + self.server.background_tasks.add(task) + task.add_done_callback(lambda t: self.server.background_tasks.discard(t)) def error_received(self, exc: Exception): print(f"UDP Error: {exc}") @@ -98,6 +100,7 @@ def __init__( self.running = False self.transport: Optional[asyncio.DatagramTransport] = None self.protocol: Optional[UDPInputProtocol] = None + self.background_tasks: Set[asyncio.Task] = set() # Track background tasks for cleanup # Player tracking self.player_addresses: Dict[str, tuple] = {} # player_id -> (host, port) @@ -136,8 +139,19 @@ async def stop(self): """Stop the UDP server and cleanup.""" self.running = False + # Cancel all background tasks + for task in self.background_tasks: + task.cancel() + + # Wait for all background tasks to complete with a timeout + if self.background_tasks: + await asyncio.gather(*self.background_tasks, return_exceptions=True) + + # Close the transport if self.transport: self.transport.close() + # Give the transport a moment to flush and close cleanly + await asyncio.sleep(0.1) self.transport = None # Log final metrics