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