diff --git a/api_server.py b/api_server.py index ac4aa22..2e162c8 100644 --- a/api_server.py +++ b/api_server.py @@ -917,7 +917,12 @@ async def get_raw_post_json(channel: str, post_id: int, request: Request, token: if isinstance(channel, str) and channel.startswith('-100'): channel_id = int(channel) - message = await client.client.get_messages(channel_id, post_id) + # Bound the RPC (stage-1 DoD: every Telegram RPC has a timeout). This + # endpoint is not under the tg_rpc gate, so a hang here only blocks this + # one request, but leaving it unbounded still violates the invariant. + message = await asyncio.wait_for( + client.client.get_messages(channel_id, post_id), timeout=30 + ) if not message: raise HTTPException(status_code=404, detail="Post not found") diff --git a/dockercompose.yml b/dockercompose.yml index a9156fa..384a8f2 100644 --- a/dockercompose.yml +++ b/dockercompose.yml @@ -24,6 +24,7 @@ services: # TG_CHAT_CACHE_TTL_HOURS: 12 # TTL for cached channel info (title/username/id); removes GetFullChannel from the poll hot path (default: 12) # TG_RPC_CONCURRENCY: 1 # Max concurrent live Telegram RPC calls — global throttle (default: 1) # TG_RPC_MIN_INTERVAL_MS: 500 # Minimum gap between live Telegram RPC starts, ms (default: 500) + # TG_RPC_TIMEOUT: 60 # Max seconds a single live Telegram RPC may run before timing out (default: 60) PYROGRAM_BRIDGE_URL: https://pgbridge.example.com API_PORT: 80 TOKEN: ХХХ diff --git a/requirements.txt b/requirements.txt index 2c19d8d..aafaac3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -11,3 +11,4 @@ python-dateutil==2.9.0.post0 python-magic==0.4.27 bleach[css]==6.1.0 types-bleach +pytest-asyncio==1.4.0 diff --git a/rss_generator.py b/rss_generator.py index b3cb822..a26ca11 100644 --- a/rss_generator.py +++ b/rss_generator.py @@ -22,7 +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 tg_throttle import tg_rpc_bounded from bleach.css_sanitizer import CSSSanitizer from bleach import clean as HTMLSanitizer @@ -503,13 +503,10 @@ 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: - # 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"], - ) + # Throttle under the global RPC gate and bound the call via the shared + # tg_rpc_bounded so a hung get_messages cannot pin the gate. + async with tg_rpc_bounded(Config["tg_rpc_timeout"]): + fetched = await client.get_messages(chat_id, ids_to_fetch) # get_messages may return a single Message or a list if not isinstance(fetched, list): fetched = [fetched] diff --git a/tests/pytest.ini b/tests/pytest.ini index 7a0abd3..6d6f517 100644 --- a/tests/pytest.ini +++ b/tests/pytest.ini @@ -1,4 +1,8 @@ [pytest] addopts = -ra +# The stage-1 anti-hang tests are async; run them without per-test event loops +# being set up by hand. (The tests also carry @pytest.mark.asyncio explicitly, so +# they work in strict mode too — but auto keeps the plugin's requirement obvious.) +asyncio_mode = auto filterwarnings = ignore::DeprecationWarning \ No newline at end of file diff --git a/tests/test_stage1_hangs.py b/tests/test_stage1_hangs.py index b4e43b5..a1b3104 100644 --- a/tests/test_stage1_hangs.py +++ b/tests/test_stage1_hangs.py @@ -97,9 +97,12 @@ async def test_gate_cancel_during_spacing_releases_permit(monkeypatch): 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 + # Record (and skip) the worker's sleeps so we can assert the FloodWait branch + # actually backed off — otherwise deleting `except FloodWait` (dropping it into + # the generic handler) leaves this test green. + sleeps = [] + async def _fast_sleep(delay=0, *_a, **_k): + sleeps.append(delay) monkeypatch.setattr(api_server.asyncio, "sleep", _fast_sleep) # Fresh queue so we don't interfere with (or depend on) module state. @@ -129,6 +132,13 @@ async def test_worker_survives_errors_and_balances_task_done(monkeypatch): # The worker stayed alive across the Exception and the FloodWait and drained all items. assert processed == ["boom", "flood", "ok"] + + # The FloodWait branch (caught BEFORE the generic Exception) backed off by + # min(value + 5, 900) = min(1 + 5, 900) = 6, distinct from the success path's + # post-download sleep of 2. Asserting the 6 pins the dedicated branch: if it + # were removed (FloodWait falling into `except Exception`), no 6 would appear. + assert 6 in sleeps, sleeps # FloodWait(value=1) -> backoff 6 + assert 2 in sleeps, sleeps # "ok" success -> post-download sleep 2 finally: worker.cancel() with pytest.raises(asyncio.CancelledError): diff --git a/tg_cache.py b/tg_cache.py index 176f7c5..da1cd8f 100644 --- a/tg_cache.py +++ b/tg_cache.py @@ -21,7 +21,7 @@ from types import SimpleNamespace from typing import Any, Optional, Union, List from pyrogram import Client from pyrogram.types import Chat, Message -from tg_throttle import tg_rpc +from tg_throttle import tg_rpc_bounded from config import get_settings logger = logging.getLogger(__name__) @@ -141,16 +141,12 @@ 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}") - # 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 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"]) + # Hold the global RPC gate for the live fetch and bound the RPC body with the + # timeout — gate outside, timeout inside — via the shared tg_rpc_bounded (so the + # tricky nesting is not re-derived here). The timeout covers the whole paginated + # fetch; see the note in tg_rpc_bounded. + async with tg_rpc_bounded(Config["tg_rpc_timeout"]): + messages = [m async for m in client.get_chat_history(channel_id, limit=limit)] await asyncio.to_thread(_save_history_to_cache, channel_id, messages, limit) return messages @@ -222,9 +218,8 @@ async def cached_get_chat(client: Client, channel_id: Union[str, int]) -> Simple return SimpleNamespace(**cached) logger.info(f"chatinfo_cache_request: fetching fresh chat info for channel {channel_id}") - async with tg_rpc(): - # 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"]) + async with tg_rpc_bounded(Config["tg_rpc_timeout"]): + chat = await client.get_chat(channel_id) data = { 'id': getattr(chat, 'id', None), diff --git a/tg_throttle.py b/tg_throttle.py index b96075c..fa120b0 100644 --- a/tg_throttle.py +++ b/tg_throttle.py @@ -9,6 +9,7 @@ import os import time import asyncio import logging +from contextlib import asynccontextmanager logger = logging.getLogger(__name__) @@ -67,3 +68,25 @@ class _TgRpcGate: def tg_rpc(): """Return an async context manager that throttles a single live Telegram RPC call.""" return _TgRpcGate() + + +@asynccontextmanager +async def tg_rpc_bounded(timeout: float): + """Throttle a live Telegram RPC AND bound it with a timeout, correctly nested. + + The single tricky invariant this centralizes: the timeout must bound ONLY the + RPC body, never the gate ENTRY — timing out the `_sem.acquire()` / spacing wait + would turn legitimate queue backpressure (e.g. ~47 feeds queueing) into false + timeouts. So the gate is the OUTER context and the timeout is the INNER one; a + TimeoutError raised inside propagates out through the gate's `__aexit__`, which + releases the permit (no leak). Call as: + + async with tg_rpc_bounded(Config["tg_rpc_timeout"]): + result = await client.get_chat(channel_id) + + Every gated+bounded RPC uses this so no call site re-derives the nesting by hand + (getting it wrong silently reopens the hang-under-backpressure class). + """ + async with _TgRpcGate(): + async with asyncio.timeout(timeout): + yield