fix: recover proxy automatically after suspend/resume

linux.py

@@ -13,14 +13,14 @@
 from PIL import Image, ImageTk
 
 import proxy.tg_ws_proxy as tg_ws_proxy
-
 from utils.tray_common import (
     APP_NAME, DEFAULT_CONFIG, FIRST_RUN_MARKER, LOG_FILE,
     acquire_lock, bootstrap, check_ipv6_warning, ctk_run_dialog,
     ensure_ctk_thread, ensure_dirs, load_config, load_icon, log,
     maybe_notify_update, quit_ctk, release_lock, restart_proxy,
     save_config, start_proxy, stop_proxy, tg_proxy_url,
 )
+from utils.resume_watchdog import ResumeWatchdog
 from ui.ctk_tray_ui import (
     install_tray_config_buttons, install_tray_config_form,
     populate_first_run_window, tray_settings_scroll_and_footer,
@@ -32,6 +32,7 @@
 )
 
 _tray_icon: Optional[object] = None
+_resume_watchdog: Optional[ResumeWatchdog] = None
 _config: dict = {}
 _exiting = False
 
@@ -65,6 +66,21 @@ def _ask_yes_no(text: str, title: str = "TG WS Proxy") -> bool:
     return bool(_msgbox("askyesno", text, title))
 
 
+def _start_resume_watchdog():
+    global _resume_watchdog
+    if _resume_watchdog is None:
+        _resume_watchdog = ResumeWatchdog(
+            lambda: restart_proxy(_config, _show_error),
+            log,
+        )
+    _resume_watchdog.start()
+
+
+def _stop_resume_watchdog():
+    if _resume_watchdog is not None:
+        _resume_watchdog.stop()
+
+
 def _apply_window_icon(root) -> None:
     icon_img = load_icon()
     if icon_img:
@@ -257,19 +273,25 @@ def run_tray() -> None:
                 time.sleep(1)
         except KeyboardInterrupt:
             stop_proxy()
+        finally:
+            _stop_resume_watchdog()
         return
 
-    start_proxy(_config, _show_error)
-    maybe_notify_update(_config, lambda: _exiting, _ask_yes_no)
-    _show_first_run()
-    check_ipv6_warning(_show_info)
+    _start_resume_watchdog()
+    try:
+        start_proxy(_config, _show_error)
+        maybe_notify_update(_config, lambda: _exiting, _ask_yes_no)
+        _show_first_run()
+        check_ipv6_warning(_show_info)
 
-    _tray_icon = pystray.Icon(APP_NAME, load_icon(), "TG WS Proxy", menu=_build_menu())
-    log.info("Tray icon running")
-    _tray_icon.run()
+        _tray_icon = pystray.Icon(APP_NAME, load_icon(), "TG WS Proxy", menu=_build_menu())
 
-    stop_proxy()
-    log.info("Tray app exited")
+        log.info("Tray icon running")
+        _tray_icon.run()
+    finally:
+        stop_proxy()
+        _stop_resume_watchdog()
+        log.info("Tray app exited")
 
 
 def main() -> None:

macos.py

@@ -26,6 +26,7 @@
 
 import proxy.tg_ws_proxy as tg_ws_proxy
 from proxy import __version__
+from utils.resume_watchdog import ResumeWatchdog
 
 from utils.tray_common import (
     APP_DIR, APP_NAME, DEFAULT_CONFIG, FIRST_RUN_MARKER, IPV6_WARN_MARKER,
@@ -38,6 +39,7 @@
 _proxy_thread: Optional[threading.Thread] = None
 _async_stop: Optional[object] = None
 _app: Optional[object] = None
+_resume_watchdog: Optional[ResumeWatchdog] = None
 _config: dict = {}
 _exiting: bool = False
 
@@ -178,6 +180,7 @@ def _start_proxy() -> None:
         return
     pc = tg_ws_proxy.proxy_config
     log.info("Starting proxy on %s:%d ...", pc.host, pc.port)
+    tg_ws_proxy.reset_runtime_state()
     _proxy_thread = threading.Thread(target=_run_proxy_thread, daemon=True, name="proxy")
     _proxy_thread.start()
 
@@ -200,6 +203,18 @@ def _restart_proxy() -> None:
     _start_proxy()
 
 
+def _start_resume_watchdog():
+    global _resume_watchdog
+    if _resume_watchdog is None:
+        _resume_watchdog = ResumeWatchdog(_restart_proxy, log)
+    _resume_watchdog.start()
+
+
+def _stop_resume_watchdog():
+    if _resume_watchdog is not None:
+        _resume_watchdog.stop()
+
+
 # menu callbacks
 
 
@@ -595,19 +610,24 @@ def run_menubar() -> None:
                 time.sleep(1)
         except KeyboardInterrupt:
             _stop_proxy()
+        finally:
+            _stop_resume_watchdog()
         return
 
-    _start_proxy()
-    _maybe_notify_update_async()
-    _show_first_run()
-    _check_ipv6_warning()
-
-    _app = TgWsProxyApp()
-    log.info("Menubar app running")
-    _app.run()
+    _start_resume_watchdog()
+    try:
+        _start_proxy()
+        _maybe_notify_update_async()
+        _show_first_run()
+        _check_ipv6_warning()
 
-    _stop_proxy()
-    log.info("Menubar app exited")
+        _app = TgWsProxyApp()
+        log.info("Menubar app running")
+        _app.run()
+    finally:
+        _stop_proxy()
+        _stop_resume_watchdog()
+        log.info("Menubar app exited")
 
 
 def main() -> None:

proxy/tg_ws_proxy.py

@@ -516,6 +516,16 @@ def summary(self) -> str:
 _stats = Stats()
 
 
+def reset_runtime_state() -> None:
+    global _dc_opt, _stats, _server_instance, _server_stop_event
+    _dc_opt = {}
+    _ws_blacklist.clear()
+    _dc_fail_until.clear()
+    _stats = Stats()
+    _server_instance = None
+    _server_stop_event = None
+
+
 class _WsPool:
     WS_POOL_MAX_AGE = 120.0
 
@@ -612,6 +622,17 @@ def reset(self):
         self._idle.clear()
         self._refilling.clear()
 
+    async def close_all(self):
+        tasks = []
+        for bucket in self._idle.values():
+            while bucket:
+                ws, _created = bucket.popleft()
+                tasks.append(asyncio.create_task(self._quiet_close(ws)))
+        self._idle.clear()
+        self._refilling.clear()
+        if tasks:
+            await asyncio.gather(*tasks, return_exceptions=True)
+
 _ws_pool = _WsPool()
 
 
@@ -1194,6 +1215,7 @@ async def log_stats():
             await log_stats_task
         except asyncio.CancelledError:
             pass
+        await _ws_pool.close_all()
     _server_instance = None
 
 

utils/resume_watchdog.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+import logging
+import threading
+import time
+from typing import Callable, Optional
+
+
+class ResumeWatchdog:
+    def __init__(
+        self,
+        on_resume: Callable[[], None],
+        logger: logging.Logger,
+        *,
+        interval: float = 15.0,
+        resume_gap: float = 45.0,
+        cooldown: float = 30.0,
+        name: str = "resume-watchdog",
+    ) -> None:
+        self._on_resume = on_resume
+        self._log = logger
+        self._interval = interval
+        self._resume_gap = resume_gap
+        self._cooldown = cooldown
+        self._name = name
+        self._stop_event = threading.Event()
+        self._thread: Optional[threading.Thread] = None
+        self._last_trigger = 0.0
+
+    def start(self) -> None:
+        if self._thread and self._thread.is_alive():
+            return
+        self._stop_event.clear()
+        self._thread = threading.Thread(
+            target=self._run,
+            daemon=True,
+            name=self._name,
+        )
+        self._thread.start()
+
+    def stop(self, timeout: float = 1.0) -> None:
+        self._stop_event.set()
+        if self._thread:
+            self._thread.join(timeout=timeout)
+        self._thread = None
+
+    def _run(self) -> None:
+        last_seen = time.time()
+        while not self._stop_event.wait(self._interval):
+            now = time.time()
+            gap = now - last_seen
+            last_seen = now
+
+            if gap < self._resume_gap:
+                continue
+
+            if now - self._last_trigger < self._cooldown:
+                continue
+
+            self._last_trigger = now
+            self._log.warning(
+                "Detected a %.1fs pause; restarting proxy to recover after resume",
+                gap,
+            )
+            try:
+                self._on_resume()
+            except Exception:
+                self._log.exception("Failed to recover proxy after resume")
+            finally:
+                last_seen = time.time()

utils/tray_common.py

@@ -287,6 +287,7 @@ def start_proxy(cfg: dict, on_error: Callable[[str], None]) -> None:
 
     pc = tg_ws_proxy.proxy_config
     log.info("Starting proxy on %s:%d ...", pc.host, pc.port)
+    tg_ws_proxy.reset_runtime_state()
     _proxy_thread = threading.Thread(
         target=_run_proxy_thread, args=(on_error,), daemon=True, name="proxy"
     )

windows.py

@@ -31,9 +31,8 @@
     Image = None
 
 import proxy.tg_ws_proxy as tg_ws_proxy
-
 from utils.win32_theme import (
-    is_windows_dark_theme, 
+    is_windows_dark_theme,
     apply_windows_dark_theme,
 )
 from utils.tray_common import (
@@ -43,6 +42,7 @@
     maybe_notify_update, quit_ctk, release_lock, restart_proxy,
     save_config, start_proxy, stop_proxy, tg_proxy_url,
 )
+from utils.resume_watchdog import ResumeWatchdog
 from ui.ctk_tray_ui import (
     install_tray_config_buttons, install_tray_config_form,
     populate_first_run_window, tray_settings_scroll_and_footer,
@@ -54,6 +54,7 @@
 )
 
 _tray_icon: Optional[object] = None
+_resume_watchdog: Optional[ResumeWatchdog] = None
 _config: dict = {}
 _exiting = False
 
@@ -126,6 +127,19 @@ def set_autostart_enabled(enabled: bool) -> None:
 
 # tray callbacks
 
+def _start_resume_watchdog():
+    global _resume_watchdog
+    if _resume_watchdog is None:
+        _resume_watchdog = ResumeWatchdog(
+            lambda: restart_proxy(_config, _show_error),
+            log,
+        )
+    _resume_watchdog.start()
+
+
+def _stop_resume_watchdog():
+    if _resume_watchdog is not None:
+        _resume_watchdog.stop()
 def _on_open_in_telegram(icon=None, item=None) -> None:
     url = tg_proxy_url(_config)
     log.info("Opening %s", url)
@@ -334,19 +348,25 @@ def run_tray() -> None:
                 time.sleep(1)
         except KeyboardInterrupt:
             stop_proxy()
+        finally:
+            _stop_resume_watchdog()
         return
 
-    start_proxy(_config, _show_error)
-    maybe_notify_update(_config, lambda: _exiting, _ask_yes_no)
-    _show_first_run()
-    check_ipv6_warning(_show_info)
+    _start_resume_watchdog()
+    try:
+        start_proxy(_config, _show_error)
+        maybe_notify_update(_config, lambda: _exiting, _ask_yes_no)
+        _show_first_run()
+        check_ipv6_warning(_show_info)
 
-    _tray_icon = pystray.Icon(APP_NAME, load_icon(), "TG WS Proxy", menu=_build_menu())
-    log.info("Tray icon running")
-    _tray_icon.run()
+        _tray_icon = pystray.Icon(APP_NAME, load_icon(), "TG WS Proxy", menu=_build_menu())
 
-    stop_proxy()
-    log.info("Tray app exited")
+        log.info("Tray icon running")
+        _tray_icon.run()
+    finally:
+        stop_proxy()
+        _stop_resume_watchdog()
+        log.info("Tray app exited")
 
 
 def main() -> None: