From 4daf611f05f2c6258da8849ef8936c748dd42f1d Mon Sep 17 00:00:00 2001 From: claude code agent Date: Sun, 5 Jul 2026 06:51:08 +0300 Subject: [PATCH 01/14] =?UTF-8?q?fix(stability):=20stage=201=20=E2=80=94?= =?UTF-8?q?=20timeouts=20on=20all=20TG=20RPC=20+=20resilient=20download=20?= =?UTF-8?q?worker=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), From 03d1de2954e7d3d095fb579795bc12b3ab2dc081 Mon Sep 17 00:00:00 2001 From: claude code agent Date: Sun, 5 Jul 2026 07:12:10 +0300 Subject: [PATCH 02/14] fix(stability): declare pytest-asyncio, bound /raw_json RPC, assert FloodWait backoff, dedupe gate+timeout (review round 1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit F1 [WARNING] pytest-asyncio was not declared, so the 6 new async stage-1 tests ERROR on a clean checkout (async def not natively supported) — zero regression protection, masked by a globally-installed plugin. Added pytest-asyncio to requirements.txt + asyncio_mode=auto to tests/pytest.ini. Verified on a fresh venv from requirements alone: the tests collect and run (180 passed). F2 [WARNING] The /raw_json endpoint's get_messages was the one remaining live RPC without a timeout, violating the stage-1 DoD ('every Telegram RPC is bounded'). Wrapped it in asyncio.wait_for(..., 30) mirroring PostParser.get_post. (It is not under the tg_rpc gate, so its blast radius was one request, not the app.) F3 [WARNING] The worker test globally no-op'd asyncio.sleep, so the dedicated except FloodWait branch was indistinguishable from the generic handler — deleting it kept the test green. The sleep stub now records delays and the test asserts the FloodWait backoff of 6 (=min(1+5,900)), distinct from the success path's 2. F5 [low] The tricky 'gate outside, timeout inside' nesting was open-coded at 3 sites (each re-deriving the invariant). Extracted tg_rpc_bounded(timeout) into tg_throttle (using asyncio.timeout()); the 3 sites now use it, so a future call site cannot silently wrap the gate entry and reopen the hang-under-backpressure. F4 [low] Documented TG_RPC_TIMEOUT in the dockercompose.yml environment block next to the other TG_RPC_* knobs. Co-Authored-By: Claude Opus 4.8 (1M context) --- api_server.py | 7 ++++++- dockercompose.yml | 1 + requirements.txt | 1 + rss_generator.py | 13 +++++-------- tests/pytest.ini | 4 ++++ tests/test_stage1_hangs.py | 16 +++++++++++++--- tg_cache.py | 23 +++++++++-------------- tg_throttle.py | 23 +++++++++++++++++++++++ 8 files changed, 62 insertions(+), 26 deletions(-) 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 From 444bd3e42f29987e72b6d23438d31e6c749eb91f Mon Sep 17 00:00:00 2001 From: claude code agent Date: Sun, 5 Jul 2026 07:40:47 +0300 Subject: [PATCH 03/14] =?UTF-8?q?fix(stability):=20stage=202=20=E2=80=94?= =?UTF-8?q?=20atomic=20media=20downloads,=20in-flight=20dedup,=20FloodWait?= =?UTF-8?q?->429,=20bounded=20semaphore?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes flaky static media serving: partial files served as 'ready', truncated stubs living for an hour, and app-wide hangs under a saturated download path. 2.1 _download_atomic(file_id, final_path, timeout): downloads to a unique {final}.part.{hex}, validates size>0, and publishes solely via os.rename (atomic on POSIX); a finally ALWAYS removes the partial (timeout/cancel/ zero-size/race-loser). Every downloader call routes through it, so a file at a FINAL name ({fid} / temp_{fid}) is GUARANTEED complete (grep-proven: the only safe_download_media caller passes a .part. path). Big-video timeout scales with size (min/max/min-speed knobs, documented in dockercompose.yml). The sweeper regex now matches both .part. and legacy .tmp. stubs. 2.2 In-flight dedup registry: the first request for a key runs the download in a DETACHED task sharing a Future; the task's finally sets the Future AND pops the key. Waiters await asyncio.shield(fut) so a client disconnect / waiter timeout cancels only the waiter, never the download — no hung waiters, no stuck key, a failed download frees the key for retry. 2.3 FloodWait->429: handler before except RPCError (FloodWait subclasses it), Retry-After = min(value + rand(1,30), 300); propagates from the detached task through the Future. 2.4 Serving a temp_* file touches its mtime so the 1h sweeper can't delete a video out from under a viewer. 2.5 The HTTP download semaphore acquire is bounded (wait_for 30 -> 503 + Retry-After); the permit is released only if the acquire succeeded. The request-scoped-permit trade-off (a disconnect can transiently exceed the download count) is documented inline for stage-7 prod observation. Tests (tests/test_stage2_static.py, 13): atomic publish/clean on every exit incl. FloodWait-through-the-finally; concurrent big-video serves no partial; dedup runs one download, a cancelled waiter doesn't hang others, a failed download frees the key; FloodWait -> 429; mtime touch; sweeper cleans .part./.tmp./stale temp_ but keeps fresh files. 193 passed (180 baseline + 13). Co-Authored-By: Claude Opus 4.8 (1M context) --- api_server.py | 205 ++++++++++++++++----- config.py | 6 + dockercompose.yml | 3 + telegram_client.py | 9 +- tests/mock_config.py | 3 + tests/test_stage2_static.py | 342 ++++++++++++++++++++++++++++++++++++ 6 files changed, 523 insertions(+), 45 deletions(-) create mode 100644 tests/test_stage2_static.py diff --git a/api_server.py b/api_server.py index 2e162c8..e0e55ff 100644 --- a/api_server.py +++ b/api_server.py @@ -72,6 +72,16 @@ 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) +# In-flight download dedup registry: maps (channel, post_id, file_unique_id) to the +# shared Future of an ongoing download. The FIRST request for a key runs the download in +# a DETACHED task and shares its Future; concurrent requests await that Future instead of +# starting their own download. Detaching the download means a client disconnecting (which +# cancels its request coroutine) can never cancel the download nor leave the Future +# forever-pending — so waiters can't hang. See _download_deduped. +_inflight: dict[tuple[str, int, str], asyncio.Future] = {} +# Strong refs to detached download tasks so they are not garbage-collected mid-flight. +_inflight_tasks: set[asyncio.Task] = set() + 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. @@ -239,6 +249,14 @@ async def prepare_file_response(file_path: str, request: Request, delete_after: if not os.path.exists(file_path): raise HTTPException(status_code=404, detail="File not found") + # Keep an actively-viewed large-video temp file alive: refresh its mtime so the 1h + # sweeper (which deletes temp_* by mtime) won't remove it out from under a viewer. + if os.path.basename(file_path).startswith("temp_"): + try: + await asyncio.to_thread(os.utime, file_path, None) + except OSError as e: + logger.debug(f"Failed to refresh mtime for {file_path}: {e}") + media_type: str | None = None if media_key is not None: @@ -367,6 +385,96 @@ async def prepare_file_response(file_path: str, request: Request, delete_after: background=background, ) +def _media_download_timeout(file_size: int) -> float: + """Scale the download timeout with file size. + + Effective floor of `media_download_min_speed` bytes/s (timeout ≈ size / min_speed), + clamped to [media_download_timeout_min, media_download_timeout_max] seconds. + A non-positive/unknown size falls back to the minimum timeout. + """ + min_t = Config["media_download_timeout_min"] + max_t = Config["media_download_timeout_max"] + min_speed = Config["media_download_min_speed"] + if file_size <= 0 or min_speed <= 0: + return float(min_t) + return float(min(max_t, max(min_t, file_size // min_speed))) + + +async def _download_atomic(file_id: str, final_path: str, timeout: float) -> str: + """Download a media file through a unique partial path and atomically publish it. + + Invariant enforced here: a file that exists at a FINAL name (`{file_unique_id}` or + `temp_{file_unique_id}`) is ALWAYS complete. The downloader only ever writes to a + unique `{final_path}.part.{hex}` path; the file appears at its final name solely via + os.rename (atomic on POSIX). The finally block ALWAYS removes our partial (on timeout, + cancel, zero-size, or losing a rename race), so no stub is ever served or left behind. + + Raises ZeroSizeFileError if the download produced a missing/zero-size file. + """ + part_path = f"{final_path}.part.{uuid.uuid4().hex}" + try: + await client.safe_download_media(file_id, part_path, timeout=timeout) + if not os.path.exists(part_path) or os.path.getsize(part_path) == 0: + raise ZeroSizeFileError( + f"Downloaded file for {final_path} is zero size or missing after download attempt." + ) + # Publish atomically, but only if nobody else already produced the final file + # (a concurrent request that won the race). rename is atomic on POSIX. + if not os.path.exists(final_path): + os.rename(part_path, final_path) + return final_path + finally: + # Always clean up our partial: timeout, cancellation, zero-size, or race loser + # (rename skipped because final_path already existed). + if os.path.exists(part_path): + try: + os.remove(part_path) + except OSError as e: + logger.warning(f"cleanup_error: Failed to remove partial file {part_path}: {e}") + + +async def _download_deduped(channel: Union[str, int], post_id: int, file_unique_id: str) -> tuple[Union[str, None], bool]: + """Deduplicate concurrent downloads of the same media by (channel, post_id, fid). + + The first request for a key runs download_media_file in a DETACHED task and shares its + Future; concurrent requests await the same Future (bounded by wait_for). The detached + task's finally BOTH sets the Future's result/exception AND pops the key — so a + completed Future never leaves the key stuck, and a client disconnect (cancelling only + the awaiting request coroutine) can neither cancel the download nor hang other waiters. + """ + key = (str(channel), post_id, file_unique_id) + fut = _inflight.get(key) + if fut is None: + loop = asyncio.get_running_loop() + fut = loop.create_future() + _inflight[key] = fut + + async def _runner(): + try: + result = await download_media_file(channel, post_id, file_unique_id) + if not fut.done(): + fut.set_result(result) + except BaseException as e: # noqa: BLE001 — must forward ANY failure to waiters + if not fut.done(): + fut.set_exception(e) + finally: + _inflight.pop(key, None) + + task = asyncio.create_task(_runner()) + _inflight_tasks.add(task) + task.add_done_callback(_inflight_tasks.discard) + # Retrieve the exception in a done-callback to silence "exception was never + # retrieved" if every waiter disconnects before awaiting the Future. + fut.add_done_callback(lambda f: f.cancelled() or f.exception()) + + # Waiter timeout: a generous safety net above the maximum per-download timeout (the + # download is itself internally bounded), so a live download always completes first. + waiter_timeout = float(Config["media_download_timeout_max"]) + 120.0 + # shield so a timed-out / cancelled waiter does NOT cancel the shared Future that the + # detached task owns and other waiters depend on. + return await asyncio.wait_for(asyncio.shield(fut), timeout=waiter_timeout) + + async def download_media_file(channel: Union[str, int], post_id: int, file_unique_id: str) -> tuple[Union[str, None], bool]: """ Download media file from Telegram and save to cache @@ -410,18 +518,27 @@ async def download_media_file(channel: Union[str, int], post_id: int, file_uniqu logger.error(f"Failed to get video file size for message {post_id}: {str(e)}") if is_large_video: - # For large video, download without permanent caching; use a temporary file + # For large video, download without permanent caching; use a temporary file. temp_file_path = os.path.join(post_dir, f"temp_{file_unique_id}") - if os.path.exists(temp_file_path): + # A temp_{fid} file now only ever appears via atomic rename (see _download_atomic), + # so its presence GUARANTEES a complete file — serve it directly. + if os.path.exists(temp_file_path) and os.path.getsize(temp_file_path) > 0: logger.info(f"Temporary file {temp_file_path} already exists, serving cached large video") return temp_file_path, False file_id = await find_file_id_in_message(message, file_unique_id) if not file_id: logger.error(f"Media with file_unique_id {file_unique_id} not found in message {post_id}") raise HTTPException(status_code=404, detail="File not found in message") - logger.info(f"Downloading large video file {file_unique_id} to temporary path {temp_file_path}") + # Scale the timeout with size (≈256 KB/s floor) so slow big-video downloads aren't + # cut short at the fixed 120s used for regular files. try: - file_path = await client.safe_download_media(file_id, temp_file_path) + file_size = int(message.video.file_size or 0) + except Exception: + file_size = 0 + timeout = _media_download_timeout(file_size) + logger.info(f"Downloading large video file {file_unique_id} to temporary path {temp_file_path} (timeout={timeout:.0f}s)") + try: + file_path = await _download_atomic(file_id, temp_file_path, timeout) except asyncio.TimeoutError: logger.error(f"Timeout downloading large video {file_unique_id}") raise HTTPException(status_code=504, detail="Download timeout") @@ -470,46 +587,17 @@ async def download_media_file(channel: Union[str, int], post_id: int, file_uniqu raise HTTPException(status_code=404, detail="File not found in message") - # Use a unique temp path to avoid race conditions when concurrent requests download the same file - temp_path = cache_path + f".tmp.{uuid.uuid4().hex}" + # Download through a unique `.part.` file and atomically publish to the final cache + # path. _download_atomic owns the zero-size check, the race-loser cleanup, and the + # rename-only-if-absent logic, keeping the "final name = complete file" invariant. try: - file_path = await client.safe_download_media(file_id, temp_path) + file_path = await _download_atomic(file_id, cache_path, timeout=float(Config["media_download_timeout_min"])) except asyncio.TimeoutError: logger.error(f"Timeout downloading media {file_unique_id}") - # Clean up the temp file if it was created - if os.path.exists(temp_path): - try: - os.remove(temp_path) - except OSError: - pass raise HTTPException(status_code=504, detail="Download timeout") - # Check if the downloaded temp file exists and has a size greater than zero - if not file_path or not os.path.exists(temp_path) or os.path.getsize(temp_path) == 0: - logger.error(f"download_failed_zero_size: Downloaded file {file_unique_id} for {channel}/{post_id} is zero size or missing.") - # Attempt to clean up the invalid temp file - if os.path.exists(temp_path): - try: - os.remove(temp_path) - logger.info(f"Removed zero-size temp file: {temp_path}") - except OSError as e: - logger.error(f"cleanup_error: Failed to remove zero-size temp file {temp_path}: {e}") - # Raise specific error to indicate download failure - raise ZeroSizeFileError(f"Downloaded file {file_unique_id} for {channel}/{post_id} is zero size or missing after download attempt.") - - # Atomic rename: if another concurrent request already wrote the final file, discard our temp copy - if os.path.exists(cache_path) and os.path.getsize(cache_path) > 0: - logger.info(f"Concurrent download already completed for {file_unique_id}, discarding temp file") - try: - os.remove(temp_path) - except OSError as e: - logger.warning(f"cleanup_error: Failed to remove temp file {temp_path}: {e}") - return cache_path, False - - # Atomically move temp file to final cache path (atomic on POSIX) - os.rename(temp_path, cache_path) logger.info(f"Downloaded media file {file_unique_id} to {cache_path}") - return cache_path, False + return file_path, False def remove_old_cached_files_sync(media_files: list, cache_dir: str) -> tuple[list, int]: @@ -565,8 +653,9 @@ def remove_old_cached_files_sync(media_files: list, cache_dir: str) -> tuple[lis for root, _, files in os.walk(cache_dir): for file in files: is_large_video_temp = file.startswith("temp_") - # Match exactly the format generated in download_media_file: {name}.tmp.{32 hex chars} - is_race_temp = bool(re.match(r'^.+\.tmp\.[0-9a-f]{32}$', file)) + # Match the partial-download suffixes: the new `.part.{32 hex}` and the legacy + # `.tmp.{32 hex}` (old stubs may still be on disk after the rename to .part.). + is_race_temp = bool(re.match(r'^.+\.(part|tmp)\.[0-9a-f]{32}$', file)) if not (is_large_video_temp or is_race_temp): continue file_path = os.path.join(root, file) @@ -1020,14 +1109,39 @@ async def get_media(channel: str, post_id: int, file_unique_id: str, request: Re media_key=(str(channel), post_id, file_unique_id)) _sem_wait_start = _time.monotonic() - async with HTTP_DOWNLOAD_SEMAPHORE: # limit concurrent live HTTP downloads + # Bound the wait for a live-download permit: a saturated semaphore must not + # hang the request indefinitely. wait_for wraps ONLY the acquire; the permit + # is released in the finally below only if the acquire actually succeeded (a + # timed-out acquire never holds a permit, so it must not release one). + # + # CONSCIOUS TRADE-OFF (stage-2 review): the permit is held by the REQUEST, not + # the detached download task. So this bounds concurrent *requests* (admission + # control with a fast 503 on saturation), not strictly concurrent *downloads*: + # if a client disconnects, its request coroutine releases the permit while the + # detached download keeps running, so a burst of disconnects across DIFFERENT + # files can transiently exceed the limit of live Telegram downloads. This is + # deliberate — the download surviving a disconnect is the whole point of the + # in-flight registry, and moving the permit into the runner would forfeit the + # fast 503 (a saturated request would instead hang on the future for up to the + # waiter timeout). Each download is itself timeout-bounded, so the overrun is + # transient. If prod observation (stage 7) shows real download-count overrun + # under disconnect churn, strengthen with a separate download-side limiter. + try: + await asyncio.wait_for(HTTP_DOWNLOAD_SEMAPHORE.acquire(), timeout=30) + except asyncio.TimeoutError: + logger.warning(f"http_semaphore_timeout: {channel}/{post_id}/{file_unique_id} waited >30s for HTTP_DOWNLOAD_SEMAPHORE") + return Response(status_code=503, content="Server busy, please retry", + headers={"Retry-After": "30"}) + try: _sem_wait = _time.monotonic() - _sem_wait_start if _sem_wait > 0.5: logger.warning(f"diag_semaphore_wait: {channel}/{post_id}/{file_unique_id} waited {_sem_wait:.3f}s for HTTP_DOWNLOAD_SEMAPHORE") _dl_start = _time.monotonic() - file_path, delete_after = await download_media_file(channel_id, post_id, file_unique_id) + file_path, delete_after = await _download_deduped(channel_id, post_id, file_unique_id) _dl_elapsed = _time.monotonic() - _dl_start logger.info(f"diag_download_timing: {channel}/{post_id}/{file_unique_id} download_media_file took {_dl_elapsed:.3f}s (semaphore_wait={_sem_wait:.3f}s)") + finally: + HTTP_DOWNLOAD_SEMAPHORE.release() if not file_path: raise HTTPException(status_code=404, detail="File not found") if file_path: @@ -1043,6 +1157,13 @@ async def get_media(channel: str, post_id: int, file_unique_id: str, request: Re except HTTPException: raise + except errors.FloodWait as e: + # MUST precede `except errors.RPCError` (FloodWait subclasses RPCError): a temporary + # Telegram throttle must return a retryable 429, never a permanent 404. + retry_after = min(int(e.value) + random.randint(1, 30), 300) + logger.warning(f"media_flood_wait: {channel}/{post_id}/{file_unique_id} retry after {retry_after}s") + return Response(status_code=429, content="Telegram flood wait", + headers={"Retry-After": str(retry_after)}) except errors.RPCError as e: logger.error(f"Media request RPC error for {channel}/{post_id}/{file_unique_id}: {type(e).__name__} - {str(e)}") raise HTTPException(status_code=404, detail="File not found in Telegram") from e diff --git a/config.py b/config.py index 03e47bd..0053cb0 100644 --- a/config.py +++ b/config.py @@ -123,4 +123,10 @@ def get_settings() -> dict[str, Any]: "tg_watchdog_heartbeat_every": _parse_int_env("TG_WATCHDOG_HEARTBEAT_EVERY", 30), "tg_disconnect_flap_limit": _parse_int_env("TG_DISCONNECT_FLAP_LIMIT", 3), "tg_disconnect_flap_window": _parse_int_env("TG_DISCONNECT_FLAP_WINDOW", 120), + # Media download timeout scales with file size (large videos): the per-download + # timeout is clamped to [min, max] seconds, with an effective floor of + # `media_download_min_speed` bytes/s (timeout ≈ file_size / min_speed). + "media_download_timeout_min": _parse_int_env("MEDIA_DOWNLOAD_TIMEOUT_MIN", 120), + "media_download_timeout_max": _parse_int_env("MEDIA_DOWNLOAD_TIMEOUT_MAX", 1800), + "media_download_min_speed": _parse_int_env("MEDIA_DOWNLOAD_MIN_SPEED", 256 * 1024), } diff --git a/dockercompose.yml b/dockercompose.yml index 384a8f2..800c685 100644 --- a/dockercompose.yml +++ b/dockercompose.yml @@ -25,6 +25,9 @@ services: # 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) + # MEDIA_DOWNLOAD_TIMEOUT_MIN: 120 # Min per-download timeout, seconds — also the timeout for regular (non-large) files (default: 120) + # MEDIA_DOWNLOAD_TIMEOUT_MAX: 1800 # Max per-download timeout, seconds — cap for the largest videos (default: 1800) + # MEDIA_DOWNLOAD_MIN_SPEED: 262144 # Assumed floor download speed, bytes/s — large-video timeout ≈ file_size / this, clamped to [MIN,MAX] (default: 262144 = 256 KB/s) PYROGRAM_BRIDGE_URL: https://pgbridge.example.com API_PORT: 80 TOKEN: ХХХ diff --git a/telegram_client.py b/telegram_client.py index abba027..da9793e 100644 --- a/telegram_client.py +++ b/telegram_client.py @@ -274,13 +274,16 @@ class TelegramClient: continue raise - async def safe_download_media(self, file_id, file_name, max_retries=2): - """Wrapper with retry logic for download errors""" + async def safe_download_media(self, file_id, file_name, max_retries=2, timeout: float = 120.0): + """Wrapper with retry logic for download errors. + + `timeout` bounds each download attempt; for large videos the caller scales it + with file size (see api_server._media_download_timeout).""" for attempt in range(max_retries): try: return await asyncio.wait_for( self.client.download_media(file_id, file_name=file_name), - timeout=120.0 + timeout=timeout ) except Exception as e: if isinstance(e, KeyError) and attempt < max_retries - 1: diff --git a/tests/mock_config.py b/tests/mock_config.py index bd22f42..6668616 100644 --- a/tests/mock_config.py +++ b/tests/mock_config.py @@ -33,4 +33,7 @@ def get_settings(): "tg_watchdog_heartbeat_every": 30, "tg_disconnect_flap_limit": 3, "tg_disconnect_flap_window": 120, + "media_download_timeout_min": 120, + "media_download_timeout_max": 1800, + "media_download_min_speed": 256 * 1024, } diff --git a/tests/test_stage2_static.py b/tests/test_stage2_static.py new file mode 100644 index 0000000..28893b3 --- /dev/null +++ b/tests/test_stage2_static.py @@ -0,0 +1,342 @@ +# 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 2 (static-media stability) regression tests. + +Covers the four scenarios from the plan plus the subtle in-flight-dedup lifecycle: +- _download_atomic: success publishes atomically and leaves no `.part.`; timeout / zero-size + / race-loser all remove the partial and never leave a stub at the final name. +- Concurrent large-video downloads never expose a partial file (atomic rename). +- _download_deduped: one shared download for concurrent requests; a cancelled waiter + (client disconnect) can't hang the others; a failed download propagates to all waiters + and frees the key (no stuck registry entry). +- FloodWait in /media returns a retryable 429 (not a permanent 404) — ordering pinned. +- A served temp_* file gets its mtime refreshed (sweeper won't delete it under a viewer). +- The sweeper removes both `.part.` and legacy `.tmp.` stubs (and stale temp_*) but keeps + fresh partials and non-temp files. +""" +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 api_server +from pyrogram import errors +from pyrogram.enums import MessageMediaType + + +def _part_files(dir_path): + return [p for p in os.listdir(dir_path) if ".part." in p or ".tmp." in p] + + +# --------------------------------------------------------------------------- # +# 2.1 — _download_atomic: publish-on-success, always clean the partial. +# --------------------------------------------------------------------------- # +async def test_download_atomic_success_publishes_and_cleans(monkeypatch, tmp_path): + final = str(tmp_path / "final") + + async def fake(file_id, file_name, timeout=None, **kw): + with open(file_name, "wb") as f: + f.write(b"hello") + return file_name + + monkeypatch.setattr(api_server.client, "safe_download_media", fake) + + res = await api_server._download_atomic("fid", final, timeout=10) + assert res == final + with open(final, "rb") as f: + assert f.read() == b"hello" + assert _part_files(tmp_path) == [] # no leftover partial + + +async def test_download_atomic_timeout_removes_part_no_final(monkeypatch, tmp_path): + final = str(tmp_path / "final") + + async def fake(file_id, file_name, timeout=None, **kw): + # A partial write, then the download times out mid-flight. + with open(file_name, "wb") as f: + f.write(b"partial") + raise asyncio.TimeoutError() + + monkeypatch.setattr(api_server.client, "safe_download_media", fake) + + with pytest.raises(asyncio.TimeoutError): + await api_server._download_atomic("fid", final, timeout=10) + + assert not os.path.exists(final) # no stub at the final name + assert _part_files(tmp_path) == [] # partial cleaned by the finally + + +async def test_download_atomic_zero_size_raises_and_cleans(monkeypatch, tmp_path): + final = str(tmp_path / "final") + + async def fake(file_id, file_name, timeout=None, **kw): + open(file_name, "wb").close() # zero-size result + return file_name + + monkeypatch.setattr(api_server.client, "safe_download_media", fake) + + with pytest.raises(api_server.ZeroSizeFileError): + await api_server._download_atomic("fid", final, timeout=10) + + assert not os.path.exists(final) + assert _part_files(tmp_path) == [] + + +async def test_download_atomic_race_loser_keeps_existing_final(monkeypatch, tmp_path): + final = str(tmp_path / "final") + with open(final, "wb") as f: + f.write(b"WINNER") # a concurrent request already published the final file + + async def fake(file_id, file_name, timeout=None, **kw): + with open(file_name, "wb") as f: + f.write(b"loser") + return file_name + + monkeypatch.setattr(api_server.client, "safe_download_media", fake) + + res = await api_server._download_atomic("fid", final, timeout=10) + assert res == final + with open(final, "rb") as f: + assert f.read() == b"WINNER" # not clobbered by the race loser + assert _part_files(tmp_path) == [] # loser's partial removed + + +# --------------------------------------------------------------------------- # +# 2.1 — Concurrent large-video downloads never expose a partial file. +# --------------------------------------------------------------------------- # +async def test_concurrent_large_video_no_partial_served(monkeypatch, tmp_path): + monkeypatch.chdir(tmp_path) # download_media_file writes under ./data/cache + + msg = SimpleNamespace( + media=MessageMediaType.VIDEO, + video=SimpleNamespace(file_size=200 * 1024 * 1024), # > 100 MB -> large-video path + empty=False, + ) + + async def fake_get(channel_id, post_id): + return msg + + monkeypatch.setattr(api_server.client, "safe_get_messages", fake_get) + + async def fake_find(message, fid): + return "fileid" + + monkeypatch.setattr(api_server, "find_file_id_in_message", fake_find) + + started = [] + + async def fake_dl(file_id, file_name, timeout=None, **kw): + started.append(file_name) + await asyncio.sleep(0.1) # slow download -> real overlap between the two requests + with open(file_name, "wb") as f: + f.write(b"VIDEODATA") + return file_name + + monkeypatch.setattr(api_server.client, "safe_download_media", fake_dl) + + t1 = asyncio.create_task(api_server.download_media_file("chan", 7, "vidfid")) + t2 = asyncio.create_task(api_server.download_media_file("chan", 7, "vidfid")) + (p1, _), (p2, _) = await asyncio.gather(t1, t2) + + assert p1 == p2 + assert os.path.basename(p1) == "temp_vidfid" # large videos keep the TTL-cleaned name + with open(p1, "rb") as f: + assert f.read() == b"VIDEODATA" # a complete file, never a partial + + post_dir = os.path.dirname(p1) + assert _part_files(post_dir) == [] # both partials resolved, none left on disk + # Both requests really downloaded to distinct `.part.` paths (a genuine race that the + # atomic rename resolved), so neither ever saw the other's partial. + assert len(started) == 2 and started[0] != started[1] + + +# --------------------------------------------------------------------------- # +# 2.2 — In-flight dedup: one shared download for concurrent requests. +# --------------------------------------------------------------------------- # +async def test_dedup_single_download_for_concurrent_requests(monkeypatch): + api_server._inflight.clear() + calls = [] + + async def fake_dl(channel, post_id, fid): + calls.append(fid) + await asyncio.sleep(0.05) + return (f"/final/{fid}", False) + + monkeypatch.setattr(api_server, "download_media_file", fake_dl) + + tasks = [asyncio.create_task(api_server._download_deduped("c", 1, "f")) for _ in range(5)] + results = await asyncio.gather(*tasks) + + assert all(r == ("/final/f", False) for r in results) + assert calls == ["f"] # exactly one real download, shared by all 5 + assert not api_server._inflight # key popped after completion + + +# --------------------------------------------------------------------------- # +# 2.2 — A cancelled waiter (client disconnect) must not hang the others, and the +# detached download runs to completion regardless. +# --------------------------------------------------------------------------- # +async def test_dedup_waiter_cancel_does_not_hang_others(monkeypatch): + api_server._inflight.clear() + gate = asyncio.Event() + calls = [] + + async def fake_dl(channel, post_id, fid): + calls.append(fid) + await gate.wait() # hold the download open until the test releases it + return (f"/path/{fid}", False) + + monkeypatch.setattr(api_server, "download_media_file", fake_dl) + + t1 = asyncio.create_task(api_server._download_deduped("c", 1, "fidX")) + await asyncio.sleep(0.05) # let t1 register the future + start the detached task + t2 = asyncio.create_task(api_server._download_deduped("c", 1, "fidX")) + await asyncio.sleep(0.05) + + # First client disconnects -> its request coroutine is cancelled. + t1.cancel() + with pytest.raises(asyncio.CancelledError): + await t1 + + # The detached download is unaffected; complete it and the surviving waiter resolves. + gate.set() + res = await asyncio.wait_for(t2, timeout=2) + assert res == ("/path/fidX", False) + assert calls == ["fidX"] # download ran exactly once (shared, not restarted) + assert not api_server._inflight # key popped + + +# --------------------------------------------------------------------------- # +# 2.2 — A failed download propagates to all waiters and frees the key (not stuck). +# --------------------------------------------------------------------------- # +async def test_dedup_exception_propagates_and_frees_key(monkeypatch): + api_server._inflight.clear() + state = {"mode": "boom"} + + async def fake_dl(channel, post_id, fid): + if state["mode"] == "boom": + raise RuntimeError("kaboom") + return ("/ok", False) + + monkeypatch.setattr(api_server, "download_media_file", fake_dl) + + t1 = asyncio.create_task(api_server._download_deduped("c", 1, "f")) + t2 = asyncio.create_task(api_server._download_deduped("c", 1, "f")) + results = await asyncio.gather(t1, t2, return_exceptions=True) + assert all(isinstance(r, RuntimeError) for r in results), results + assert not api_server._inflight # failed future did NOT leave the key stuck + + # A subsequent request with a working download succeeds (proves the key was freed). + state["mode"] = "ok" + res = await api_server._download_deduped("c", 1, "f") + assert res == ("/ok", False) + + +# --------------------------------------------------------------------------- # +# 2.3 — FloodWait in /media -> 429 with Retry-After (ordering pinned: not a 404). +# --------------------------------------------------------------------------- # +async def test_media_floodwait_returns_429_not_404(monkeypatch, tmp_path): + api_server._inflight.clear() + monkeypatch.chdir(tmp_path) # so the pre-semaphore cache check misses -> download path + monkeypatch.setattr(api_server, "verify_media_digest", lambda url, digest: True) + + async def fake_dl(channel, post_id, fid): + raise errors.FloodWait(value=42) + + monkeypatch.setattr(api_server, "download_media_file", fake_dl) + + req = SimpleNamespace(headers={}) + resp = await api_server.get_media("chan", 5, "floodfid", request=req, digest="x") + + assert resp.status_code == 429 + retry_after = resp.headers.get("retry-after") + assert retry_after is not None + assert 43 <= int(retry_after) <= 72 # 42 + random(1..30) + + +async def test_download_atomic_reraises_floodwait_and_cleans_part(monkeypatch, tmp_path): + # Stage-2 review: prove a FloodWait raised by the downloader propagates OUT of + # _download_atomic (its finally must NOT swallow it) AND the partial is cleaned. + # This is the path the wholesale-mocked 429 test above skips. + final = str(tmp_path / "temp_vid") + seen_part = {} + + async def fake(file_id, part_path, timeout=120): + # The partial path exists mid-download, then FloodWait strikes. + with open(part_path, "wb") as fh: + fh.write(b"partial") + seen_part["path"] = part_path + raise errors.FloodWait(value=7) + + monkeypatch.setattr(api_server.client, "safe_download_media", fake) + + with pytest.raises(errors.FloodWait): + await api_server._download_atomic("fid", final, timeout=10) + + # The finally cleaned the partial, and no final file was published. + assert not os.path.exists(seen_part["path"]) + assert not os.path.exists(final) + + +# --------------------------------------------------------------------------- # +# 2.4 — Serving a temp_* file refreshes its mtime; non-temp files are untouched. +# --------------------------------------------------------------------------- # +async def test_temp_file_mtime_touched_on_serve(tmp_path): + temp_file = tmp_path / "temp_abc" + temp_file.write_bytes(b"videodata") + stale = time.time() - 10000 + os.utime(temp_file, (stale, stale)) + + req = SimpleNamespace(headers={}) + resp = await api_server.prepare_file_response(str(temp_file), request=req) + assert resp is not None + assert os.path.getmtime(temp_file) > time.time() - 100 # refreshed + + +async def test_non_temp_file_mtime_not_touched(tmp_path): + normal = tmp_path / "regularfile" + normal.write_bytes(b"data") + stale = time.time() - 10000 + os.utime(normal, (stale, stale)) + + req = SimpleNamespace(headers={}) + await api_server.prepare_file_response(str(normal), request=req) + assert os.path.getmtime(normal) < time.time() - 5000 # left alone + + +# --------------------------------------------------------------------------- # +# 2.1/2.4 — Sweeper cleans both `.part.` and legacy `.tmp.` stubs and stale temp_*, +# but keeps fresh partials and ordinary cached files. +# --------------------------------------------------------------------------- # +def test_sweeper_cleans_part_and_tmp_not_fresh(tmp_path): + cache = tmp_path / "cache" + post = cache / "chan" / "10" + post.mkdir(parents=True) + + stale = time.time() - 4000 # older than the 1h (3600s) threshold + hexid = "0" * 32 + + old_part = post / f"file.part.{hexid}"; old_part.write_bytes(b"x"); os.utime(old_part, (stale, stale)) + old_tmp = post / f"file.tmp.{hexid}"; old_tmp.write_bytes(b"x"); os.utime(old_tmp, (stale, stale)) + old_tempvid = post / "temp_bigvid"; old_tempvid.write_bytes(b"x"); os.utime(old_tempvid, (stale, stale)) + fresh_part = post / f"file2.part.{hexid}"; fresh_part.write_bytes(b"x") # fresh mtime + normal = post / "realfile"; normal.write_bytes(b"x") # not a temp file at all + + _, removed = api_server.remove_old_cached_files_sync([], str(cache)) + + assert not old_part.exists() # new-suffix stub swept + assert not old_tmp.exists() # legacy-suffix stub swept (transition period) + assert not old_tempvid.exists() # stale large-video temp swept + assert fresh_part.exists() # too new to sweep + assert normal.exists() # ordinary cached file untouched + assert removed >= 3 From 41d143458f8d0c926c4531a02813f2aa995b4b37 Mon Sep 17 00:00:00 2001 From: claude code agent Date: Sun, 5 Jul 2026 07:56:23 +0300 Subject: [PATCH 04/14] fix(stability): test the HTTP semaphore balance + correct the dedup docstring (review round 1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit F1 [WARNING] The HTTP_DOWNLOAD_SEMAPHORE admission-control balance was untested, and it is a plain asyncio.Semaphore — so an over-release would SILENTLY inflate the permit count and disable the limiter (no ValueError), a green suite hiding the exact bug stage 2 fixes. Added two tests (mirroring the stage-1 tg_throttle balance test): the permit is released exactly once when the download errors (count back to baseline), and a timed-out acquire (503) releases NOTHING (count unchanged). Adversarially validated: neutralizing the finally release makes the first test fail, so it genuinely catches a leak. F2 [low] Corrected the _download_deduped docstring: the detached task sets the Future in try/except, and its finally ALWAYS pops the key — the previous wording attributed the Future-set to the finally, which an auditor reading the finally block would find contradicted. Co-Authored-By: Claude Opus 4.8 (1M context) --- api_server.py | 7 +++-- tests/test_stage2_static.py | 61 +++++++++++++++++++++++++++++++++++++ 2 files changed, 65 insertions(+), 3 deletions(-) diff --git a/api_server.py b/api_server.py index e0e55ff..7f59ef4 100644 --- a/api_server.py +++ b/api_server.py @@ -438,9 +438,10 @@ async def _download_deduped(channel: Union[str, int], post_id: int, file_unique_ The first request for a key runs download_media_file in a DETACHED task and shares its Future; concurrent requests await the same Future (bounded by wait_for). The detached - task's finally BOTH sets the Future's result/exception AND pops the key — so a - completed Future never leaves the key stuck, and a client disconnect (cancelling only - the awaiting request coroutine) can neither cancel the download nor hang other waiters. + task sets the Future's result/exception (on success/failure) and its finally ALWAYS + pops the key — so both happen before the task ends: a completed Future never leaves the + key stuck, and a client disconnect (cancelling only the awaiting request coroutine) can + neither cancel the download nor hang other waiters. """ key = (str(channel), post_id, file_unique_id) fut = _inflight.get(key) diff --git a/tests/test_stage2_static.py b/tests/test_stage2_static.py index 28893b3..a8a6538 100644 --- a/tests/test_stage2_static.py +++ b/tests/test_stage2_static.py @@ -264,6 +264,67 @@ async def test_media_floodwait_returns_429_not_404(monkeypatch, tmp_path): assert 43 <= int(retry_after) <= 72 # 42 + random(1..30) +# --------------------------------------------------------------------------- # +# 2.5 — HTTP_DOWNLOAD_SEMAPHORE balance: released exactly once on success/error, +# never released when the acquire itself timed out (503). It is a plain +# asyncio.Semaphore, so an over-release would SILENTLY inflate the permit count +# and disable the limiter — assert the count, mirroring the stage-1 gate test. +# --------------------------------------------------------------------------- # +async def test_media_semaphore_released_on_download_error(monkeypatch, tmp_path): + api_server._inflight.clear() + monkeypatch.chdir(tmp_path) # pre-semaphore cache miss -> the acquire path runs + monkeypatch.setattr(api_server, "verify_media_digest", lambda url, digest: True) + + permits_before = api_server.HTTP_DOWNLOAD_SEMAPHORE._value + + async def boom(channel_id, post_id, fid): + raise RuntimeError("download blew up") + + monkeypatch.setattr(api_server, "_download_deduped", boom) + + req = SimpleNamespace(headers={}) + # get_media swallows the error into a 4xx/5xx response; the point is the permit. + try: + await api_server.get_media("chan", 5, "errfid", request=req, digest="x") + except Exception: + pass + + # Permit released exactly once (back to baseline) — not leaked, not double-released. + assert api_server.HTTP_DOWNLOAD_SEMAPHORE._value == permits_before + + +async def test_media_semaphore_not_released_on_acquire_timeout(monkeypatch, tmp_path): + api_server._inflight.clear() + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(api_server, "verify_media_digest", lambda url, digest: True) + + permits_before = api_server.HTTP_DOWNLOAD_SEMAPHORE._value + + # Make the bounded acquire time out: the first wait_for (wrapping the acquire) + # raises TimeoutError; we close the passed coroutine to avoid a "never awaited" + # warning. The 503 short-circuits before _download_deduped, so only this one + # wait_for is hit. + real_wait_for = asyncio.wait_for + + async def fake_wait_for(coro, timeout=None): + if hasattr(coro, "close"): + coro.close() + raise asyncio.TimeoutError() + + monkeypatch.setattr(api_server.asyncio, "wait_for", fake_wait_for) + + req = SimpleNamespace(headers={}) + resp = await api_server.get_media("chan", 5, "busyfid", request=req, digest="x") + + assert resp.status_code == 503 + assert resp.headers.get("retry-after") == "30" + # A timed-out acquire never held a permit, so nothing must have been released: + # the count is UNCHANGED (an erroneous release would inflate it above baseline). + assert api_server.HTTP_DOWNLOAD_SEMAPHORE._value == permits_before + + monkeypatch.setattr(api_server.asyncio, "wait_for", real_wait_for) + + async def test_download_atomic_reraises_floodwait_and_cleans_part(monkeypatch, tmp_path): # Stage-2 review: prove a FloodWait raised by the downloader propagates OUT of # _download_atomic (its finally must NOT swallow it) AND the partial is cleaned. From 870e0a40d85ab49954f62eb80fd916f5d217bdf8 Mon Sep 17 00:00:00 2001 From: claude code agent Date: Sun, 5 Jul 2026 08:26:35 +0300 Subject: [PATCH 05/14] =?UTF-8?q?fix(stability):=20stage=203=20=E2=80=94?= =?UTF-8?q?=20serve=20via=20FileResponse,=20pure-ASGI=20logging,=20bigger?= =?UTF-8?q?=20executor?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the hand-rolled media streaming with Starlette FileResponse, drops the BaseHTTPMiddleware, and enlarges the default threadpool. 3.1 prepare_file_response now returns FileResponse (handles Range/If-Range/206/ 416/multipart, sets Accept-Ranges/ETag/Last-Modified, reads efficiently — no per-64KB to_thread hop that starved the pool). Kept: the early 404 pre-check, the MIME logic (python-magic + SQLite cache), and every stage-2 behavior — the temp_* mtime touch (now DEBOUNCED, see below), delete_after -> BackgroundTask (passed as FileResponse background=), the media_key MIME cache. Removed the manual Range parsing, file_chunk_generator, and hand-built headers; Content-Disposition is formed by FileResponse from filename= (no double-set). 206 slices are byte-identical to the old code; accepted RFC-7233 deltas documented in the tests. 3.2 RequestLoggingMiddleware rewritten as a pure-ASGI class (wraps only send to observe the status line, never buffers the body, passes non-http scopes through) — the streaming body flows untouched. 3.3 lifespan sets a larger default executor (ThreadPoolExecutor, IO_THREAD_POOL_SIZE default 32) and shuts it down on exit. Review round-1 fixes folded in: the temp_* mtime touch is DEBOUNCED (TEMP_MTIME_REFRESH_INTERVAL=300s) so FileResponse's mtime-derived ETag stays stable across a resume/seek session (an every-serve touch broke If-Range resume); starlette pinned to 0.45.3; the io executor is shut down on lifespan exit; the ASGI logger includes the query string. Tests (tests/test_stage3_fileresponse.py, 18): the Range matrix vs FileResponse with every delta documented; temp_* mtime refreshed when stale AND stable when fresh (ETag identical); delete_after background runs and removes the file; media_key MIME cache hit/miss; the ASGI middleware passes the body and logs. 213 passed (195 baseline + 18). Co-Authored-By: Claude Opus 4.8 (1M context) --- api_server.py | 190 +++++++--------- config.py | 4 + dockercompose.yml | 1 + requirements.txt | 4 + tests/mock_config.py | 1 + tests/test_stage3_fileresponse.py | 363 ++++++++++++++++++++++++++++++ 6 files changed, 451 insertions(+), 112 deletions(-) create mode 100644 tests/test_stage3_fileresponse.py diff --git a/api_server.py b/api_server.py index 7f59ef4..7fa691b 100644 --- a/api_server.py +++ b/api_server.py @@ -15,7 +15,6 @@ import re import uuid import mimetypes from typing import List, Union, Any -from urllib.parse import quote import json from datetime import datetime @@ -23,9 +22,9 @@ import time from contextlib import asynccontextmanager import random import asyncio +from concurrent.futures import ThreadPoolExecutor import uvloop asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) -from starlette.middleware.base import BaseHTTPMiddleware from starlette.background import BackgroundTask import sys @@ -51,14 +50,35 @@ magic_mime = magic.Magic(mime=True) class ZeroSizeFileError(Exception): """Custom exception for zero-size files found or downloaded.""" -class RequestLoggingMiddleware(BaseHTTPMiddleware): - async def dispatch(self, request: Request, call_next): - # Log only method and URL at debug level to avoid flooding logs on active RSS polling - logger.debug(f"Request: {request.method} {request.url}") +class RequestLoggingMiddleware: + """Pure-ASGI request logger (no BaseHTTPMiddleware). + + BaseHTTPMiddleware runs the downstream app in a separate anyio task and pumps the + response through an in-memory stream, which adds per-request overhead and interacts + badly with streaming bodies, background tasks and client cancellation. This plain + ASGI middleware only wraps `send` to observe the response status line, so it never + buffers the body — the FileResponse stream flows straight through untouched. + """ + def __init__(self, app): + self.app = app + + async def __call__(self, scope, receive, send): + if scope["type"] != "http": + await self.app(scope, receive, send) + return + # Log only method and path (with query, matching the old request.url logging) at + # debug level to avoid flooding logs on active RSS polling. + _qs = scope.get("query_string") or b"" + _path = scope["path"] + (f"?{_qs.decode('latin-1')}" if _qs else "") + logger.debug(f"Request: {scope['method']} {_path}") + + async def send_wrapper(message): + if message["type"] == "http.response.start": + logger.debug(f"Response status: {message['status']}") + await send(message) + try: - response = await call_next(request) - logger.debug(f"Response status: {response.status_code}") - return response + await self.app(scope, receive, send_wrapper) except Exception as e: logger.error(f"Request processing error: {str(e)}") raise @@ -71,6 +91,11 @@ Config = get_settings() HTTP_DOWNLOAD_SEMAPHORE = asyncio.Semaphore(3) # semaphore for live HTTP media requests BACKGROUND_DOWNLOAD_SEMAPHORE = asyncio.Semaphore(2) # semaphore for background cache worker download_queue = asyncio.Queue(maxsize=100) +# How stale a temp_* file's mtime must be before a serve refreshes it (keeps the 1h +# sweeper from deleting an actively-viewed video). Well below 1h so the file stays alive, +# but large enough that the mtime — and thus FileResponse's ETag — is stable across the +# rapid requests of one resume/seek session. +TEMP_MTIME_REFRESH_INTERVAL = 300 # seconds # In-flight download dedup registry: maps (channel, post_id, file_unique_id) to the # shared Future of an ongoing download. The FIRST request for a key runs the download in @@ -117,7 +142,13 @@ async def _supervised(factory, name: str, min_restart_interval: float = 60.0): async def lifespan(_: FastAPI): setup_logging(Config["log_level"]) - + # Enlarge the default threadpool: SQLite/python-magic/pickle/os.walk all run via + # asyncio.to_thread, and the interpreter default (min(32, cpu+4) = 5-6 on a 1-2 CPU + # container) is too small under load. Configurable via IO_THREAD_POOL_SIZE. + loop = asyncio.get_running_loop() + io_executor = ThreadPoolExecutor(max_workers=Config["io_thread_pool_size"], thread_name_prefix="io") + loop.set_default_executor(io_executor) + base_cache_dir = os.path.abspath("./data/cache") os.makedirs(base_cache_dir, exist_ok=True) # Create cache directory @@ -141,6 +172,8 @@ async def lifespan(_: FastAPI): except asyncio.CancelledError: pass await client.stop() + # Shut the io threadpool down so its threads don't linger past a reload/restart. + io_executor.shutdown(wait=False) app = FastAPI(title="Pyrogram Bridge", lifespan=lifespan) app.add_middleware(RequestLoggingMiddleware) @@ -244,16 +277,33 @@ async def delayed_delete_file(file_path: str, delay: int = 300) -> None: async def prepare_file_response(file_path: str, request: Request, delete_after: bool = False, - media_key: tuple[str, int, str] | None = None) -> StreamingResponse: - """Prepare a streaming file response with HTTP Range request support.""" + media_key: tuple[str, int, str] | None = None) -> Response: + """Serve a cached media file via Starlette's FileResponse. + + FileResponse handles Range/If-Range/206/416/multipart and sets + Accept-Ranges/ETag/Last-Modified itself, and reads the file efficiently (no per-64KB + to_thread hop that starved the threadpool). We keep: the early 404 pre-check, the MIME + logic (python-magic + SQLite type cache), the stage-2 temp_* mtime touch, and the + delete_after BackgroundTask. + """ + # `request` is unused now that FileResponse parses the Range header itself, but the + # signature is kept for call-site compatibility (and future needs). if not os.path.exists(file_path): raise HTTPException(status_code=404, detail="File not found") # Keep an actively-viewed large-video temp file alive: refresh its mtime so the 1h # sweeper (which deletes temp_* by mtime) won't remove it out from under a viewer. + # DEBOUNCED: FileResponse derives ETag/Last-Modified from mtime, so touching on EVERY + # serve would make the validators change per request and break a player's `If-Range` + # resume (it would restart with a 200 instead of a 206). We only refresh when the + # mtime is already older than TEMP_MTIME_REFRESH_INTERVAL — far below the 1h sweeper + # window, so the file still stays alive, but the ETag is stable across the rapid + # requests of a single resume/seek session. if os.path.basename(file_path).startswith("temp_"): try: - await asyncio.to_thread(os.utime, file_path, None) + age = time.time() - await asyncio.to_thread(os.path.getmtime, file_path) + if age > TEMP_MTIME_REFRESH_INTERVAL: + await asyncio.to_thread(os.utime, file_path, None) except OSError as e: logger.debug(f"Failed to refresh mtime for {file_path}: {e}") @@ -281,107 +331,23 @@ async def prepare_file_response(file_path: str, request: Request, delete_after: logger.debug(f"Determined media type for {os.path.basename(file_path)}: {media_type}") - try: - total_size = os.path.getsize(file_path) - except FileNotFoundError: - raise HTTPException(status_code=404, detail="File not found") - range_header = request.headers.get("range") - - if range_header: - # Parse Range header according to RFC 7233 - try: - range_value = range_header.strip() - if not range_value.startswith("bytes="): - raise ValueError("Only bytes ranges are supported") - range_spec = range_value[len("bytes="):] - if range_spec.startswith("-"): - # Suffix range: bytes=-N (last N bytes) - suffix_length = int(range_spec[1:]) - start = max(0, total_size - suffix_length) - end = total_size - 1 - elif range_spec.endswith("-"): - # Open-ended range: bytes=START- - start = int(range_spec[:-1]) - end = total_size - 1 - else: - # Full range: bytes=START-END - start_str, end_str = range_spec.split("-", 1) - start = int(start_str) - end = int(end_str) - except (ValueError, IndexError) as e: - logger.debug(f"Invalid Range header '{range_header}': {e}") - return Response( - status_code=416, - headers={"Content-Range": f"bytes */{total_size}"} - ) - - # Clamp end to file size - 1 (RFC 7233 allows end >= total_size) - end = min(end, total_size - 1) - - # If start is beyond file size, return 416 - if start >= total_size: - return Response( - status_code=416, - headers={"Content-Range": f"bytes */{total_size}"} - ) - - content_length = end - start + 1 - status_code = 206 - headers = { - "Content-Disposition": ( - f"inline; filename=\"{os.path.basename(file_path)}\"; " - f"filename*=UTF-8''{quote(os.path.basename(file_path))}" - ), - # Files are addressed by file_unique_id which is immutable in Telegram, - # so it is safe to cache them aggressively on the client side. - "Cache-Control": "public, max-age=86400, immutable", - "Accept-Ranges": "bytes", - "Content-Range": f"bytes {start}-{end}/{total_size}", - "Content-Length": str(content_length), - } - else: - # No Range header — serve full file with status 200 - start = 0 - end = total_size - 1 - status_code = 200 - headers = { - "Content-Disposition": ( - f"inline; filename=\"{os.path.basename(file_path)}\"; " - f"filename*=UTF-8''{quote(os.path.basename(file_path))}" - ), - # Files are addressed by file_unique_id which is immutable in Telegram, - # so it is safe to cache them aggressively on the client side. - "Cache-Control": "public, max-age=86400, immutable", - "Accept-Ranges": "bytes", - "Content-Length": str(total_size), - } - - chunk_size = 64 * 1024 # 64 KB chunks - - async def file_chunk_generator(): - """Async generator that reads file in chunks; each chunk is a separate open/seek/read/close.""" - bytes_remaining = end - start + 1 - offset = start - while bytes_remaining > 0: - to_read = min(chunk_size, bytes_remaining) - def read_at(path, off, size): - # Open, seek, read, and close within a single thread call to avoid fd leaks - with open(path, "rb") as f: - f.seek(off) - return f.read(size) - data = await asyncio.to_thread(read_at, file_path, offset, to_read) - if not data: - break - bytes_remaining -= len(data) - offset += len(data) - yield data - + # Delete the temporary file once the response has been fully sent (stage-2 delete_after). + # FileResponse runs this BackgroundTask after streaming the body. background = BackgroundTask(delayed_delete_file, file_path) if delete_after else None - return StreamingResponse( - content=file_chunk_generator(), - status_code=status_code, + + # FileResponse handles Range/If-Range/206/416/multipart and sets + # Accept-Ranges/ETag/Last-Modified itself. Do NOT hand-build Content-Disposition: + # passing filename= makes FileResponse emit `inline; filename*=UTF-8''...` on its own + # (setting it manually would double the header). + # + # Files are addressed by file_unique_id which is immutable in Telegram, so it is safe + # to cache them aggressively on the client side. + return FileResponse( + file_path, media_type=media_type, - headers=headers, + filename=os.path.basename(file_path), + content_disposition_type="inline", + headers={"Cache-Control": "public, max-age=86400, immutable"}, background=background, ) diff --git a/config.py b/config.py index 0053cb0..a8fc395 100644 --- a/config.py +++ b/config.py @@ -129,4 +129,8 @@ def get_settings() -> dict[str, Any]: "media_download_timeout_min": _parse_int_env("MEDIA_DOWNLOAD_TIMEOUT_MIN", 120), "media_download_timeout_max": _parse_int_env("MEDIA_DOWNLOAD_TIMEOUT_MAX", 1800), "media_download_min_speed": _parse_int_env("MEDIA_DOWNLOAD_MIN_SPEED", 256 * 1024), + # Size of the asyncio default ThreadPoolExecutor. SQLite/python-magic/pickle/os.walk + # all run via asyncio.to_thread; the interpreter default (min(32, cpu+4)) is only 5-6 + # on a 1-2 CPU container, which starves those under load. 32 gives ample headroom. + "io_thread_pool_size": _parse_int_env("IO_THREAD_POOL_SIZE", 32), } diff --git a/dockercompose.yml b/dockercompose.yml index 800c685..c22711a 100644 --- a/dockercompose.yml +++ b/dockercompose.yml @@ -28,6 +28,7 @@ services: # MEDIA_DOWNLOAD_TIMEOUT_MIN: 120 # Min per-download timeout, seconds — also the timeout for regular (non-large) files (default: 120) # MEDIA_DOWNLOAD_TIMEOUT_MAX: 1800 # Max per-download timeout, seconds — cap for the largest videos (default: 1800) # MEDIA_DOWNLOAD_MIN_SPEED: 262144 # Assumed floor download speed, bytes/s — large-video timeout ≈ file_size / this, clamped to [MIN,MAX] (default: 262144 = 256 KB/s) + # IO_THREAD_POOL_SIZE: 32 # Size of the asyncio default threadpool for blocking I/O (SQLite/python-magic/pickle/os.walk); raise on a busy 1-2 CPU box (default: 32) PYROGRAM_BRIDGE_URL: https://pgbridge.example.com API_PORT: 80 TOKEN: ХХХ diff --git a/requirements.txt b/requirements.txt index aafaac3..a30a961 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,7 @@ fastapi==0.115.8 +# Pinned explicitly (fastapi only requires a range): the stage-3 FileResponse Range +# behaviour and the multipart-byteranges test are tied to this Starlette version. +starlette==0.45.3 uvicorn==0.34.0 python-multipart==0.0.20 Kurigram==2.2.22 @@ -12,3 +15,4 @@ python-magic==0.4.27 bleach[css]==6.1.0 types-bleach pytest-asyncio==1.4.0 +httpx==0.28.1 diff --git a/tests/mock_config.py b/tests/mock_config.py index 6668616..d6bacfe 100644 --- a/tests/mock_config.py +++ b/tests/mock_config.py @@ -36,4 +36,5 @@ def get_settings(): "media_download_timeout_min": 120, "media_download_timeout_max": 1800, "media_download_min_speed": 256 * 1024, + "io_thread_pool_size": 32, } diff --git a/tests/test_stage3_fileresponse.py b/tests/test_stage3_fileresponse.py new file mode 100644 index 0000000..03bbc4e --- /dev/null +++ b/tests/test_stage3_fileresponse.py @@ -0,0 +1,363 @@ +# 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 3 (FileResponse + HTTP-layer cleanup) regression tests. + +Covers: +- The HTTP Range matrix asserted against the FINAL (Starlette FileResponse) behavior, + each case documented, including the RFC-7233-permitted deltas vs the old hand-rolled + streaming (garbage header -> 400 not 416; multi-range -> proper multipart 206 not 416; + 416 Content-Range now `*/size`; ETag/Last-Modified now present; ASCII filename no longer + gets a redundant filename*= form). +- Stage-2 behavior preserved through the rewrite: a served temp_* file still has its mtime + refreshed; delete_after still schedules the temp-file BackgroundTask (and it actually + runs); the media_key MIME cache is still consulted and populated. +- The pure-ASGI RequestLoggingMiddleware: a normal request still returns (body intact, not + buffered/truncated) and is logged. + +Before/after Range matrix (2048-byte file), old = hand-rolled streaming, new = FileResponse: + + request | old | new (FileResponse) + -----------------------+-----------------------------+----------------------------------- + no Range | 200 full | 200 full (+ETag/Last-Modified) + bytes=0-499 | 206 bytes 0-499/2048 | 206 bytes 0-499/2048 (same bytes) + bytes=500- | 206 bytes 500-2047/2048 | 206 bytes 500-2047/2048 (same) + bytes=-500 | 206 bytes 1548-2047/2048 | 206 bytes 1548-2047/2048 (same) + bytes=999999- (>EOF) | 416 `bytes */2048` | 416 `*/2048` (Starlette omits the + | | `bytes ` prefix; RFC-acceptable) + garbage header | 416 | 400 Bad Request (RFC 7233 lets a + | | server reject/ignore a malformed + | | Range; Starlette returns 400) + bytes=0-9,20-29 | 416 (old parser bug) | 206 multipart/byteranges (correct) + +All 206 responses that carry data return byte-for-byte identical slices old vs new, so no +real regression is hidden behind an "accepted difference". +""" +import os +import sys +import time +import logging + +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']) + +from fastapi import FastAPI, Request +from fastapi.testclient import TestClient + +import api_server + + +# --------------------------------------------------------------------------- # +# Helpers +# --------------------------------------------------------------------------- # +BODY = bytes(range(256)) * 8 # 2048 deterministic bytes +SIZE = len(BODY) + + +def _make_client(file_path, delete_after=False, media_key=None): + """Mount prepare_file_response on a tiny app so it is driven through real ASGI + (FileResponse computes Range/206/416 at send time, so it must be exercised via a client).""" + app = FastAPI() + + @app.get("/f") + async def _serve(request: Request): + return await api_server.prepare_file_response( + file_path, request=request, delete_after=delete_after, media_key=media_key + ) + + return TestClient(app) + + +@pytest.fixture +def sample_file(tmp_path): + fp = tmp_path / "myfile.bin" + fp.write_bytes(BODY) + return str(fp) + + +# --------------------------------------------------------------------------- # +# 3.1 — Range matrix against the FINAL FileResponse behavior. +# --------------------------------------------------------------------------- # +def test_no_range_returns_full_200(sample_file): + c = _make_client(sample_file) + r = c.get("/f") + assert r.status_code == 200 + assert r.content == BODY + assert r.headers["content-length"] == str(SIZE) + assert r.headers["accept-ranges"] == "bytes" + assert r.headers["cache-control"] == "public, max-age=86400, immutable" + # Content-Disposition: inline, filename set by FileResponse itself (no manual double-set). + assert r.headers["content-disposition"].startswith("inline") + assert 'filename="myfile.bin"' in r.headers["content-disposition"] + # NEW vs old: FileResponse adds validators. Old streaming had neither. + assert r.headers.get("etag") + assert r.headers.get("last-modified") + + +def test_range_prefix_0_499(sample_file): + c = _make_client(sample_file) + r = c.get("/f", headers={"Range": "bytes=0-499"}) + assert r.status_code == 206 + assert r.headers["content-range"] == f"bytes 0-499/{SIZE}" + assert r.headers["content-length"] == "500" + assert r.content == BODY[:500] # exact slice, identical to old behavior + + +def test_range_open_ended_500_to_eof(sample_file): + c = _make_client(sample_file) + r = c.get("/f", headers={"Range": "bytes=500-"}) + assert r.status_code == 206 + assert r.headers["content-range"] == f"bytes 500-{SIZE - 1}/{SIZE}" + assert r.headers["content-length"] == str(SIZE - 500) + assert r.content == BODY[500:] + + +def test_range_suffix_last_500(sample_file): + c = _make_client(sample_file) + r = c.get("/f", headers={"Range": "bytes=-500"}) + assert r.status_code == 206 + assert r.headers["content-range"] == f"bytes {SIZE - 500}-{SIZE - 1}/{SIZE}" + assert r.headers["content-length"] == "500" + assert r.content == BODY[-500:] + + +def test_range_start_past_eof_returns_416(sample_file): + c = _make_client(sample_file) + r = c.get("/f", headers={"Range": "bytes=999999-"}) + assert r.status_code == 416 + # DELTA (RFC-acceptable): Starlette's 416 Content-Range is `*/size` (it omits the + # `bytes ` unit prefix the old code emitted). Same information, permitted by RFC 7233. + assert r.headers["content-range"] == f"*/{SIZE}" + + +def test_garbage_range_header_returns_400(sample_file): + c = _make_client(sample_file) + r = c.get("/f", headers={"Range": "somethinggarbage"}) + # DELTA (RFC-acceptable): the old hand-rolled parser returned 416 on an unparseable + # header; Starlette rejects a malformed Range with 400 Bad Request. RFC 7233 permits a + # server to reject/ignore a bad Range; this is a conscious accepted difference, not a + # data regression (no bytes are served either way). + assert r.status_code == 400 + + +def test_multi_range_returns_multipart_206(sample_file): + c = _make_client(sample_file) + r = c.get("/f", headers={"Range": "bytes=0-9,20-29"}) + # DELTA (improvement): the old parser mis-split multi-ranges and returned 416. + # FileResponse serves a proper multipart/byteranges 206 containing both slices. + # Starlette 0.45.3 quirk: it advertises the multipart boundary via the `content-range` + # header (rather than content-type, which stays the file's own media_type); the body is + # a real multipart/byteranges document. We assert on the boundary marker + both slices. + assert r.status_code == 206 + assert r.headers["content-range"].startswith("multipart/byteranges") + body = r.content + assert BODY[0:10] in body + assert BODY[20:30] in body + + +# --------------------------------------------------------------------------- # +# 3.1 (stage-2 preserved) — temp_* mtime touch survives the FileResponse swap. +# --------------------------------------------------------------------------- # +def test_temp_file_mtime_refreshed_when_stale(tmp_path): + # A temp_* file whose mtime is OLDER than the refresh interval gets touched, so the + # 1h sweeper won't delete an actively-viewed video. + temp_file = tmp_path / "temp_bigvid" + temp_file.write_bytes(b"videodata") + stale = time.time() - 10000 # >> TEMP_MTIME_REFRESH_INTERVAL (300s) + os.utime(temp_file, (stale, stale)) + + c = _make_client(str(temp_file)) + r = c.get("/f") + assert r.status_code == 200 + assert r.content == b"videodata" + assert os.path.getmtime(temp_file) > time.time() - 100 # refreshed by the serve + + +def test_temp_file_mtime_stable_when_fresh(tmp_path): + # DEBOUNCE: a temp_* file touched RECENTLY (within the refresh interval) is NOT + # re-touched — so FileResponse's mtime-derived ETag stays stable across the rapid + # requests of a single resume/seek session and `If-Range` resume keeps working. + temp_file = tmp_path / "temp_freshvid" + temp_file.write_bytes(b"videodata") + recent = time.time() - 30 # << TEMP_MTIME_REFRESH_INTERVAL (300s) + os.utime(temp_file, (recent, recent)) + mtime_before = os.path.getmtime(temp_file) + + c = _make_client(str(temp_file)) + r1 = c.get("/f") + r2 = c.get("/f") + assert r1.status_code == 200 and r2.status_code == 200 + # mtime unchanged -> the ETag is identical across the two serves. + assert os.path.getmtime(temp_file) == mtime_before + assert r1.headers.get("etag") == r2.headers.get("etag") + + +def test_non_temp_file_mtime_untouched(tmp_path): + normal = tmp_path / "regularfile" + normal.write_bytes(b"data") + stale = time.time() - 10000 + os.utime(normal, (stale, stale)) + + c = _make_client(str(normal)) + c.get("/f") + assert os.path.getmtime(normal) < time.time() - 5000 # left alone + + +# --------------------------------------------------------------------------- # +# 3.1 (stage-2 preserved) — delete_after still schedules the temp-file BackgroundTask. +# --------------------------------------------------------------------------- # +async def test_delete_after_attaches_background_task(tmp_path): + fp = tmp_path / "temp_todelete" + fp.write_bytes(b"x") + + class _Req: + headers = {} + + resp = await api_server.prepare_file_response(str(fp), request=_Req(), delete_after=True) + # FileResponse must carry a non-None background that deletes exactly this file. + assert resp.background is not None + assert resp.background.func is api_server.delayed_delete_file + assert resp.background.args == (str(fp),) + + +async def test_no_delete_after_has_no_background(tmp_path): + fp = tmp_path / "keepme" + fp.write_bytes(b"x") + + class _Req: + headers = {} + + resp = await api_server.prepare_file_response(str(fp), request=_Req(), delete_after=False) + assert resp.background is None + + +def test_delete_after_background_runs_and_removes_file(tmp_path, monkeypatch): + fp = tmp_path / "temp_gone" + fp.write_bytes(BODY) + + # Patch the module-level deleter to remove immediately (real one sleeps 300s); the + # BackgroundTask picks up this reference at prepare_file_response call time. + async def _delete_now(path, delay=300): + os.remove(path) + + monkeypatch.setattr(api_server, "delayed_delete_file", _delete_now) + + c = _make_client(str(fp), delete_after=True) + r = c.get("/f") + assert r.status_code == 200 + assert r.content == BODY + # TestClient blocks until the background task has run. + assert not os.path.exists(fp) + + +# --------------------------------------------------------------------------- # +# 3.1 (stage-2 preserved) — media_key MIME cache is consulted and populated. +# --------------------------------------------------------------------------- # +def test_media_key_mime_cache_hit_used(sample_file, monkeypatch): + calls = {"get": 0, "set": 0, "magic": 0} + + def fake_get(db, ch, pid, fid): + calls["get"] += 1 + return "video/mp4" # cache HIT + + def fake_set(*a, **k): + calls["set"] += 1 + + def fake_magic(_path): + calls["magic"] += 1 + return "application/octet-stream" + + monkeypatch.setattr(api_server, "get_mime_type_sync", fake_get) + monkeypatch.setattr(api_server, "set_mime_type_sync", fake_set) + monkeypatch.setattr(api_server.magic_mime, "from_file", fake_magic) + + c = _make_client(sample_file, media_key=("chan", 42, "fid")) + r = c.get("/f") + assert r.status_code == 200 + # Cache hit: python-magic never invoked, nothing re-written to the cache, MIME applied. + assert calls["get"] == 1 + assert calls["magic"] == 0 + assert calls["set"] == 0 + assert r.headers["content-type"].startswith("video/mp4") + + +def test_media_key_mime_cache_miss_populates(sample_file, monkeypatch): + written = {} + + def fake_get(db, ch, pid, fid): + return None # cache MISS + + def fake_set(db, ch, pid, fid, mime): + written["mime"] = mime + + def fake_magic(_path): + return "image/png" + + monkeypatch.setattr(api_server, "get_mime_type_sync", fake_get) + monkeypatch.setattr(api_server, "set_mime_type_sync", fake_set) + monkeypatch.setattr(api_server.magic_mime, "from_file", fake_magic) + + c = _make_client(sample_file, media_key=("chan", 42, "fid")) + r = c.get("/f") + assert r.status_code == 200 + # Miss -> python-magic result is both applied and persisted for next time. + assert written.get("mime") == "image/png" + assert r.headers["content-type"].startswith("image/png") + + +# --------------------------------------------------------------------------- # +# 3.1 — 404 pre-check kept. +# --------------------------------------------------------------------------- # +def test_missing_file_returns_404(tmp_path): + c = _make_client(str(tmp_path / "does_not_exist")) + r = c.get("/f") + assert r.status_code == 404 + + +# --------------------------------------------------------------------------- # +# 3.2 — pure-ASGI RequestLoggingMiddleware: request returns intact and is logged. +# --------------------------------------------------------------------------- # +def test_asgi_logging_middleware_passes_body_and_logs(tmp_path, caplog): + fp = tmp_path / "streamed.bin" + fp.write_bytes(BODY) + + app = FastAPI() + app.add_middleware(api_server.RequestLoggingMiddleware) + + @app.get("/f") + async def _serve(request: Request): + return await api_server.prepare_file_response(str(fp), request=request) + + client = TestClient(app) + with caplog.at_level(logging.DEBUG, logger="api_server"): + r = client.get("/f") + + # Body flows straight through the middleware — not buffered/truncated. + assert r.status_code == 200 + assert r.content == BODY + messages = " ".join(rec.getMessage() for rec in caplog.records) + assert "Request: GET /f" in messages + assert "Response status: 200" in messages + + +def test_asgi_logging_middleware_range_still_works(tmp_path): + fp = tmp_path / "streamed.bin" + fp.write_bytes(BODY) + + app = FastAPI() + app.add_middleware(api_server.RequestLoggingMiddleware) + + @app.get("/f") + async def _serve(request: Request): + return await api_server.prepare_file_response(str(fp), request=request) + + client = TestClient(app) + r = client.get("/f", headers={"Range": "bytes=0-99"}) + # 206 streaming still works through the pure-ASGI middleware (send not buffered). + assert r.status_code == 206 + assert r.content == BODY[:100] From 1e24fd535470c1b9e85e79f205c0104cdd79e1ba Mon Sep 17 00:00:00 2001 From: claude code agent Date: Sun, 5 Jul 2026 08:42:54 +0300 Subject: [PATCH 06/14] fix(stability): pass FileResponse a stat_result (404 not 500 on race) + doc/annotation fixes (review round 1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit #1 [medium] FileResponse with stat_result=None re-stats at send-time and raises a RuntimeError (-> 500 + stacktrace, escaping get_media's try/except) if the file was swept between the early exists() check and the send. The old streaming code caught FileNotFoundError on getsize -> 404. Now prepare_file_response takes ONE authoritative os.stat (in try/except FileNotFoundError -> 404) after the temp_* touch and passes it as FileResponse(stat_result=...): restores the 404 semantics, collapses the double stat into one, and makes the ETag reflect the observed mtime. The remaining narrow deleted-between-our-stat-and-open window pre-existed. #2/#3 [doc] Corrected the Content-Disposition comment (FileResponse emits filename="x" for ASCII, filename*=UTF-8''x only for non-ASCII, via setdefault — a manual header would override, not double) and the ETag-stability comment (stable within a 300s window; a longer view costs at most one safe 200 If-Range restart per interval). #4 [simplification] Removed the now-dead StreamingResponse import and corrected get_media's return annotation to -> Response (nothing returns StreamingResponse after the FileResponse migration). #5 [test] Added a test for the middleware's non-http (lifespan/websocket) branch — it only runs in a real deploy, never through TestClient — asserting it delegates to the inner app and logs nothing. 214 passed (195 baseline + 19). Co-Authored-By: Claude Opus 4.8 (1M context) --- api_server.py | 41 ++++++++++++++++++++----------- tests/test_stage3_fileresponse.py | 20 +++++++++++++++ 2 files changed, 47 insertions(+), 14 deletions(-) diff --git a/api_server.py b/api_server.py index 7fa691b..6201bcf 100644 --- a/api_server.py +++ b/api_server.py @@ -33,7 +33,7 @@ from pyrogram import errors from pyrogram.types import Message from pyrogram.enums import MessageMediaType from fastapi import FastAPI, HTTPException, Response, Request -from fastapi.responses import HTMLResponse, FileResponse, StreamingResponse +from fastapi.responses import HTMLResponse, FileResponse from telegram_client import TelegramClient from config import get_settings, setup_logging from rss_generator import generate_channel_rss, generate_channel_html @@ -93,8 +93,9 @@ BACKGROUND_DOWNLOAD_SEMAPHORE = asyncio.Semaphore(2) # semaphore for background download_queue = asyncio.Queue(maxsize=100) # How stale a temp_* file's mtime must be before a serve refreshes it (keeps the 1h # sweeper from deleting an actively-viewed video). Well below 1h so the file stays alive, -# but large enough that the mtime — and thus FileResponse's ETag — is stable across the -# rapid requests of one resume/seek session. +# but large enough that the mtime — and thus FileResponse's ETag — is stable within any +# such window; a view running longer than one interval costs at most one safe 200 +# If-Range restart per interval (a full re-fetch, never corruption). TEMP_MTIME_REFRESH_INTERVAL = 300 # seconds # In-flight download dedup registry: maps (channel, post_id, file_unique_id) to the @@ -288,17 +289,14 @@ async def prepare_file_response(file_path: str, request: Request, delete_after: """ # `request` is unused now that FileResponse parses the Range header itself, but the # signature is kept for call-site compatibility (and future needs). - if not os.path.exists(file_path): - raise HTTPException(status_code=404, detail="File not found") # Keep an actively-viewed large-video temp file alive: refresh its mtime so the 1h # sweeper (which deletes temp_* by mtime) won't remove it out from under a viewer. # DEBOUNCED: FileResponse derives ETag/Last-Modified from mtime, so touching on EVERY - # serve would make the validators change per request and break a player's `If-Range` - # resume (it would restart with a 200 instead of a 206). We only refresh when the - # mtime is already older than TEMP_MTIME_REFRESH_INTERVAL — far below the 1h sweeper - # window, so the file still stays alive, but the ETag is stable across the rapid - # requests of a single resume/seek session. + # serve would change the validators per request. We only refresh when the mtime is + # already older than TEMP_MTIME_REFRESH_INTERVAL — far below the 1h sweeper window, so + # the file stays alive, and the ETag is stable within any such window (a view running + # longer than one interval costs at most one safe 200 If-Range restart per interval). if os.path.basename(file_path).startswith("temp_"): try: age = time.time() - await asyncio.to_thread(os.path.getmtime, file_path) @@ -307,6 +305,18 @@ async def prepare_file_response(file_path: str, request: Request, delete_after: except OSError as e: logger.debug(f"Failed to refresh mtime for {file_path}: {e}") + # Take ONE authoritative stat and hand it to FileResponse as stat_result. This both + # (a) preserves the 404 semantics: FileResponse with stat_result=None re-stats at + # send-time and raises a RuntimeError (-> 500, escaping this handler's try/except) if + # the file was swept between here and the send; and (b) makes the ETag/Last-Modified + # reflect exactly the mtime observed after the optional touch above. The remaining + # narrow window (deleted between this stat and FileResponse's own open) truncates the + # body — that pre-existed the FileResponse migration and is not handled here. + try: + stat_result = await asyncio.to_thread(os.stat, file_path) + except FileNotFoundError: + raise HTTPException(status_code=404, detail="File not found") + media_type: str | None = None if media_key is not None: @@ -336,9 +346,11 @@ async def prepare_file_response(file_path: str, request: Request, delete_after: background = BackgroundTask(delayed_delete_file, file_path) if delete_after else None # FileResponse handles Range/If-Range/206/416/multipart and sets - # Accept-Ranges/ETag/Last-Modified itself. Do NOT hand-build Content-Disposition: - # passing filename= makes FileResponse emit `inline; filename*=UTF-8''...` on its own - # (setting it manually would double the header). + # Accept-Ranges/ETag/Last-Modified itself (from the stat_result we pass). Do NOT + # hand-build Content-Disposition: FileResponse forms it from filename= — + # `inline; filename="x"` for an ASCII name, adding `filename*=UTF-8''x` only for a + # non-ASCII name. It uses setdefault, so a manual header would OVERRIDE it, not double + # it; letting FileResponse own it keeps the RFC 5987 encoding correct. # # Files are addressed by file_unique_id which is immutable in Telegram, so it is safe # to cache them aggressively on the client side. @@ -347,6 +359,7 @@ async def prepare_file_response(file_path: str, request: Request, delete_after: media_type=media_type, filename=os.path.basename(file_path), content_disposition_type="inline", + stat_result=stat_result, headers={"Cache-Control": "public, max-age=86400, immutable"}, background=background, ) @@ -1035,7 +1048,7 @@ async def health_check(request: Request, token: str | None = None) -> Response: @app.get("/media/{channel}/{post_id}/{file_unique_id}/{digest}", response_model=None) @app.get("/media/{channel}/{post_id}/{file_unique_id}", response_model=None) -async def get_media(channel: str, post_id: int, file_unique_id: str, request: Request, digest: str | None = None) -> StreamingResponse|Response: +async def get_media(channel: str, post_id: int, file_unique_id: str, request: Request, digest: str | None = None) -> Response: try: url = f"{channel}/{post_id}/{file_unique_id}" if not verify_media_digest(url, digest): diff --git a/tests/test_stage3_fileresponse.py b/tests/test_stage3_fileresponse.py index 03bbc4e..9dc054c 100644 --- a/tests/test_stage3_fileresponse.py +++ b/tests/test_stage3_fileresponse.py @@ -361,3 +361,23 @@ def test_asgi_logging_middleware_range_still_works(tmp_path): # 206 streaming still works through the pure-ASGI middleware (send not buffered). assert r.status_code == 206 assert r.content == BODY[:100] + + +async def test_asgi_logging_middleware_passes_non_http_scope_through(caplog): + # The `scope["type"] != "http"` branch (lifespan/websocket) only runs in a real + # deploy, never through TestClient's plain request path — so pin it directly: + # a non-http scope must delegate to the inner app untouched and log nothing. + called = {} + + async def inner(scope, receive, send): + called["scope_type"] = scope["type"] + + mw = api_server.RequestLoggingMiddleware(inner) + with caplog.at_level(logging.DEBUG, logger="api_server"): + await mw({"type": "lifespan"}, None, None) + + assert called["scope_type"] == "lifespan" # delegated to the inner app + # Nothing request/response-ish was logged for a non-http scope. + messages = " ".join(rec.getMessage() for rec in caplog.records) + assert "Request:" not in messages + assert "Response status:" not in messages From 9f9091f6fba29dd2a4e9a740a97be5e8abdca684 Mon Sep 17 00:00:00 2001 From: claude code agent Date: Sun, 5 Jul 2026 09:16:14 +0300 Subject: [PATCH 07/14] =?UTF-8?q?fix(stability):=20stage=204=20=E2=80=94?= =?UTF-8?q?=20event-loop=20hygiene=20for=20feed=20generation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Feed generation blocked the loop (deepcopy of hundreds of Messages + per-message bleach x4 + str(message) per post). This moves the CPU work off the loop and cuts it down. 4.1 raw_message is lazy: process_message(include_raw=False) omits str(message) entirely; only /json and debug-HTML compute it. 4.2 (done BEFORE 4.3) _save_media_file_ids no longer touches asyncio/DB — it appends to the per-request PostParser._pending_media_ids, and the caller flushes once via to_thread(upsert_media_file_ids_bulk_sync) (a new executemany fn in file_io.py). Removed _persist_pending_count + the create_task machinery. This had to precede 4.3 or get_running_loop() inside the render thread would raise and silently kill media-id persistence. 4.3 The render pipeline (_create_time_based_media_groups, _create_messages_groups, _trim_messages_groups, _render_messages_groups) is now plain-sync and runs in ONE asyncio.to_thread(_render_pipeline) from both feed generators; deepcopy moved into the thread. The render path is verified free of asyncio primitives. 4.4 Sanitize is now ONE pass per output boundary (removed the 4 internal per-fragment bleach passes): RSS/HTML feeds sanitize once at the whole-feed pass; /html and /json sanitize body+footer once in process_message; the debug 
 raw dump is html.escape'd. Media embeds in body, reactions in footer, so
    both are covered by the boundary pass. XSS tests (click"
+
+
+class _Str(str):
+    """Minimal stand-in for Pyrogram's Str: .html returns the raw string unchanged,
+    so a malicious payload reaches the pre-sanitization html body just like real text
+    carrying entities would."""
+    @property
+    def html(self):
+        return str(self)
+
+
+def make_message(mid, text=None, media=None, photo_uid=None, username="testchan",
+                 date=None):
+    m = SimpleNamespace()
+    m.id = mid
+    m.date = date or datetime(2024, 1, 1, 12, 0, mid % 60, tzinfo=timezone.utc)
+    m.text = _Str(text) if text is not None else None
+    m.caption = None
+    m.media = media
+    m.web_page = None
+    m.poll = None
+    m.service = None
+    m.forward_origin = None
+    m.reply_to_message = None
+    m.reply_to_message_id = None
+    m.sender_chat = None
+    m.from_user = None
+    m.reactions = None
+    m.views = 100
+    m.media_group_id = None
+    m.show_caption_above_media = False
+    m.chat = SimpleNamespace(id=-1001234567890, username=username)
+    # media sub-objects default to None
+    for attr in ("photo", "video", "document", "audio", "voice",
+                 "video_note", "animation", "sticker"):
+        setattr(m, attr, None)
+    if media == MessageMediaType.PHOTO and photo_uid:
+        m.photo = SimpleNamespace(file_unique_id=photo_uid)
+    return m
+
+
+def _co_names(func):
+    return set(func.__code__.co_names)
+
+
+# ---------------------------------------------------------------------------
+# 4.3 — render functions are plain sync and run off the main thread
+# ---------------------------------------------------------------------------
+
+def test_render_functions_are_sync():
+    for fn in (_create_time_based_media_groups, _create_messages_groups,
+               _trim_messages_groups, _render_messages_groups, _render_pipeline):
+        assert not asyncio.iscoroutinefunction(fn), f"{fn.__name__} must be a plain sync function"
+
+
+@pytest.mark.asyncio
+async def test_pipeline_runs_in_worker_thread(monkeypatch):
+    main_ident = threading.get_ident()
+    seen = {}
+
+    real_render = rss_module._render_messages_groups
+
+    def spy(*args, **kwargs):
+        seen["ident"] = threading.get_ident()
+        return real_render(*args, **kwargs)
+
+    monkeypatch.setattr(rss_module, "_render_messages_groups", spy)
+
+    async def fake_get_chat(client, channel):
+        return SimpleNamespace(title="Test", username="testchan", id=-1001234567890)
+
+    async def fake_get_history(client, channel, limit=20):
+        return [make_message(i, text=f"post {i}") for i in range(5)]
+
+    monkeypatch.setattr("tg_cache.cached_get_chat", fake_get_chat, raising=False)
+    monkeypatch.setattr("tg_cache.cached_get_chat_history", fake_get_history, raising=False)
+
+    await generate_channel_rss("testchan", client=SimpleNamespace(), limit=5)
+    assert "ident" in seen
+    assert seen["ident"] != main_ident, "render pipeline must run in a worker thread, not the loop thread"
+
+
+def test_deepcopy_of_pickled_message_does_not_crash():
+    from pyrogram.types import Message, Chat
+    from pyrogram.enums import ChatType
+    m = Message(id=7, date=datetime.now(timezone.utc), text="hello",
+                chat=Chat(id=-1001, type=ChatType.CHANNEL, username="testchan"))
+    roundtripped = pickle.loads(pickle.dumps(m))   # mimics the pickle cache
+    clone = copy.deepcopy(roundtripped)            # what _create_time_based_media_groups does
+    assert clone.id == 7
+    assert clone.chat.username == "testchan"
+
+
+# ---------------------------------------------------------------------------
+# 4.2 — no asyncio in the render path; bulk upsert after render
+# ---------------------------------------------------------------------------
+
+def test_render_path_has_no_asyncio_side_effects():
+    banned = {"create_task", "get_running_loop", "to_thread", "ensure_future"}
+    funcs = [
+        _render_pipeline, _create_time_based_media_groups, _create_messages_groups,
+        _trim_messages_groups, _render_messages_groups,
+        PostParser.process_message, PostParser._generate_html_body,
+        PostParser._generate_html_media, PostParser.generate_html_footer,
+        PostParser._reactions_views_links, PostParser._save_media_file_ids,
+        PostParser._sanitize_html,
+    ]
+    for fn in funcs:
+        offenders = _co_names(fn) & banned
+        assert not offenders, f"{fn.__qualname__} references forbidden asyncio names: {offenders}"
+
+
+@pytest.mark.asyncio
+async def test_media_ids_persisted_via_bulk_upsert(monkeypatch):
+    calls = []
+
+    def fake_bulk(db_path, entries):
+        calls.append(list(entries))
+
+    monkeypatch.setattr(pp_module, "upsert_media_file_ids_bulk_sync", fake_bulk)
+
+    async def fake_get_chat(client, channel):
+        return SimpleNamespace(title="Test", username="testchan", id=-1001234567890)
+
+    async def fake_get_history(client, channel, limit=20):
+        return [
+            make_message(1, text="just text"),
+            make_message(2, media=MessageMediaType.PHOTO, photo_uid="uid_abc"),
+        ]
+
+    monkeypatch.setattr("tg_cache.cached_get_chat", fake_get_chat, raising=False)
+    monkeypatch.setattr("tg_cache.cached_get_chat_history", fake_get_history, raising=False)
+
+    await generate_channel_rss("testchan", client=SimpleNamespace(), limit=10)
+
+    assert len(calls) == 1, "bulk upsert must be called exactly once after render"
+    entries = calls[0]
+    assert len(entries) == 1, "only the photo message contributes a media id"
+    channel, post_id, fid, _ts = entries[0]
+    assert (channel, post_id, fid) == ("testchan", 2, "uid_abc")
+
+
+@pytest.mark.asyncio
+async def test_save_media_file_ids_only_appends(monkeypatch):
+    # Even with a running loop, _save_media_file_ids must not create tasks — just append.
+    parser = PostParser(SimpleNamespace())
+    monkeypatch.setattr(pp_module, "upsert_media_file_ids_bulk_sync",
+                        lambda *a, **k: (_ for _ in ()).throw(AssertionError("must not be called directly")))
+    msg = make_message(3, media=MessageMediaType.PHOTO, photo_uid="uid_x")
+    parser._save_media_file_ids(msg)
+    assert parser._pending_media_ids == [("testchan", 3, "uid_x", parser._pending_media_ids[0][3])]
+
+
+# ---------------------------------------------------------------------------
+# 4.1 — raw_message laziness
+# ---------------------------------------------------------------------------
+
+def test_raw_message_lazy_for_feed():
+    parser = PostParser(SimpleNamespace())
+    result = parser.process_message(make_message(10, text="hi"), include_raw=False, sanitize=False)
+    assert "raw_message" not in result
+
+
+def test_raw_message_present_for_json_and_debug():
+    parser = PostParser(SimpleNamespace())
+    result = parser.process_message(make_message(11, text="hi"), include_raw=True)
+    assert "raw_message" in result
+    assert isinstance(result["raw_message"], str)
+
+
+@pytest.mark.asyncio
+async def test_100_message_feed_generates(monkeypatch):
+    async def fake_get_chat(client, channel):
+        return SimpleNamespace(title="Test", username="testchan", id=-1001234567890)
+
+    async def fake_get_history(client, channel, limit=20):
+        return [make_message(i, text=f"post number {i}") for i in range(1, 101)]
+
+    monkeypatch.setattr("tg_cache.cached_get_chat", fake_get_chat, raising=False)
+    monkeypatch.setattr("tg_cache.cached_get_chat_history", fake_get_history, raising=False)
+
+    rss = await generate_channel_rss("testchan", client=SimpleNamespace(), limit=100)
+    assert rss.count("") == 100
+    assert "post number 50" in rss
+
+
+# ---------------------------------------------------------------------------
+# 4.4 — XSS: payload stripped in all four outputs
+# ---------------------------------------------------------------------------
+
+def _assert_clean(html_str, where):
+    assert "