Merge pull request 'fix(media): follow-up self-heal — чистка негативкэша на рестарте + LRU-кап + тесты' (#48) from fix/43-selfheal-followup into main
Reviewed-on: #48
This commit was merged in pull request #48.
This commit is contained in:
+41
-3
@@ -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.
|
||||
|
||||
+33
-1
@@ -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)
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user