From 4daf611f05f2c6258da8849ef8936c748dd42f1d Mon Sep 17 00:00:00 2001 From: claude code agent Date: Sun, 5 Jul 2026 06:51:08 +0300 Subject: [PATCH] =?UTF-8?q?fix(stability):=20stage=201=20=E2=80=94=20timeo?= =?UTF-8?q?uts=20on=20all=20TG=20RPC=20+=20resilient=20download=20worker?= =?UTF-8?q?=20+=20supervision?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Eliminates the app-wide hangs where one stuck Telegram RPC holds the single tg_rpc gate permit forever (freezing every RSS/HTML request) and where the background download worker dies silently. 1.1 Every RPC under the global gate is now bounded by asyncio.wait_for (config tg_rpc_timeout, env TG_RPC_TIMEOUT, default 60). The wait_for wraps ONLY the RPC body, never the gate acquire, so queue backpressure stays legitimate; the paginated get_chat_history is collected in an inner coroutine so the async-for can be bounded. A timeout propagates out of the gate so its permit is released (no leak). 1.2 The other live RPC paths get the same treatment: _reply_enrichment's get_messages now runs under the gate + timeout; PostParser.get_post is bounded at 30s (and a stray print removed); /health's get_me at 10s. 1.3 background_download_worker: queue.get() moved out of the try so task_done() in finally balances exactly one get (no ValueError that killed the worker); FloodWait is caught BEFORE the generic Exception (sleeps, does not flood); download_new_files uses put_nowait so a full queue no longer blocks the sweeper forever. 1.4 Both background tasks run under a _supervised() wrapper: a crash or an unexpected return is logged CRITICAL and restarted (restarts rate-limited to once per 60s so a hard-failing task can't spin), while CancelledError is propagated to the child for a clean shutdown. Tests (tests/test_stage1_hangs.py): gate timeout releases the permit + a second call succeeds; gate cancel mid-spacing releases the permit; the worker survives Exception and FloodWait with balanced task_done; and _supervised restarts a crashing task (rate-limited) and an unexpected return, and propagates cancellation to its child. 180 passed (174 baseline + 6). Co-Authored-By: Claude Opus 4.8 (1M context) --- api_server.py | 72 ++++++++++--- config.py | 3 + post_parser.py | 7 +- rss_generator.py | 9 +- tests/mock_config.py | 5 + tests/test_stage1_hangs.py | 214 +++++++++++++++++++++++++++++++++++++ tg_cache.py | 18 +++- 7 files changed, 305 insertions(+), 23 deletions(-) create mode 100644 tests/test_stage1_hangs.py diff --git a/api_server.py b/api_server.py index ea8ce36..ac4aa22 100644 --- a/api_server.py +++ b/api_server.py @@ -72,6 +72,37 @@ HTTP_DOWNLOAD_SEMAPHORE = asyncio.Semaphore(3) # semaphore for live HTTP media BACKGROUND_DOWNLOAD_SEMAPHORE = asyncio.Semaphore(2) # semaphore for background cache worker download_queue = asyncio.Queue(maxsize=100) +async def _supervised(factory, name: str, min_restart_interval: float = 60.0): + """Run factory() forever, restarting it if it dies with a non-cancellation error. + + A background loop that returns or raises (anything except CancelledError) is logged + at CRITICAL and restarted, but successive (re)starts are spaced at least + min_restart_interval seconds apart so a hard-failing task can't spin the event loop. + CancelledError (shutdown) is propagated to the child and stops supervision. + """ + while True: + start = time.monotonic() + task = asyncio.create_task(factory(), name=name) + try: + await task + # A supervised background loop is not expected to return on its own. + logger.critical(f"supervised_task_exited: {name} returned unexpectedly; restarting") + except asyncio.CancelledError: + # Shutdown or external cancel: propagate to the child, then stop supervising. + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + raise + except Exception as e: + logger.critical(f"supervised_task_crashed: {name} died with {e!r}; restarting", exc_info=True) + # Rate-limit restarts: keep successive starts at least min_restart_interval apart. + elapsed = time.monotonic() - start + if elapsed < min_restart_interval: + await asyncio.sleep(min_restart_interval - elapsed) + + @asynccontextmanager async def lifespan(_: FastAPI): setup_logging(Config["log_level"]) @@ -84,8 +115,10 @@ async def lifespan(_: FastAPI): await asyncio.to_thread(init_db_sync, DB_PATH) await client.start() - background_task = asyncio.create_task(cache_media_files()) # Start background task - worker_task = asyncio.create_task(background_download_worker()) # Start download worker + # Supervise the background tasks: if either dies (not via cancellation) it is logged + # CRITICAL and restarted, so a crash can no longer silently stop cache sweeping or downloads. + background_task = asyncio.create_task(_supervised(cache_media_files, "cache_media_files")) + worker_task = asyncio.create_task(_supervised(background_download_worker, "background_download_worker")) yield background_task.cancel() # Cleanup worker_task.cancel() @@ -602,7 +635,10 @@ async def download_new_files(media_files: list, cache_dir: str) -> None: cache_path = os.path.join(post_dir, file_unique_id) if not os.path.exists(cache_path): try: - await download_queue.put((channel, post_id, file_unique_id)) + # put_nowait so a full queue raises QueueFull instead of blocking + # cache_media_files (and thus the sweeper) forever. `await put()` + # never raises QueueFull, which made the except below dead code. + download_queue.put_nowait((channel, post_id, file_unique_id)) files_queued += 1 logger.debug(f"Queued for background download: {channel}/{post_id}/{file_unique_id}") except asyncio.QueueFull: @@ -620,19 +656,22 @@ async def download_new_files(media_files: list, cache_dir: str) -> None: async def background_download_worker(): """Worker that processes downloads from queue""" while True: + # get() is OUTSIDE the try so task_done() in finally always balances exactly one + # successful get(). Cancellation propagates cleanly here (nothing to unbalance). + item = await download_queue.get() + channel, post_id, file_unique_id = item + logger.info(f"Background download: {channel}/{post_id}/{file_unique_id}") try: - channel, post_id, file_unique_id = await download_queue.get() - logger.info(f"Background download: {channel}/{post_id}/{file_unique_id}") - - try: - async with BACKGROUND_DOWNLOAD_SEMAPHORE: # limit concurrent background downloads - await download_media_file(channel, post_id, file_unique_id) - await asyncio.sleep(2) - except Exception as e: - logger.error(f"Background download error for {channel}/{post_id}/{file_unique_id}: {e}") - + async with BACKGROUND_DOWNLOAD_SEMAPHORE: # limit concurrent background downloads + await download_media_file(channel, post_id, file_unique_id) + await asyncio.sleep(2) + except errors.FloodWait as e: + # Must be caught BEFORE the generic Exception (FloodWait subclasses RPCError), + # otherwise the worker would hammer Telegram while under a flood wait. + logger.warning(f"bg_download_floodwait: {channel}/{post_id}/{file_unique_id} sleeping {e.value}s") + await asyncio.sleep(min(int(e.value) + 5, 900)) except Exception as e: - logger.error(f"Background download worker error: {e}") + logger.error(f"Background download error for {channel}/{post_id}/{file_unique_id}: {e}") finally: download_queue.task_done() @@ -903,8 +942,9 @@ async def health_check(request: Request, token: str | None = None) -> Response: logger.info(f"Local request, skipping token check for health check.") try: - me = await client.client.get_me() - + # Bound the Telegram RPC so a hung get_me cannot hang the healthcheck. + me = await asyncio.wait_for(client.client.get_me(), timeout=10) + # Offload heavy filesystem scanning to threadpool cache_stats = await asyncio.to_thread(calculate_cache_stats) diff --git a/config.py b/config.py index 8ef91c6..03e47bd 100644 --- a/config.py +++ b/config.py @@ -112,6 +112,9 @@ 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, + # Hard cap (seconds) on any single live Telegram RPC held under the global RPC gate, + # so a hung MTProto call can never pin the gate (and the whole app) indefinitely. + "tg_rpc_timeout": _parse_int_env("TG_RPC_TIMEOUT", 60), "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), diff --git a/post_parser.py b/post_parser.py index 1e054e3..3e75f49 100644 --- a/post_parser.py +++ b/post_parser.py @@ -92,10 +92,13 @@ class PostParser: post_id: int, output_type: str = 'json', debug: bool = False) -> Union[str, Dict[Any, Any], None]: - print(f"Getting post {channel}, {post_id}") try: prepared_channel_id: Union[str, int] = self.channel_name_prepare(channel) - message = await self.client.get_messages(prepared_channel_id, post_id) + # Bound the single-post fetch so a hung RPC cannot block the request forever. + message = await asyncio.wait_for( + self.client.get_messages(prepared_channel_id, post_id), + timeout=30, + ) if Config["debug"]: print(message) diff --git a/rss_generator.py b/rss_generator.py index 0545b0b..b3cb822 100644 --- a/rss_generator.py +++ b/rss_generator.py @@ -22,6 +22,7 @@ from pyrogram import errors, Client from pyrogram.types import Message from post_parser import PostParser from config import get_settings +from tg_throttle import tg_rpc from bleach.css_sanitizer import CSSSanitizer from bleach import clean as HTMLSanitizer @@ -502,7 +503,13 @@ async def _reply_enrichment(client: Client, messages: list[Message]) -> list[Mes for chat_id, chat_msgs in chat_messages.items(): ids_to_fetch = [m.id for m in chat_msgs] try: - fetched = await client.get_messages(chat_id, ids_to_fetch) + # Throttle under the global RPC gate and bound the call so a hung + # get_messages cannot pin the gate (wait_for wraps the RPC, not the gate entry). + async with tg_rpc(): + fetched = await asyncio.wait_for( + client.get_messages(chat_id, ids_to_fetch), + timeout=Config["tg_rpc_timeout"], + ) # get_messages may return a single Message or a list if not isinstance(fetched, list): fetched = [fetched] diff --git a/tests/mock_config.py b/tests/mock_config.py index 2d7e1eb..bd22f42 100644 --- a/tests/mock_config.py +++ b/tests/mock_config.py @@ -1,6 +1,10 @@ #!/usr/bin/env python3 # -*- coding: utf-8 -*- +def setup_logging(level_name: str = "INFO") -> None: + """No-op logging setup for tests (mirrors config.setup_logging signature).""" + return None + def get_settings(): """ Mock config for testing without requiring TG_API_ID and TG_API_HASH @@ -20,6 +24,7 @@ def get_settings(): "show_post_flags": True, "proxy": None, "trusted_proxies": [], + "tg_rpc_timeout": 60, "tg_watchdog_enabled": True, "tg_watchdog_interval": 60, "tg_watchdog_timeout": 10, diff --git a/tests/test_stage1_hangs.py b/tests/test_stage1_hangs.py new file mode 100644 index 0000000..b4e43b5 --- /dev/null +++ b/tests/test_stage1_hangs.py @@ -0,0 +1,214 @@ +# flake8: noqa +# pylint: disable=protected-access, missing-function-docstring, missing-class-docstring +# pylint: disable=redefined-outer-name, logging-fstring-interpolation, line-too-long +# pylance: disable=reportMissingImports, reportMissingModuleSource +""" +Stage 1 (anti-hang) regression tests: +- RPC gate timeout: a hung RPC times out and the gate permit is NOT leaked; a + subsequent call still succeeds. +- Background worker: a download that raises Exception / FloodWait does not kill the + worker, and task_done stays balanced so queue.join() completes. +- Gate cancellation during the spacing wait does not lose the permit. +""" +import os +import sys +import time +import asyncio +from types import SimpleNamespace + +import pytest + +# Add project root to sys.path and mock the config module (same pattern as the other tests). +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) +sys.modules['config'] = __import__('tests.mock_config', fromlist=['get_settings']) + +import tg_throttle +import tg_cache +from pyrogram import errors + + +# --------------------------------------------------------------------------- # +# 1.1 — RPC gate timeout releases the permit; second call succeeds. +# --------------------------------------------------------------------------- # +@pytest.mark.asyncio +async def test_gate_timeout_releases_permit(monkeypatch): + # Short timeout so the hung RPC fails fast, and no min-interval spacing to slow the test. + monkeypatch.setitem(tg_cache.Config, "tg_rpc_timeout", 0.1) + monkeypatch.setattr(tg_throttle, "_MIN_INTERVAL", 0.0) + # Bypass the on-disk chat cache so we always hit the live RPC path. + monkeypatch.setattr(tg_cache, "_get_chat_from_cache", lambda *a, **k: None) + monkeypatch.setattr(tg_cache, "_save_chat_to_cache", lambda *a, **k: None) + + permits_before = tg_throttle._sem._value + never = asyncio.Event() # never set -> RPC hangs forever + + class HungClient: + async def get_chat(self, channel_id): + await never.wait() + + # First call: the RPC hangs and must time out (not hang the whole app). + with pytest.raises(asyncio.TimeoutError): + await tg_cache.cached_get_chat(HungClient(), "hung_channel") + + # The permit was released on timeout (gate fully available again) — no leak. + assert tg_throttle._sem._value == permits_before + + # A second call still goes through the gate and succeeds. + class OkClient: + async def get_chat(self, channel_id): + return SimpleNamespace(id=42, title="ok", username="okchan") + + res = await tg_cache.cached_get_chat(OkClient(), "ok_channel") + assert res.id == 42 + # Permit released again after the successful call. + assert tg_throttle._sem._value == permits_before + + +# --------------------------------------------------------------------------- # +# 1.4 — Gate cancellation during the spacing wait does not lose the permit. +# --------------------------------------------------------------------------- # +@pytest.mark.asyncio +async def test_gate_cancel_during_spacing_releases_permit(monkeypatch): + # Force a spacing wait long enough to cancel inside it. + monkeypatch.setattr(tg_throttle, "_MIN_INTERVAL", 1.0) + monkeypatch.setattr(tg_throttle, "_last_start", time.monotonic()) + + permits_before = tg_throttle._sem._value + + async def enter_gate(): + async with tg_throttle.tg_rpc(): + pass + + task = asyncio.create_task(enter_gate()) + # Let it acquire the semaphore and start sleeping for the spacing interval. + await asyncio.sleep(0.05) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + # Cancelled mid-spacing: the acquired permit must be returned. + assert tg_throttle._sem._value == permits_before + + +# --------------------------------------------------------------------------- # +# 1.3 — Background worker survives download errors; task_done stays balanced. +# --------------------------------------------------------------------------- # +@pytest.mark.asyncio +async def test_worker_survives_errors_and_balances_task_done(monkeypatch): + import api_server + + # Speed up the worker's post-download / flood-wait sleeps. + async def _fast_sleep(*_a, **_k): + return + monkeypatch.setattr(api_server.asyncio, "sleep", _fast_sleep) + + # Fresh queue so we don't interfere with (or depend on) module state. + q = asyncio.Queue(maxsize=100) + monkeypatch.setattr(api_server, "download_queue", q) + + processed = [] + + async def fake_download(channel, post_id, file_unique_id): + processed.append(file_unique_id) + if file_unique_id == "boom": + raise RuntimeError("simulated download failure") + if file_unique_id == "flood": + raise errors.FloodWait(value=1) + # "ok" succeeds + + monkeypatch.setattr(api_server, "download_media_file", fake_download) + + worker = asyncio.create_task(api_server.background_download_worker()) + try: + for fid in ("boom", "flood", "ok"): + q.put_nowait(("chan", 1, fid)) + + # If the worker died on the first error, or task_done() were unbalanced, + # join() would never complete and this wait_for would time out. + await asyncio.wait_for(q.join(), timeout=5) + + # The worker stayed alive across the Exception and the FloodWait and drained all items. + assert processed == ["boom", "flood", "ok"] + finally: + worker.cancel() + with pytest.raises(asyncio.CancelledError): + await worker + + +# --------------------------------------------------------------------------- # +# 1.4 — _supervised: restart on crash (rate-limited), restart on unexpected +# return, and clean cancellation propagation on shutdown. +# --------------------------------------------------------------------------- # +@pytest.mark.asyncio +async def test_supervised_restarts_crashing_task_rate_limited(): + import api_server + + starts = [] + + async def crashing(): + starts.append(time.monotonic()) + raise RuntimeError("always dies") + + # Tiny min interval so the test is fast, but non-zero so we can assert spacing. + sup = asyncio.create_task( + api_server._supervised(crashing, "crasher", min_restart_interval=0.05) + ) + # Let it crash-and-restart a few times. + await asyncio.sleep(0.28) + sup.cancel() + with pytest.raises(asyncio.CancelledError): + await sup + + # It restarted several times (did not give up after the first crash)... + assert len(starts) >= 3 + # ...but restarts were spaced at least ~min_restart_interval apart (no spin). + gaps = [b - a for a, b in zip(starts, starts[1:])] + assert all(g >= 0.045 for g in gaps), gaps + + +@pytest.mark.asyncio +async def test_supervised_restarts_on_unexpected_return(): + import api_server + + runs = [] + + async def returns_immediately(): + runs.append(1) + # A supervised background loop is not meant to return; _supervised must + # log CRITICAL and restart it rather than stop supervising. + return + + sup = asyncio.create_task( + api_server._supervised(returns_immediately, "returner", min_restart_interval=0.02) + ) + await asyncio.sleep(0.1) + sup.cancel() + with pytest.raises(asyncio.CancelledError): + await sup + + assert len(runs) >= 2 # restarted after the unexpected return + + +@pytest.mark.asyncio +async def test_supervised_cancellation_propagates_to_child(): + import api_server + + child_cancelled = asyncio.Event() + + async def long_running(): + try: + await asyncio.Event().wait() # runs until cancelled + except asyncio.CancelledError: + child_cancelled.set() + raise + + sup = asyncio.create_task( + api_server._supervised(long_running, "longrun") + ) + await asyncio.sleep(0.05) # let the child start + sup.cancel() + # Cancelling the supervisor must propagate CancelledError out (clean shutdown)... + with pytest.raises(asyncio.CancelledError): + await sup + # ...and the child task must have been cancelled too (no leaked background task). + assert child_cancelled.is_set() diff --git a/tg_cache.py b/tg_cache.py index af83019..176f7c5 100644 --- a/tg_cache.py +++ b/tg_cache.py @@ -22,9 +22,12 @@ from typing import Any, Optional, Union, List from pyrogram import Client from pyrogram.types import Chat, Message from tg_throttle import tg_rpc +from config import get_settings logger = logging.getLogger(__name__) +Config = get_settings() + # Path to cache directory CACHE_DIR = os.path.join('data', 'tgcache') @@ -138,10 +141,16 @@ async def cached_get_chat_history(client: Client, channel_id: Union[str, int], l try: logger.info(f"history_cache_request: fetching fresh history for channel {channel_id}, limit {limit}") - messages = [] + # Hold the global RPC gate only for the live fetch. The paginated `async for` + # cannot be wrapped in wait_for directly, so collect it in an inner coroutine + # and bound THAT with the timeout. wait_for wraps only the RPC body, never the + # gate entry (`async with tg_rpc()`), so legitimate queue backpressure is not + # mistaken for a hang. async with tg_rpc(): - async for message in client.get_chat_history(channel_id, limit=limit): - messages.append(message) + async def _collect(): + # Full paginated history; bounded by the outer wait_for. + return [m async for m in client.get_chat_history(channel_id, limit=limit)] + messages = await asyncio.wait_for(_collect(), timeout=Config["tg_rpc_timeout"]) await asyncio.to_thread(_save_history_to_cache, channel_id, messages, limit) return messages @@ -214,7 +223,8 @@ async def cached_get_chat(client: Client, channel_id: Union[str, int]) -> Simple logger.info(f"chatinfo_cache_request: fetching fresh chat info for channel {channel_id}") async with tg_rpc(): - chat = await client.get_chat(channel_id) + # Bound the RPC itself, not the gate entry above it. + chat = await asyncio.wait_for(client.get_chat(channel_id), timeout=Config["tg_rpc_timeout"]) data = { 'id': getattr(chat, 'id', None),