fix(stability): stage 2 — atomic media downloads, in-flight dedup, FloodWait->429, bounded semaphore
Fixes flaky static media serving: partial files served as 'ready', truncated
stubs living for an hour, and app-wide hangs under a saturated download path.
2.1 _download_atomic(file_id, final_path, timeout): downloads to a unique
{final}.part.{hex}, validates size>0, and publishes solely via os.rename
(atomic on POSIX); a finally ALWAYS removes the partial (timeout/cancel/
zero-size/race-loser). Every downloader call routes through it, so a file at a
FINAL name ({fid} / temp_{fid}) is GUARANTEED complete (grep-proven: the only
safe_download_media caller passes a .part. path). Big-video timeout scales with
size (min/max/min-speed knobs, documented in dockercompose.yml). The sweeper
regex now matches both .part. and legacy .tmp. stubs.
2.2 In-flight dedup registry: the first request for a key runs the download in a
DETACHED task sharing a Future; the task's finally sets the Future AND pops the
key. Waiters await asyncio.shield(fut) so a client disconnect / waiter timeout
cancels only the waiter, never the download — no hung waiters, no stuck key, a
failed download frees the key for retry.
2.3 FloodWait->429: handler before except RPCError (FloodWait subclasses it),
Retry-After = min(value + rand(1,30), 300); propagates from the detached task
through the Future.
2.4 Serving a temp_* file touches its mtime so the 1h sweeper can't delete a
video out from under a viewer.
2.5 The HTTP download semaphore acquire is bounded (wait_for 30 -> 503 +
Retry-After); the permit is released only if the acquire succeeded. The
request-scoped-permit trade-off (a disconnect can transiently exceed the
download count) is documented inline for stage-7 prod observation.
Tests (tests/test_stage2_static.py, 13): atomic publish/clean on every exit incl.
FloodWait-through-the-finally; concurrent big-video serves no partial; dedup runs
one download, a cancelled waiter doesn't hang others, a failed download frees the
key; FloodWait -> 429; mtime touch; sweeper cleans .part./.tmp./stale temp_ but
keeps fresh files. 193 passed (180 baseline + 13).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
claude code agent
+163
-42
@@ -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
|
||||
|
||||
@@ -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),
|
||||
}
|
||||
|
||||
@@ -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: ХХХ
|
||||
|
||||
+6
-3
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user