diff --git a/api_server.py b/api_server.py index 14234c2..95fb23c 100644 --- a/api_server.py +++ b/api_server.py @@ -92,6 +92,36 @@ 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) +# --- Failing-download backoff (negative cache) -------------------------------- +# A media file whose download keeps failing (hang/timeout/not-found) must not be +# retried on every 60s cache sweep, nor keep occupying a scarce Pyrogram +# transmission slot. We remember recent failures per (channel, post_id, +# file_unique_id) and skip / fast-reject re-attempts until an exponentially +# growing backoff elapses. Cleared on the first successful download. All access is +# from the single event-loop thread, so a plain dict needs no lock. +_DOWNLOAD_BACKOFF_BASE = 60.0 # seconds; backoff after the first failure +_DOWNLOAD_BACKOFF_MAX = 3600.0 # seconds; cap on the backoff +# key -> (consecutive_failures, retry_not_before_monotonic) +_download_failures: dict[tuple[str, int, str], tuple[int, float]] = {} + +def _download_backoff_remaining(key: tuple[str, int, str]) -> float: + """Seconds until `key` may be retried; 0.0 if allowed now (or never failed).""" + entry = _download_failures.get(key) + if entry is None: + return 0.0 + return max(0.0, entry[1] - time.monotonic()) + +def _record_download_failure(key: tuple[str, int, str]) -> None: + """Register a failed download and (re)arm an exponential backoff for `key`.""" + fails = _download_failures.get(key, (0, 0.0))[0] + 1 + backoff = min(_DOWNLOAD_BACKOFF_MAX, _DOWNLOAD_BACKOFF_BASE * (2 ** (fails - 1))) + _download_failures[key] = (fails, time.monotonic() + backoff) + logger.warning(f"download_backoff_armed: {key[0]}/{key[1]}/{key[2]} failed {fails}x, next retry in {backoff:.0f}s") + +def _clear_download_failure(key: tuple[str, int, str]) -> None: + """Forget any recorded failure for `key` after a successful download.""" + if _download_failures.pop(key, None) is not None: + logger.info(f"download_backoff_cleared: {key[0]}/{key[1]}/{key[2]} recovered") # 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 within any @@ -530,9 +560,13 @@ async def _download_deduped(channel: Union[str, int], post_id: int, file_unique_ async def _runner(): try: result = await download_media_file(channel, post_id, file_unique_id) + _clear_download_failure(key) if not fut.done(): fut.set_result(result) except BaseException as e: # noqa: BLE001 — must forward ANY failure to waiters + # Arm backoff for real failures only, never for a shutdown-time cancel. + if isinstance(e, Exception): + _record_download_failure(key) if not fut.done(): fut.set_exception(e) finally: @@ -784,7 +818,13 @@ async def download_new_files(media_files: list, cache_dir: str) -> None: if not all([channel, post_id, file_unique_id]): logger.error(f"Invalid file data: {file_data}") continue - + + # Skip files that recently kept failing — do not re-queue them on every 60s + # sweep (that retry-storm is what kept the download slots jammed). The backoff + # expires on its own; a successful download elsewhere clears it. + if _download_backoff_remaining((str(channel), int(post_id), file_unique_id)) > 0: + continue + channel_dir = os.path.join(cache_dir, str(channel)) post_dir = os.path.join(channel_dir, str(post_id)) os.makedirs(post_dir, exist_ok=True) @@ -822,17 +862,21 @@ async def background_download_worker(): # successful get(). Cancellation propagates cleanly here (nothing to unbalance). item = await download_queue.get() channel, post_id, file_unique_id = item + bg_key = (str(channel), int(post_id), file_unique_id) 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) + _clear_download_failure(bg_key) 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. + # otherwise the worker would hammer Telegram while under a flood wait. A flood + # wait is a global throttle, not a per-file fault, so it does NOT arm backoff. 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: + _record_download_failure(bg_key) logger.error(f"Background download error for {channel}/{post_id}/{file_unique_id}: {e}") finally: download_queue.task_done() @@ -1213,6 +1257,16 @@ async def get_media(channel: str, post_id: int, file_unique_id: str, request: Re return await prepare_file_response(cache_path, request=request, media_key=(str(channel), post_id, file_unique_id)) + # A file that recently kept failing is in backoff: fast-reject instead of + # occupying a scarce download slot (and a Pyrogram transmission permit) on a + # request that will very likely hang again. Cached files already returned above, + # so this only guards the live-download path. + backoff_remaining = _download_backoff_remaining((str(channel), post_id, file_unique_id)) + if backoff_remaining > 0: + logger.info(f"media_backoff_skip: {channel}/{post_id}/{file_unique_id} in backoff {backoff_remaining:.0f}s") + return Response(status_code=503, content="Media temporarily unavailable, retry later", + headers={"Retry-After": str(int(backoff_remaining) + 1)}) + _sem_wait_start = _time.monotonic() # 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 diff --git a/config.py b/config.py index 3cc8763..8e73643 100644 --- a/config.py +++ b/config.py @@ -144,4 +144,15 @@ def get_settings() -> dict[str, Any]: # 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), + # Max concurrent Telegram file transmissions (Pyrogram get_file/save_file + # semaphores). Kurigram's default is 1, which lets a single hung download block + # ALL media downloads process-wide; 3 aligns with HTTP_DOWNLOAD_SEMAPHORE. This is + # only a blast-radius limiter — the real cure for a zombie media connection is the + # download-timeout-triggered restart below. + "tg_max_concurrent_transmissions": _parse_int_env("TG_MAX_CONCURRENT_TRANSMISSIONS", 3), + # After this many CONSECUTIVE media-download timeouts, force an in-process client + # restart to rebuild the (zombie) media-DC connection. Any successful download + # resets the streak, so this only fires on a genuine death loop, not on the odd + # slow large-video timeout. The watchdog cannot catch this — it probes the main DC. + "media_timeout_restart_threshold": _parse_int_env("MEDIA_TIMEOUT_RESTART_THRESHOLD", 5), } diff --git a/dockercompose.yml b/dockercompose.yml index 4fa8bda..67ec3d5 100644 --- a/dockercompose.yml +++ b/dockercompose.yml @@ -29,6 +29,8 @@ services: # 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) + # TG_MAX_CONCURRENT_TRANSMISSIONS: 3 # Max concurrent Telegram file transmissions (Pyrogram get_file semaphore). Kurigram default is 1, so one hung download blocks ALL media (default: 3) + # MEDIA_TIMEOUT_RESTART_THRESHOLD: 5 # Consecutive media-download timeouts before an in-process restart rebuilds the zombie media-DC connection (the main-DC watchdog can't see this) (default: 5) PYROGRAM_BRIDGE_URL: https://pgbridge.example.com API_PORT: 80 TOKEN: ХХХ diff --git a/telegram_client.py b/telegram_client.py index 6d4f547..e18f2ea 100644 --- a/telegram_client.py +++ b/telegram_client.py @@ -33,10 +33,18 @@ class TelegramClient: api_hash=settings["tg_api_hash"], workdir=settings["session_path"], proxy=settings["proxy"], # MTProto proxy config, None if not set + max_concurrent_transmissions=settings["tg_max_concurrent_transmissions"], ) self.max_disconnects = settings["tg_disconnect_flap_limit"] # Max disconnects within the flap window before restart self._shutting_down = False # Guard to prevent re-triggering restart during shutdown self._restarting = False # Guard: an intentional in-process restart is in progress + # Consecutive media-download timeouts. A zombie media-DC connection makes every + # download time out while the main-DC watchdog stays green; when this streak + # reaches the threshold we reuse _restart_client() to rebuild the connection. Any + # successful download resets it (see note_download_ok / note_download_timeout). + self._download_timeout_streak = 0 + self.media_timeout_restart_threshold = settings["media_timeout_restart_threshold"] + self._media_recovery_task = None # strong ref to a scheduled recovery restart task self._disconnect_times = [] # Monotonic timestamps of recent disconnects (sliding window) self.disconnect_window = settings["tg_disconnect_flap_window"] # Seconds; window for flap detection self._watchdog_task = None @@ -270,6 +278,39 @@ class TelegramClient: return None return time.monotonic() - self._wd_last_ok_monotonic + def note_download_ok(self) -> None: + """Reset the media-download timeout streak after any successful download.""" + if self._download_timeout_streak: + logger.info(f"media_download: recovered, resetting timeout streak (was {self._download_timeout_streak})") + self._download_timeout_streak = 0 + + def note_download_timeout(self) -> None: + """Count a media-download timeout; force a connection-rebuilding restart on a streak. + + A zombie media-DC connection makes EVERY download time out while the main-DC + watchdog probe (get_me) stays green, so this is the only signal that can trigger + recovery. The restart runs as a detached task so it never blocks the download path; + the streak is reset immediately so we schedule at most one restart per streak. + """ + self._download_timeout_streak += 1 + logger.warning( + f"media_download_timeout: streak {self._download_timeout_streak}/{self.media_timeout_restart_threshold}" + ) + if self._download_timeout_streak < self.media_timeout_restart_threshold: + return + self._download_timeout_streak = 0 + if self._restarting or self._shutting_down: + return + if self._media_recovery_task is not None and not self._media_recovery_task.done(): + return # a recovery restart is already scheduled/running + logger.critical( + f"media_download: {self.media_timeout_restart_threshold} consecutive download timeouts — " + f"scheduling media-connection restart (main-DC watchdog cannot see this)" + ) + self._media_recovery_task = asyncio.create_task( + self._restart_client(reason="media download timeout streak") + ) + async def safe_get_messages(self, channel_id, post_id, max_retries=2): """Wrapper with retry logic for auth errors""" for attempt in range(max_retries): @@ -289,13 +330,22 @@ class TelegramClient: """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).""" + with file size (see api_server._media_download_timeout). A timeout cancels the + underlying download (freeing the Pyrogram transmission slot); a streak of timeouts + escalates to a connection-rebuilding restart via note_download_timeout().""" for attempt in range(max_retries): try: - return await asyncio.wait_for( + result = await asyncio.wait_for( self.client.download_media(file_id, file_name=file_name), timeout=timeout ) + self.note_download_ok() + return result + except asyncio.TimeoutError: + # Hung download: the wait_for above already cancelled it and released the + # get_file semaphore. Count it toward the media-connection restart streak. + self.note_download_timeout() + raise except Exception as e: if isinstance(e, KeyError) and attempt < max_retries - 1: logger.warning(f"Download auth error on attempt {attempt + 1}, retrying...") diff --git a/tests/mock_config.py b/tests/mock_config.py index cfed846..beeca8a 100644 --- a/tests/mock_config.py +++ b/tests/mock_config.py @@ -38,4 +38,6 @@ def get_settings(): "media_download_timeout_max": 1800, "media_download_min_speed": 256 * 1024, "io_thread_pool_size": 32, + "tg_max_concurrent_transmissions": 3, + "media_timeout_restart_threshold": 5, } diff --git a/tests/test_media_selfheal.py b/tests/test_media_selfheal.py new file mode 100644 index 0000000..f9fd34d --- /dev/null +++ b/tests/test_media_selfheal.py @@ -0,0 +1,100 @@ +# flake8: noqa +# pylint: disable=protected-access, missing-function-docstring, redefined-outer-name, line-too-long +# pylance: disable=reportMissingImports, reportMissingModuleSource +""" +Regression tests for the media self-healing fix (post 'static refactor' outage): + +Root cause recap: Kurigram serializes downloads through a single get_file slot; a +zombie media-DC connection makes every download time out while the main-DC watchdog +stays green, so downloads jam forever. The fix adds (a) a negative-cache backoff so a +repeatedly-failing file is not hammered, and (b) a consecutive-timeout streak that +reuses the existing in-process restart to rebuild the media connection. +""" +import asyncio + +import pytest + +import api_server +from telegram_client import TelegramClient + + +# --------------------------------------------------------------------------- # +# Negative-cache backoff helpers +# --------------------------------------------------------------------------- # +def test_backoff_arms_and_clears(): + key = ("selfheal_chan", 1, "fid_a") + api_server._download_failures.pop(key, None) + assert api_server._download_backoff_remaining(key) == 0.0 # never failed -> allowed + api_server._record_download_failure(key) + assert api_server._download_backoff_remaining(key) > 0.0 # armed -> blocked + api_server._clear_download_failure(key) + assert api_server._download_backoff_remaining(key) == 0.0 # recovered -> allowed + + +def test_backoff_failure_counter_increments(): + key = ("selfheal_chan", 2, "fid_b") + api_server._download_failures.pop(key, None) + api_server._record_download_failure(key) + first_fails = api_server._download_failures[key][0] + api_server._record_download_failure(key) + second_fails = api_server._download_failures[key][0] + assert first_fails == 1 + assert second_fails == 2 # consecutive-failure counter grows -> longer backoff + api_server._clear_download_failure(key) + + +# --------------------------------------------------------------------------- # +# Download-timeout streak -> single media-connection restart +# --------------------------------------------------------------------------- # +@pytest.mark.asyncio +async def test_timeout_streak_triggers_single_restart(monkeypatch): + c = TelegramClient() + calls = [] + + async def fake_restart(reason: str = "unspecified"): + calls.append(reason) + + monkeypatch.setattr(c, "_restart_client", fake_restart) + threshold = c.media_timeout_restart_threshold + + # One short of the threshold: no restart scheduled yet. + for _ in range(threshold - 1): + c.note_download_timeout() + assert c._media_recovery_task is None + assert c._download_timeout_streak == threshold - 1 + + # The threshold-th timeout schedules exactly one restart and resets the streak. + c.note_download_timeout() + assert c._download_timeout_streak == 0 + assert c._media_recovery_task is not None + await c._media_recovery_task + assert calls == ["media download timeout streak"] + + +@pytest.mark.asyncio +async def test_success_resets_streak(): + c = TelegramClient() + c.note_download_timeout() + c.note_download_timeout() + assert c._download_timeout_streak == 2 + c.note_download_ok() + assert c._download_timeout_streak == 0 + + +@pytest.mark.asyncio +async def test_no_restart_while_already_restarting(monkeypatch): + c = TelegramClient() + calls = [] + + async def fake_restart(reason: str = "unspecified"): + calls.append(reason) + + monkeypatch.setattr(c, "_restart_client", fake_restart) + c._restarting = True # a restart is already underway + + for _ in range(c.media_timeout_restart_threshold): + c.note_download_timeout() + + # Streak reached the threshold but no new restart is scheduled during a restart. + assert c._media_recovery_task is None + assert calls == []