diff --git a/packages/prime-tunnel/src/prime_tunnel/tunnel.py b/packages/prime-tunnel/src/prime_tunnel/tunnel.py index 1535800f..27823984 100644 --- a/packages/prime-tunnel/src/prime_tunnel/tunnel.py +++ b/packages/prime-tunnel/src/prime_tunnel/tunnel.py @@ -1,4 +1,5 @@ import asyncio +import atexit import fcntl import os import re @@ -29,6 +30,22 @@ ) _ANSI_RE = re.compile(r"\x1b\[[0-9;]*m") +# Started-but-not-stopped tunnels, so an atexit backstop can delete their backend +# registrations if the caller never stops them (forgotten cleanup, unhandled +# exception) +_active_tunnels: "set[Tunnel]" = set() + + +def _stop_active_tunnels() -> None: + for tunnel in list(_active_tunnels): + try: + tunnel.sync_stop() + except Exception: + pass + + +atexit.register(_stop_active_tunnels) + def _parse_frpc_error( output_lines: list[str], @@ -205,6 +222,8 @@ async def start(self) -> str: self._started = True + _active_tunnels.add(self) + return self.url async def stop(self) -> None: @@ -220,6 +239,8 @@ def sync_stop(self) -> None: if not self._started: return + _active_tunnels.discard(self) + if self._process is not None: try: self._process.terminate() @@ -258,6 +279,8 @@ def sync_stop(self) -> None: async def _cleanup(self) -> None: """Clean up tunnel resources.""" + _active_tunnels.discard(self) + # Stop frpc process (this will cause drain threads to exit via EOF) if self._process is not None: try: