vvzvlad/pyrogram-bridge
1
0
Fork 0
Code Issues 6 Pull Requests Actions 8 Packages Projects Releases Wiki Activity

feat(client): add watchdog and disconnect flap handling

Browse Source 
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
2026-06-04 19:02:09 +03:00
parent 577093b9fa
commit 79b127d406
5 changed files with 162 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
config.py
+22
View File
@@ -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),
}
dockercompose.yml
+9 -1
View File
@@ -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: ХХХ
requirements.txt
+1 -1
View File
@@ -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
telegram_client.py
+122 -27
View File
@@ -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}")
tests/mock_config.py
+8 -1
View File
@@ -20,4 +20,11 @@ def get_settings():
"show_post_flags": True,
"proxy": None,
"trusted_proxies": [],
}
"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,
}