api_server.py

@@ -17,6 +17,7 @@ import mimetypes
 from typing import List, Union, Any
 
 import json
+from collections import OrderedDict
 from datetime import datetime
 import time
 from contextlib import asynccontextmanager
@@ -133,9 +134,21 @@ _queued_media: set[tuple[str, int, str]] = set()
 # 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]] = {}
+# Cap on the backoff. Kept deliberately short (10 min, not 1h): the media self-heal
+# restart recovers in minutes, and a verified restart clears this cache outright (see
+# _clear_all_download_failures), so a system that just healed must not hold a stale,
+# long backoff that keeps fast-503'ing a now-downloadable file.
+_DOWNLOAD_BACKOFF_MAX = 600.0     # seconds; cap on the backoff
+# LRU cap on the negative cache. Permanently-404 / deleted files would otherwise leave
+# eternal entries and slowly leak memory on a long-uptime process. We keep at most this
+# many most-recently-failed keys and evict the oldest. Eviction is harmless: a dropped
+# key is simply treated as "never failed" again — at worst one extra retry attempt, which
+# re-arms its backoff on failure. 10k keys is a tiny footprint yet far above any realistic
+# concurrent-failure working set.
+_DOWNLOAD_FAILURES_MAX = 10000
+# key -> (consecutive_failures, retry_not_before_monotonic). OrderedDict so we can evict
+# in least-recently-updated order once the LRU cap is exceeded.
+_download_failures: "OrderedDict[tuple[str, int, str], tuple[int, float]]" = OrderedDict()
 
 def _download_backoff_remaining(key: tuple[str, int, str]) -> float:
     """Seconds until `key` may be retried; 0.0 if allowed now (or never failed)."""
@@ -149,12 +162,31 @@ def _record_download_failure(key: tuple[str, int, str]) -> None:
     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)
+    _download_failures.move_to_end(key)  # mark as most-recently-updated for LRU eviction
+    # Bound memory (see _DOWNLOAD_FAILURES_MAX): evict the oldest entries beyond the cap.
+    while len(_download_failures) > _DOWNLOAD_FAILURES_MAX:
+        _download_failures.popitem(last=False)
     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")
+
+def _clear_all_download_failures() -> None:
+    """Drop the entire download negative cache. Registered with telegram_client and invoked
+    ONLY after a VERIFIED self-heal restart (verify_get_me OK). _restart_client() rebuilds
+    the WHOLE client — all DCs, including the media DC — so every per-file backoff is stale
+    and a previously-failing file may now download. Without this a file that reached the
+    backoff cap would keep fast-503'ing for up to _DOWNLOAD_BACKOFF_MAX after recovery, and
+    nothing would retry it (the sweeper skips backed-off keys, get_media fast-503s them), so
+    it could not self-heal until the backoff expired — defeating the self-heal's purpose.
+    The retry-storm risk if the restart did not actually help is bounded: each re-download is
+    timeout-bounded and simply re-arms its backoff."""
+    n = len(_download_failures)
+    _download_failures.clear()
+    if n:
+        logger.warning(f"download_backoff_cleared_all: dropped {n} entries after verified self-heal restart")
 # 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
@@ -291,6 +323,12 @@ async def lifespan(_: FastAPI):
         except OSError as e:
             logger.warning(f"legacy_media_file_ids_remove_error: {e}")
 
+    # Wire the self-heal restart -> negative-cache clear hook. telegram_client owns the
+    # restart; the negative cache lives here. Registering a plain callback (rather than
+    # importing api_server from telegram_client) keeps the dependency one-way and avoids a
+    # circular import. The hook fires ONLY on a verified restart — see _restart_client.
+    client.set_restart_callback(_clear_all_download_failures)
+
     await client.start()
     # 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.

telegram_client.py

@@ -45,6 +45,7 @@ class TelegramClient:
         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._on_restart_verified = None   # optional callback fired after a VERIFIED restart (see set_restart_callback)
         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
@@ -188,6 +189,17 @@ class TelegramClient:
         except Exception as e:
             logger.critical(f"watchdog: loop crashed unexpectedly ({type(e).__name__}: {e}); liveness protection is now DISABLED until next start")
 
+    def set_restart_callback(self, callback) -> None:
+        """Register a callback invoked after a VERIFIED in-process restart (verify_get_me OK).
+
+        Used by api_server to clear its download negative cache: a completed restart rebuilds
+        the whole client (all DCs, including the media DC), so any per-file download backoff is
+        stale and a previously-failing file may now download. Kept as a plain hook so
+        telegram_client never has to import api_server (which would be a circular import).
+        The callback is synchronous, best-effort, and never fires on a failed/aborted restart.
+        """
+        self._on_restart_verified = callback
+
     async def _restart_client(self, reason: str = "unspecified"):
         """Recover the client without killing the process when possible.
 
@@ -214,10 +226,14 @@ class TelegramClient:
                 await asyncio.wait_for(self.client.start(), timeout=self.watchdog_restart_timeout)
             duration = time.monotonic() - restart_started
             self._wd_restart_count += 1
-            # Verification probe to prove the network layer is actually back (diagnostic only).
+            # Verification probe to prove the network layer is actually back. It also gates the
+            # verified-restart callback below: we only treat the restart as a real recovery
+            # (and clear the download negative cache) when this probe actually succeeds.
+            verify_ok = False
             try:
                 me = await asyncio.wait_for(self.client.get_me(), timeout=self.watchdog_timeout)
                 self._wd_last_ok_monotonic = time.monotonic()
+                verify_ok = True
                 verify = f", verify_get_me ok (me_id={getattr(me, 'id', None)})"
             except Exception as ve:
                 verify = f", verify_get_me FAILED ({type(ve).__name__}: {ve})"
@@ -226,6 +242,16 @@ class TelegramClient:
                 f"(is_connected={self.client.is_connected}{verify}, total in-process restarts={self._wd_restart_count})"
             )
             self._disconnect_times.clear()
+            # On a VERIFIED restart the whole client (all DCs, incl. the media DC) is
+            # re-established, so any per-file download backoff is stale — fire the registered
+            # hook (api_server clears its negative cache) so recovered media load immediately
+            # instead of fast-503'ing until the backoff expires. Only on verify success, never
+            # on a failed/aborted restart. Best-effort: a callback error must not abort recovery.
+            if verify_ok and self._on_restart_verified is not None:
+                try:
+                    self._on_restart_verified()
+                except Exception as cbe:
+                    logger.warning(f"recovery: restart-verified callback raised {type(cbe).__name__}: {cbe}")
             # Re-arm the watchdog in case it had previously crashed (self-healing).
             self._start_watchdog()
         except asyncio.CancelledError:
@@ -347,6 +373,12 @@ class TelegramClient:
                 self.note_download_timeout()
                 raise
             except Exception as e:
+                # INTENTIONAL: the restart streak counts CONSECUTIVE TIMEOUTS only. Non-timeout
+                # download errors (RPCError, FILE_REFERENCE_EXPIRED, etc.) reach here and call
+                # neither note_download_ok nor note_download_timeout, so they are streak-neutral:
+                # they neither advance nor reset it. This is deliberate — a zombie media-DC
+                # manifests as timeouts, not RPC errors, so only timeouts should escalate to a
+                # connection-rebuilding restart; a burst of unrelated RPC errors must not.
                 if isinstance(e, KeyError) and attempt < max_retries - 1:
                     logger.warning(f"Download auth error on attempt {attempt + 1}, retrying...")
                     await asyncio.sleep(5)

tests/test_media_selfheal.py

@@ -98,3 +98,242 @@ async def test_no_restart_while_already_restarting(monkeypatch):
     # Streak reached the threshold but no new restart is scheduled during a restart.
     assert c._media_recovery_task is None
     assert calls == []
+
+
+# --------------------------------------------------------------------------- #
+# Integration: jam -> auto-recovery -> negative cache cleared (Fix 1)
+# --------------------------------------------------------------------------- #
+@pytest.mark.asyncio
+async def test_hung_download_times_out_and_frees_slot(monkeypatch):
+    """A hung download must be cancelled by asyncio.wait_for, freeing its transmission slot
+    for the next download (and counting toward the restart streak)."""
+    c = TelegramClient()
+
+    # A single-permit semaphore stands in for Pyrogram's get_file transmission slot: the mock
+    # holds it for the whole (never-completing) download and releases it in the async-with
+    # __aexit__, which runs when wait_for cancels the coroutine.
+    slot = asyncio.Semaphore(1)
+    entered = asyncio.Event()
+
+    async def hung_download(file_id, file_name=None):
+        async with slot:
+            entered.set()
+            await asyncio.sleep(3600)  # never completes within the test's timeout
+
+    monkeypatch.setattr(c.client, "download_media", hung_download)
+
+    with pytest.raises(asyncio.TimeoutError):
+        await c.safe_download_media("fid", "/tmp/x", max_retries=1, timeout=0.05)
+
+    # (a) the download actually started and its slot was released by the cancellation, so a
+    # fresh acquire for the "next" download succeeds immediately.
+    assert entered.is_set()
+    await asyncio.wait_for(slot.acquire(), timeout=1.0)
+    slot.release()
+    # The timeout was counted toward the restart streak.
+    assert c._download_timeout_streak == 1
+
+
+@pytest.mark.asyncio
+async def test_streak_escalates_once_and_verified_restart_clears_cache(monkeypatch):
+    """Full self-heal chain: a timeout streak reaching the threshold escalates to EXACTLY ONE
+    restart, and the verified restart CLEARS the download negative cache (Fix 1)."""
+    c = TelegramClient()
+
+    # Pre-arm the negative cache with a persistently-failing file.
+    key = ("selfheal_chan", 7, "fid_jam")
+    api_server._download_failures.pop(key, None)
+    api_server._record_download_failure(key)
+    assert api_server._download_backoff_remaining(key) > 0.0
+
+    # Wire the REAL self-heal hook and drive a REAL _restart_client with only the network
+    # layer mocked, so the whole chain (streak -> one restart -> verify ok -> cache clear) runs.
+    c.set_restart_callback(api_server._clear_all_download_failures)
+    monkeypatch.setattr(c, "_start_watchdog", lambda: None)
+    monkeypatch.setattr(c.client, "is_connected", True)
+
+    restart_calls = []
+
+    async def fake_client_restart():
+        restart_calls.append(True)
+
+    async def fake_get_me():
+        return type("Me", (), {"id": 42})()
+
+    monkeypatch.setattr(c.client, "restart", fake_client_restart)
+    monkeypatch.setattr(c.client, "get_me", fake_get_me)
+
+    threshold = c.media_timeout_restart_threshold
+    for _ in range(threshold - 1):
+        c.note_download_timeout()
+    assert c._media_recovery_task is None  # one short of threshold: nothing scheduled yet
+
+    c.note_download_timeout()  # threshold-th timeout schedules exactly one restart
+    assert c._media_recovery_task is not None
+    await c._media_recovery_task
+
+    # (b) escalated to EXACTLY ONE underlying client.restart().
+    assert restart_calls == [True]
+    # (c) the verified restart cleared the negative cache -> the file loads again immediately.
+    assert api_server._download_backoff_remaining(key) == 0.0
+    assert len(api_server._download_failures) == 0
+
+
+@pytest.mark.asyncio
+async def test_failed_verify_does_not_clear_cache(monkeypatch):
+    """The cache-clear hook must fire ONLY on a verified restart: if verify_get_me fails, the
+    backoff must survive (a still-broken media DC should not drop the protective backoff)."""
+    c = TelegramClient()
+    key = ("selfheal_chan", 8, "fid_still_broken")
+    api_server._download_failures.pop(key, None)
+    api_server._record_download_failure(key)
+
+    c.set_restart_callback(api_server._clear_all_download_failures)
+    monkeypatch.setattr(c, "_start_watchdog", lambda: None)
+    monkeypatch.setattr(c.client, "is_connected", True)
+
+    async def fake_client_restart():
+        return None
+
+    async def failing_get_me():
+        raise RuntimeError("media DC still dead")
+
+    monkeypatch.setattr(c.client, "restart", fake_client_restart)
+    monkeypatch.setattr(c.client, "get_me", failing_get_me)
+
+    await c._restart_client(reason="test failed verify")
+
+    # verify_get_me failed -> callback NOT fired -> backoff still armed.
+    assert api_server._download_backoff_remaining(key) > 0.0
+    api_server._clear_download_failure(key)
+
+
+# --------------------------------------------------------------------------- #
+# Cross-path negative-cache key consistency
+# --------------------------------------------------------------------------- #
+def test_download_key_consistent_across_paths():
+    """The negative-cache key must be byte-identical across every producer/consumer path so a
+    failure recorded on one path is seen by the guard on another. Reproduce each path's key
+    expression for the same logical file and assert they coincide, then verify end-to-end that
+    a worker-recorded failure is visible to the get_media backoff check."""
+    channel_from_db = "durov"   # background_download_worker / download_new_files: str(channel)
+    post_id_from_db = 123       # ...                                              int(post_id)
+    fid = "AgADfid"
+
+    # get_media receives the raw (possibly differently-cased) URL channel, canonicalizes it,
+    # and uses the int path param post_id.
+    raw_url_channel = "Durov"
+    fs_channel = api_server.canonical_channel_key(raw_url_channel)
+    get_media_key = (fs_channel, 123, fid)
+
+    worker_key = (str(channel_from_db), int(post_id_from_db), fid)      # background_download_worker
+    new_files_key = (str(channel_from_db), int(post_id_from_db), fid)   # download_new_files
+    deduped_key = (str(fs_channel), 123, fid)                          # _download_deduped(fs_channel, post_id, fid)
+
+    assert worker_key == new_files_key == get_media_key == deduped_key
+
+    # End-to-end: a failure recorded by the worker path is seen by the get_media guard.
+    api_server._download_failures.pop(worker_key, None)
+    api_server._record_download_failure(worker_key)
+    try:
+        assert api_server._download_backoff_remaining(get_media_key) > 0.0
+    finally:
+        api_server._clear_download_failure(worker_key)
+
+
+# --------------------------------------------------------------------------- #
+# Fix 2: negative cache is LRU-bounded
+# --------------------------------------------------------------------------- #
+def test_download_failures_lru_bounded(monkeypatch):
+    """The negative cache must not grow unbounded (permanently-404 files leak entries).
+    Once over the cap the oldest entry is evicted."""
+    monkeypatch.setattr(api_server, "_DOWNLOAD_FAILURES_MAX", 3)
+    api_server._download_failures.clear()
+    try:
+        for i in range(5):
+            api_server._record_download_failure(("chan", i, "fid"))
+        assert len(api_server._download_failures) == 3
+        # Oldest two (i=0,1) evicted; newest three retained.
+        remaining = list(api_server._download_failures.keys())
+        assert remaining == [("chan", 2, "fid"), ("chan", 3, "fid"), ("chan", 4, "fid")]
+    finally:
+        api_server._download_failures.clear()
+
+
+# --------------------------------------------------------------------------- #
+# Backoff is capped at _DOWNLOAD_BACKOFF_MAX (min() cap at api_server.py:163)
+# --------------------------------------------------------------------------- #
+def test_backoff_capped_at_max():
+    """The exponential backoff grows as BASE * 2**(fails-1), but must never exceed
+    _DOWNLOAD_BACKOFF_MAX. With BASE=60s and MAX=600s the raw exponential overtakes the cap
+    at the 5th failure (60*2**4 = 960s > 600s), so after that many consecutive failures the
+    effective backoff must stay pinned at the cap. This test FAILS if the min(MAX, ...) cap
+    is removed (the raw exponential would then blow past 600s)."""
+    key = ("selfheal_chan", 99, "fid_cap")
+    api_server._download_failures.pop(key, None)
+
+    # Record well past the point where the raw exponential exceeds the cap: 8 failures ->
+    # raw 60*2**7 = 7680s, an order of magnitude above the 600s cap.
+    for _ in range(8):
+        api_server._record_download_failure(key)
+
+    assert api_server._download_failures[key][0] == 8  # counter really climbed that high
+    remaining = api_server._download_backoff_remaining(key)
+    # The cap holds: remaining is bounded by MAX (with a tiny slack for monotonic drift since
+    # retry_not_before was stamped). Without the min() cap this would be ~7680s.
+    assert remaining <= api_server._DOWNLOAD_BACKOFF_MAX
+    assert remaining > api_server._DOWNLOAD_BACKOFF_MAX - 5  # and it IS pinned near the cap
+
+    api_server._clear_download_failure(key)
+
+
+# --------------------------------------------------------------------------- #
+# Restart-verified callback is best-effort (try/except at telegram_client.py:250-254)
+# --------------------------------------------------------------------------- #
+@pytest.mark.asyncio
+async def test_restart_callback_error_does_not_abort(monkeypatch, caplog):
+    """A registered restart-verified callback that RAISES must not abort recovery: the error is
+    swallowed + logged, and the restart still completes normally (watchdog re-armed, no SIGTERM
+    fallback, _restarting reset). This test FAILS if the try/except around the callback is
+    removed — the exception would then propagate to the outer handler, skip the watchdog re-arm
+    and trigger the process-restart fallback instead."""
+    c = TelegramClient()
+
+    def exploding_callback():
+        raise RuntimeError("boom in verified-restart callback")
+
+    c.set_restart_callback(exploding_callback)
+
+    # Drive a VERIFIED restart: connected client, restart + get_me both succeed so verify_ok is
+    # True and the callback fires.
+    monkeypatch.setattr(c.client, "is_connected", True)
+
+    async def fake_client_restart():
+        return None
+
+    async def fake_get_me():
+        return type("Me", (), {"id": 42})()
+
+    monkeypatch.setattr(c.client, "restart", fake_client_restart)
+    monkeypatch.setattr(c.client, "get_me", fake_get_me)
+
+    # Observe the success path (watchdog re-arm) vs the failure path (SIGTERM fallback) without
+    # actually killing the test process.
+    watchdog_rearmed = []
+    monkeypatch.setattr(c, "_start_watchdog", lambda: watchdog_rearmed.append(True))
+    sigterm_calls = []
+    monkeypatch.setattr(c, "_restart_app", lambda: sigterm_calls.append(True))
+
+    with caplog.at_level("WARNING"):
+        await c._restart_client(reason="test callback raises")
+
+    # (a) Recovery completed on the SUCCESS path despite the callback raising: watchdog re-armed,
+    # no process-restart fallback, restart flag cleared.
+    assert watchdog_rearmed == [True]
+    assert sigterm_calls == []
+    assert c._restarting is False
+    # (b) The callback error was logged (not silently dropped).
+    assert any(
+        "callback raised" in r.getMessage() and "RuntimeError" in r.getMessage()
+        for r in caplog.records
+    )