feat(client): add watchdog and disconnect flap handling
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.
This commit is contained in:
vvzvlad
+122
-27
@@ -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}")
|
||||
|
||||
Reference in New Issue
Block a user