From 79b127d406bce9bb935dd0d4fdc5dea9435edbd6 Mon Sep 17 00:00:00 2001 From: vvzvlad Date: Thu, 4 Jun 2026 19:02:09 +0300 Subject: [PATCH] feat(client): add watchdog and disconnect flap handling MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduce an active watchdog that probes the Telegram client to detect zombie sessions and restart them in‑process. Add configurable disconnect flap detection with a sliding window to trigger restarts after repeated disconnects. New environment variables and config entries are added, and the Kurigram dependency is now version‑pinned. --- config.py | 22 +++++++ dockercompose.yml | 10 ++- requirements.txt | 2 +- telegram_client.py | 149 +++++++++++++++++++++++++++++++++++-------- tests/mock_config.py | 9 ++- 5 files changed, 162 insertions(+), 30 deletions(-) diff --git a/config.py b/config.py index 1b85603..9a6b807 100644 --- a/config.py +++ b/config.py @@ -82,6 +82,21 @@ def get_settings() -> dict[str, Any]: print(f"API_PORT must be a valid integer, got: {os.getenv('API_PORT')!r}", flush=True) sys.exit(1) + # Local helper to parse int env vars with a default and exit on a bad value + def _parse_int_env(name: str, default: int, minimum: int = 1) -> int: + raw = os.getenv(name) + if raw is None or raw.strip() == "": + return default + try: + value = int(raw) + except ValueError: + print(f"{name} must be a valid integer, got: {raw!r}", flush=True) + sys.exit(1) + if value < minimum: + print(f"{name} must be >= {minimum}, got: {value}", flush=True) + sys.exit(1) + return value + return { "tg_api_id": tg_api_id_int, "tg_api_hash": tg_api_hash, @@ -97,4 +112,11 @@ def get_settings() -> dict[str, Any]: "show_bridge_link": os.getenv("SHOW_BRIDGE_LINK", "False").strip() in ["True", "true"], "show_post_flags": os.getenv("SHOW_POST_FLAGS", "False").strip() in ["True", "true"], "proxy": proxy, + "tg_watchdog_enabled": os.getenv("TG_WATCHDOG_ENABLED", "true").strip().lower() not in ["false", "0", "no", "off", "disable", "disabled"], + "tg_watchdog_interval": _parse_int_env("TG_WATCHDOG_INTERVAL", 60), + "tg_watchdog_timeout": _parse_int_env("TG_WATCHDOG_TIMEOUT", 10), + "tg_watchdog_failures": _parse_int_env("TG_WATCHDOG_FAILURES", 3), + "tg_watchdog_restart_timeout": _parse_int_env("TG_WATCHDOG_RESTART_TIMEOUT", 90), + "tg_disconnect_flap_limit": _parse_int_env("TG_DISCONNECT_FLAP_LIMIT", 3), + "tg_disconnect_flap_window": _parse_int_env("TG_DISCONNECT_FLAP_WINDOW", 120), } diff --git a/dockercompose.yml b/dockercompose.yml index 81fa6eb..3813236 100644 --- a/dockercompose.yml +++ b/dockercompose.yml @@ -8,10 +8,18 @@ services: environment: TG_API_ID: XXX TG_API_HASH: XXX - # TG_PROXY_HOST: 127.0.0.1 # MTProto proxy host, SOCKS5 interface (optional) + # TG_PROXY_HOST: 10.0.0.1 # MTProto proxy host (SOCKS5). Prefer a literal IP over a hostname: + # # a hostname is re-resolved on every reconnect, so flaky DNS during a + # # network blip can wedge Pyrogram's reconnect. # TG_PROXY_PORT: 1080 # MTProto proxy SOCKS5 port (default: 1080) # TG_PROXY_USERNAME: XXX # SOCKS5 username (optional) # TG_PROXY_PASSWORD: XXX # SOCKS5 password (optional) + # TG_WATCHDOG_ENABLED: "True" # Active liveness watchdog: detects 'zombie' sessions and restarts in-process (default: True) + # TG_WATCHDOG_INTERVAL: 60 # Seconds between liveness probes (default: 60) + # TG_WATCHDOG_TIMEOUT: 10 # Seconds to wait for each get_me probe (default: 10) + # TG_WATCHDOG_FAILURES: 3 # Consecutive failed probes before restart (default: 3) + # TG_DISCONNECT_FLAP_LIMIT: 3 # Disconnect events within the flap window before an in-process restart (default: 3) + # TG_DISCONNECT_FLAP_WINDOW: 120 # Flap detection window in seconds (default: 120) PYROGRAM_BRIDGE_URL: https://pgbridge.example.com API_PORT: 80 TOKEN: ХХХ diff --git a/requirements.txt b/requirements.txt index ac8ece5..2c19d8d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,7 @@ fastapi==0.115.8 uvicorn==0.34.0 python-multipart==0.0.20 -Kurigram +Kurigram==2.2.22 #git+https://github.com/KurimuzonAkuma/pyrogram.git@dev#egg=Kurigram TgCrypto uvloop diff --git a/telegram_client.py b/telegram_client.py index 41a38d7..99f7b52 100644 --- a/telegram_client.py +++ b/telegram_client.py @@ -34,9 +34,17 @@ class TelegramClient: workdir=settings["session_path"], proxy=settings["proxy"], # MTProto proxy config, None if not set ) - self.disconnect_count = 0 - self.max_disconnects = 3 + self.max_disconnects = settings["tg_disconnect_flap_limit"] # Max disconnects within the flap window before restart self._shutting_down = False # Guard to prevent re-triggering restart during shutdown + self._restarting = False # Guard: an intentional in-process restart is in progress + self._disconnect_times = [] # Monotonic timestamps of recent disconnects (sliding window) + self.disconnect_window = settings["tg_disconnect_flap_window"] # Seconds; window for flap detection + self._watchdog_task = None + self.watchdog_enabled = settings["tg_watchdog_enabled"] + self.watchdog_interval = settings["tg_watchdog_interval"] + self.watchdog_timeout = settings["tg_watchdog_timeout"] + self.watchdog_failures = settings["tg_watchdog_failures"] + self.watchdog_restart_timeout = settings["tg_watchdog_restart_timeout"] self._setup_connection_handlers() def _ensure_session_directory(self): @@ -53,41 +61,113 @@ class TelegramClient: logger.info("connection_handlers: connection handlers set up") async def _on_disconnect(self, _client, _session=None): - """Handles disconnection from Telegram servers""" - # Ignore disconnects that are caused by the intentional shutdown sequence - if self._shutting_down: - logger.debug("connection_handler: ignoring disconnect event during shutdown") - return - self.disconnect_count += 1 - logger.warning(f"connection_handler: connection lost (#{self.disconnect_count})") + """Handles disconnection events from Telegram servers (best-effort safety net). - # Wait briefly to allow Pyrogram's internal reconnect to complete - await asyncio.sleep(10) - - # Re-check shutdown flag — it may have been set while we were sleeping - # (e.g. another concurrent _on_disconnect already triggered _restart_app) - if self._shutting_down: - logger.debug("connection_handler: shutdown started during sleep, suppressing further action") + NOTE: this handler only fires when Pyrogram calls session.stop(). It does NOT fire in + the 'zombie session' case where the recv loop dies silently — that case is handled by + the active watchdog (_watchdog_loop). Keep both paths. + """ + # Ignore disconnects caused by intentional shutdown or by our own in-process restart + # (client.restart() -> client.stop() -> session.stop() re-dispatches this handler). + if self._shutting_down or self._restarting: + logger.debug("connection_handler: ignoring disconnect event (shutdown/restart in progress)") return - # If the client reconnected successfully, reset the counter and return - if self.client.is_connected: - logger.info(f"connection_handler: reconnected successfully, resetting disconnect counter") - self.disconnect_count = 0 - return + now = time.monotonic() + # Sliding window: drop disconnects older than the window, then record this one. + self._disconnect_times = [t for t in self._disconnect_times if now - t <= self.disconnect_window] + self._disconnect_times.append(now) + count = len(self._disconnect_times) + logger.warning(f"connection_handler: connection lost ({count} within {self.disconnect_window}s window)") - if self.disconnect_count >= self.max_disconnects: - logger.critical(f"connection_handler: reached disconnect limit ({self.max_disconnects})") + if count >= self.max_disconnects: + logger.critical(f"connection_handler: disconnect flap limit reached ({self.max_disconnects} within {self.disconnect_window}s), restarting client") + await self._restart_client() + + def _start_watchdog(self): + """Starts the active liveness watchdog task (idempotent).""" + if not self.watchdog_enabled: + logger.info("watchdog: disabled via configuration") + return + if self._watchdog_task is not None and not self._watchdog_task.done(): + return + self._watchdog_task = asyncio.create_task(self._watchdog_loop()) + + async def _watchdog_loop(self): + """Active liveness probe for the 'zombie session' state. + + The disconnect-only recovery never triggers when Pyrogram's recv loop dies silently + (is_connected stays True, no Disconnect event). This loop periodically issues a real + lightweight API call (get_me) bounded by a short timeout; after N consecutive failures + it forces an in-process restart. + """ + consecutive_failures = 0 + logger.info(f"watchdog: started (interval={self.watchdog_interval}s, timeout={self.watchdog_timeout}s, failures={self.watchdog_failures})") + try: + while True: + await asyncio.sleep(self.watchdog_interval) + if self._shutting_down: + break + if self._restarting: + # A restart is already underway; skip this probe cycle. + continue + try: + await asyncio.wait_for(self.client.get_me(), timeout=self.watchdog_timeout) + if consecutive_failures: + logger.info(f"watchdog: liveness restored after {consecutive_failures} failed probe(s)") + consecutive_failures = 0 + except asyncio.CancelledError: + raise + except Exception as e: + consecutive_failures += 1 + logger.warning(f"watchdog: liveness probe failed ({consecutive_failures}/{self.watchdog_failures}): {type(e).__name__}: {e}") + if consecutive_failures >= self.watchdog_failures: + consecutive_failures = 0 + await self._restart_client() + except asyncio.CancelledError: + logger.info("watchdog: stopped") + raise + except Exception as e: + logger.critical(f"watchdog: loop crashed unexpectedly ({type(e).__name__}: {e}); liveness protection is now DISABLED until next start") + + async def _restart_client(self): + """Recover the client without killing the process when possible. + + Performs an in-process restart (rebuilds session + recv loop), bounded by a timeout so a + dead network layer cannot make it hang forever. Falls back to a full process restart + (SIGTERM) if the in-process restart fails or times out. + """ + if self._restarting or self._shutting_down: + return + self._restarting = True + try: + if self.client.is_connected: + logger.critical("recovery: restarting Telegram client in-process (client.restart)") + await asyncio.wait_for(self.client.restart(), timeout=self.watchdog_restart_timeout) + else: + logger.critical("recovery: client not connected, starting it in-process (client.start)") + await asyncio.wait_for(self.client.start(), timeout=self.watchdog_restart_timeout) + logger.info("recovery: in-process client recovery completed successfully") + self._disconnect_times.clear() + # Re-arm the watchdog in case it had previously crashed (self-healing). + self._start_watchdog() + except asyncio.CancelledError: + raise + except Exception as e: + logger.error(f"recovery: in-process restart failed ({type(e).__name__}: {e}); falling back to process restart") self._restart_app() + finally: + self._restarting = False async def start(self): try: if not self.client.is_connected: await self.client.start() logger.info("Telegram client connected successfully") - # Reset disconnect counter on successful connection - self.disconnect_count = 0 + # Reset flap history on a fresh successful connection + self._disconnect_times.clear() logger.info("connection_handler: connection established") + self._start_watchdog() except Exception as e: logger.error(f"Failed to start Telegram client: {str(e)}") raise @@ -136,6 +216,21 @@ class TelegramClient: raise async def stop(self): + # Suppress disconnect handling during intentional shutdown + # (client.stop() dispatches the DisconnectHandler). + self._shutting_down = True + if self._watchdog_task is not None: + self._watchdog_task.cancel() + try: + await self._watchdog_task + except asyncio.CancelledError: + pass + self._watchdog_task = None if self.client.is_connected: - await self.client.stop() - logger.info("Telegram client disconnected") + try: + await self.client.stop() + logger.info("Telegram client disconnected") + except Exception as e: + # During shutdown the client may be in a half-restarted state + # (e.g. a watchdog restart was cancelled mid-flight); ignore stop errors. + logger.warning(f"Telegram client stop during shutdown raised {type(e).__name__}: {e}") diff --git a/tests/mock_config.py b/tests/mock_config.py index 2552053..b166ed8 100644 --- a/tests/mock_config.py +++ b/tests/mock_config.py @@ -20,4 +20,11 @@ def get_settings(): "show_post_flags": True, "proxy": None, "trusted_proxies": [], - } \ No newline at end of file + "tg_watchdog_enabled": True, + "tg_watchdog_interval": 60, + "tg_watchdog_timeout": 10, + "tg_watchdog_failures": 3, + "tg_watchdog_restart_timeout": 90, + "tg_disconnect_flap_limit": 3, + "tg_disconnect_flap_window": 120, + }