From a0011dc886aa6e0bee19afb5abbf7d6357a4304a Mon Sep 17 00:00:00 2001 From: Vladimir Chebotarev Date: Thu, 9 Apr 2026 10:59:53 +0300 Subject: [PATCH] fix: recover proxy automatically after suspend/resume --- linux.py | 42 ++++++++++++++++++------ macos.py | 40 +++++++++++++++++------ proxy/tg_ws_proxy.py | 22 +++++++++++++ utils/resume_watchdog.py | 70 ++++++++++++++++++++++++++++++++++++++++ utils/tray_common.py | 1 + windows.py | 42 +++++++++++++++++------- 6 files changed, 186 insertions(+), 31 deletions(-) create mode 100644 utils/resume_watchdog.py diff --git a/linux.py b/linux.py index 39cf6c3d..22ba4d9a 100644 --- a/linux.py +++ b/linux.py @@ -13,7 +13,6 @@ 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, @@ -21,6 +20,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, @@ -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: diff --git a/macos.py b/macos.py index 29fbe6c5..cbc35b44 100644 --- a/macos.py +++ b/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: diff --git a/proxy/tg_ws_proxy.py b/proxy/tg_ws_proxy.py index 63bfcb10..090d3d2f 100644 --- a/proxy/tg_ws_proxy.py +++ b/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 diff --git a/utils/resume_watchdog.py b/utils/resume_watchdog.py new file mode 100644 index 00000000..12de1de9 --- /dev/null +++ b/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() diff --git a/utils/tray_common.py b/utils/tray_common.py index f888b5d3..76ae2535 100644 --- a/utils/tray_common.py +++ b/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" ) diff --git a/windows.py b/windows.py index d810b690..6a6e5a4b 100644 --- a/windows.py +++ b/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: