vvzvlad/pyrogram-bridge
1
0
Fork 0
Code Issues 6 Pull Requests Actions 8 Packages Projects Releases Wiki Activity

Merge pull request 'fix(stability): стадия 2 — атомарные скачивания медиа, in-flight дедуп, FloodWait->429, ограниченный семафор' (#10) from fix/stage-2-static into fix/stage-1-hangs

Browse Source
This commit was merged in pull request #10.
This commit is contained in:
agent_vscode
2026-07-05 16:58:44 +03:00
parent 03d1de2954 c5e50354ba
commit 9f893969ca
16 changed files with 2843 additions and 250 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api_server.py
+372 -177
View File
@@ -15,7 +15,6 @@ import re
import uuid
import mimetypes
from typing import List, Union, Any
from urllib.parse import quote
import json
from datetime import datetime
@@ -23,9 +22,9 @@ import time
from contextlib import asynccontextmanager
import random
import asyncio
from concurrent.futures import ThreadPoolExecutor
import uvloop
asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.background import BackgroundTask
import sys
@@ -34,14 +33,15 @@ from pyrogram import errors
from pyrogram.types import Message
from pyrogram.enums import MessageMediaType
from fastapi import FastAPI, HTTPException, Response, Request
from fastapi.responses import HTMLResponse, FileResponse, StreamingResponse
from fastapi.responses import HTMLResponse, FileResponse, JSONResponse
from telegram_client import TelegramClient
from config import get_settings, setup_logging
from rss_generator import generate_channel_rss, generate_channel_html
from post_parser import PostParser
from url_signer import verify_media_digest, generate_media_digest
from file_io import (DB_PATH, init_db_sync, get_all_media_file_ids_sync,
update_media_file_access_sync, remove_media_file_ids_sync,
update_media_file_access_sync, update_media_file_access_bulk_sync,
remove_media_file_ids_sync,
get_mime_type_sync, set_mime_type_sync)
# Global python-magic instance for MIME type detection
@@ -51,14 +51,35 @@ magic_mime = magic.Magic(mime=True)
class ZeroSizeFileError(Exception):
"""Custom exception for zero-size files found or downloaded."""
class RequestLoggingMiddleware(BaseHTTPMiddleware):
async def dispatch(self, request: Request, call_next):
# Log only method and URL at debug level to avoid flooding logs on active RSS polling
logger.debug(f"Request: {request.method} {request.url}")
class RequestLoggingMiddleware:
"""Pure-ASGI request logger (no BaseHTTPMiddleware).
BaseHTTPMiddleware runs the downstream app in a separate anyio task and pumps the
response through an in-memory stream, which adds per-request overhead and interacts
badly with streaming bodies, background tasks and client cancellation. This plain
ASGI middleware only wraps `send` to observe the response status line, so it never
buffers the body — the FileResponse stream flows straight through untouched.
"""
def __init__(self, app):
self.app = app
async def __call__(self, scope, receive, send):
if scope["type"] != "http":
await self.app(scope, receive, send)
return
# Log only method and path (with query, matching the old request.url logging) at
# debug level to avoid flooding logs on active RSS polling.
_qs = scope.get("query_string") or b""
_path = scope["path"] + (f"?{_qs.decode('latin-1')}" if _qs else "")
logger.debug(f"Request: {scope['method']} {_path}")
async def send_wrapper(message):
if message["type"] == "http.response.start":
logger.debug(f"Response status: {message['status']}")
await send(message)
try:
response = await call_next(request)
logger.debug(f"Response status: {response.status_code}")
return response
await self.app(scope, receive, send_wrapper)
except Exception as e:
logger.error(f"Request processing error: {str(e)}")
raise
@@ -71,6 +92,22 @@ Config = get_settings()
HTTP_DOWNLOAD_SEMAPHORE = asyncio.Semaphore(3) # semaphore for live HTTP media requests
BACKGROUND_DOWNLOAD_SEMAPHORE = asyncio.Semaphore(2) # semaphore for background cache worker
download_queue = asyncio.Queue(maxsize=100)
# 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
# such window; a view running longer than one interval costs at most one safe 200
# If-Range restart per interval (a full re-fetch, never corruption).
TEMP_MTIME_REFRESH_INTERVAL = 300 # seconds
# 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.
@@ -103,11 +140,70 @@ async def _supervised(factory, name: str, min_restart_interval: float = 60.0):
await asyncio.sleep(min_restart_interval - elapsed)
# Access-time write accumulator. A /media cache hit used to touch SQLite on the hot path
# (a threadpool hop + connect + UPDATE per request), which starves the threadpool under
# active RSS polling. Instead a cache hit just records the timestamp here — a dict write
# on the single-threaded event loop is cheap and atomic — and a periodic background task
# flushes the whole batch to SQLite in one executemany. Keys use str(channel) to stay
# consistent with the string form written at insert time and to not lean on SQLite's
# implicit column-affinity coercion (the channel column is TEXT, so a bound int would be
# coerced and still match — but we key the accumulator by the same type we store, rather
# than depend on that).
ACCESS_FLUSH_INTERVAL = 60 # seconds between access-time flushes
_access_updates: dict[tuple[str, int, str], float] = {}
async def _flush_access_updates() -> None:
"""Flush the accumulated access timestamps to SQLite in one bulk UPDATE.
Snapshot-then-clear atomically on the loop: capture the current dict reference and
replace the module global with a fresh empty dict in ONE synchronous step (before any
await), so cache-hit writes arriving DURING the flush land in the new dict and are not
lost. The bulk UPDATE runs off-loop via asyncio.to_thread. An empty batch is a no-op.
"""
global _access_updates
if not _access_updates:
return
pending = _access_updates
_access_updates = {}
entries = [(channel, post_id, file_unique_id, added)
for (channel, post_id, file_unique_id), added in pending.items()]
try:
await asyncio.to_thread(update_media_file_access_bulk_sync, DB_PATH, entries)
except Exception:
# Bulk write failed: re-queue this batch so the access-times are not lost (a lost
# timestamp would eventually evict a still-used file from the 20-day cache). Use
# setdefault so any FRESHER write accumulated during the flush is never overwritten
# by our stale snapshot. Runs on the loop with no await before the mutation, so this
# is race-free. Re-raise so the flush loop logs it.
for key, added in pending.items():
_access_updates.setdefault(key, added)
raise
async def _access_flush_loop() -> None:
"""Periodically flush the access-time accumulator (runs under _supervised)."""
while True:
await asyncio.sleep(ACCESS_FLUSH_INTERVAL)
try:
await _flush_access_updates()
except Exception as e:
# Log and keep looping: a transient SQLite error must not drop the batch's
# successors. (_supervised still restarts us if this ever raises out.)
logger.error(f"access_flush_error: {e}")
@asynccontextmanager
async def lifespan(_: FastAPI):
setup_logging(Config["log_level"])
# Enlarge the default threadpool: SQLite/python-magic/pickle/os.walk all run via
# asyncio.to_thread, and the interpreter default (min(32, cpu+4) = 5-6 on a 1-2 CPU
# container) is too small under load. Configurable via IO_THREAD_POOL_SIZE.
loop = asyncio.get_running_loop()
io_executor = ThreadPoolExecutor(max_workers=Config["io_thread_pool_size"], thread_name_prefix="io")
loop.set_default_executor(io_executor)
base_cache_dir = os.path.abspath("./data/cache")
os.makedirs(base_cache_dir, exist_ok=True) # Create cache directory
@@ -119,9 +215,11 @@ async def lifespan(_: FastAPI):
# CRITICAL and restarted, so a crash can no longer silently stop cache sweeping or downloads.
background_task = asyncio.create_task(_supervised(cache_media_files, "cache_media_files"))
worker_task = asyncio.create_task(_supervised(background_download_worker, "background_download_worker"))
access_flush_task = asyncio.create_task(_supervised(_access_flush_loop, "access_flush_loop"))
yield
background_task.cancel() # Cleanup
worker_task.cancel()
access_flush_task.cancel()
try:
await background_task
except asyncio.CancelledError:
@@ -130,7 +228,20 @@ async def lifespan(_: FastAPI):
await worker_task
except asyncio.CancelledError:
pass
try:
await access_flush_task
except asyncio.CancelledError:
pass
# Final flush AFTER the loop task is cancelled (no race with a loop-driven flush) and
# BEFORE the threadpool is shut down (to_thread still has its executor), so the last
# <=ACCESS_FLUSH_INTERVAL seconds of access-times are persisted on shutdown.
try:
await _flush_access_updates()
except Exception as e:
logger.error(f"access_flush_shutdown_error: {e}")
await client.stop()
# Shut the io threadpool down so its threads don't linger past a reload/restart.
io_executor.shutdown(wait=False)
app = FastAPI(title="Pyrogram Bridge", lifespan=lifespan)
app.add_middleware(RequestLoggingMiddleware)
@@ -234,9 +345,43 @@ async def delayed_delete_file(file_path: str, delay: int = 300) -> None:
async def prepare_file_response(file_path: str, request: Request, delete_after: bool = False,
media_key: tuple[str, int, str] | None = None) -> StreamingResponse:
"""Prepare a streaming file response with HTTP Range request support."""
if not os.path.exists(file_path):
media_key: tuple[str, int, str] | None = None) -> Response:
"""Serve a cached media file via Starlette's FileResponse.
FileResponse handles Range/If-Range/206/416/multipart and sets
Accept-Ranges/ETag/Last-Modified itself, and reads the file efficiently (no per-64KB
to_thread hop that starved the threadpool). We keep: the early 404 pre-check, the MIME
logic (python-magic + SQLite type cache), the stage-2 temp_* mtime touch, and the
delete_after BackgroundTask.
"""
# `request` is unused now that FileResponse parses the Range header itself, but the
# signature is kept for call-site compatibility (and future needs).
# 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.
# DEBOUNCED: FileResponse derives ETag/Last-Modified from mtime, so touching on EVERY
# serve would change the validators per request. We only refresh when the mtime is
# already older than TEMP_MTIME_REFRESH_INTERVAL — far below the 1h sweeper window, so
# the file stays alive, and the ETag is stable within any such window (a view running
# longer than one interval costs at most one safe 200 If-Range restart per interval).
if os.path.basename(file_path).startswith("temp_"):
try:
age = time.time() - await asyncio.to_thread(os.path.getmtime, file_path)
if age > TEMP_MTIME_REFRESH_INTERVAL:
await asyncio.to_thread(os.utime, file_path, None)
except OSError as e:
logger.debug(f"Failed to refresh mtime for {file_path}: {e}")
# Take ONE authoritative stat and hand it to FileResponse as stat_result. This both
# (a) preserves the 404 semantics: FileResponse with stat_result=None re-stats at
# send-time and raises a RuntimeError (-> 500, escaping this handler's try/except) if
# the file was swept between here and the send; and (b) makes the ETag/Last-Modified
# reflect exactly the mtime observed after the optional touch above. The remaining
# narrow window (deleted between this stat and FileResponse's own open) truncates the
# body — that pre-existed the FileResponse migration and is not handled here.
try:
stat_result = await asyncio.to_thread(os.stat, file_path)
except FileNotFoundError:
raise HTTPException(status_code=404, detail="File not found")
media_type: str | None = None
@@ -263,110 +408,120 @@ async def prepare_file_response(file_path: str, request: Request, delete_after:
logger.debug(f"Determined media type for {os.path.basename(file_path)}: {media_type}")
try:
total_size = os.path.getsize(file_path)
except FileNotFoundError:
raise HTTPException(status_code=404, detail="File not found")
range_header = request.headers.get("range")
if range_header:
# Parse Range header according to RFC 7233
try:
range_value = range_header.strip()
if not range_value.startswith("bytes="):
raise ValueError("Only bytes ranges are supported")
range_spec = range_value[len("bytes="):]
if range_spec.startswith("-"):
# Suffix range: bytes=-N (last N bytes)
suffix_length = int(range_spec[1:])
start = max(0, total_size - suffix_length)
end = total_size - 1
elif range_spec.endswith("-"):
# Open-ended range: bytes=START-
start = int(range_spec[:-1])
end = total_size - 1
else:
# Full range: bytes=START-END
start_str, end_str = range_spec.split("-", 1)
start = int(start_str)
end = int(end_str)
except (ValueError, IndexError) as e:
logger.debug(f"Invalid Range header '{range_header}': {e}")
return Response(
status_code=416,
headers={"Content-Range": f"bytes */{total_size}"}
)
# Clamp end to file size - 1 (RFC 7233 allows end >= total_size)
end = min(end, total_size - 1)
# If start is beyond file size, return 416
if start >= total_size:
return Response(
status_code=416,
headers={"Content-Range": f"bytes */{total_size}"}
)
content_length = end - start + 1
status_code = 206
headers = {
"Content-Disposition": (
f"inline; filename=\"{os.path.basename(file_path)}\"; "
f"filename*=UTF-8''{quote(os.path.basename(file_path))}"
),
# Files are addressed by file_unique_id which is immutable in Telegram,
# so it is safe to cache them aggressively on the client side.
"Cache-Control": "public, max-age=86400, immutable",
"Accept-Ranges": "bytes",
"Content-Range": f"bytes {start}-{end}/{total_size}",
"Content-Length": str(content_length),
}
else:
# No Range header — serve full file with status 200
start = 0
end = total_size - 1
status_code = 200
headers = {
"Content-Disposition": (
f"inline; filename=\"{os.path.basename(file_path)}\"; "
f"filename*=UTF-8''{quote(os.path.basename(file_path))}"
),
# Files are addressed by file_unique_id which is immutable in Telegram,
# so it is safe to cache them aggressively on the client side.
"Cache-Control": "public, max-age=86400, immutable",
"Accept-Ranges": "bytes",
"Content-Length": str(total_size),
}
chunk_size = 64 * 1024 # 64 KB chunks
async def file_chunk_generator():
"""Async generator that reads file in chunks; each chunk is a separate open/seek/read/close."""
bytes_remaining = end - start + 1
offset = start
while bytes_remaining > 0:
to_read = min(chunk_size, bytes_remaining)
def read_at(path, off, size):
# Open, seek, read, and close within a single thread call to avoid fd leaks
with open(path, "rb") as f:
f.seek(off)
return f.read(size)
data = await asyncio.to_thread(read_at, file_path, offset, to_read)
if not data:
break
bytes_remaining -= len(data)
offset += len(data)
yield data
# Delete the temporary file once the response has been fully sent (stage-2 delete_after).
# FileResponse runs this BackgroundTask after streaming the body.
background = BackgroundTask(delayed_delete_file, file_path) if delete_after else None
return StreamingResponse(
content=file_chunk_generator(),
status_code=status_code,
# FileResponse handles Range/If-Range/206/416/multipart and sets
# Accept-Ranges/ETag/Last-Modified itself (from the stat_result we pass). Do NOT
# hand-build Content-Disposition: FileResponse forms it from filename= —
# `inline; filename="x"` for an ASCII name, adding `filename*=UTF-8''x` only for a
# non-ASCII name. It uses setdefault, so a manual header would OVERRIDE it, not double
# it; letting FileResponse own it keeps the RFC 5987 encoding correct.
#
# Files are addressed by file_unique_id which is immutable in Telegram, so it is safe
# to cache them aggressively on the client side.
return FileResponse(
file_path,
media_type=media_type,
headers=headers,
filename=os.path.basename(file_path),
content_disposition_type="inline",
stat_result=stat_result,
headers={"Cache-Control": "public, max-age=86400, immutable"},
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 sets the Future's result/exception (on success/failure) and its finally ALWAYS
pops the key — so both happen before the task ends: 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 +565,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")
@@ -441,16 +605,11 @@ async def download_media_file(channel: Union[str, int], post_id: int, file_uniqu
logger.error(f"cleanup_error: Failed to remove zero-size cached file {cache_path}: {e}")
# Do not raise error here, proceed to download below
else:
# File exists and is not zero size, update access timestamp and return
# File exists and is not zero size, record access timestamp and return.
# Record into the accumulator instead of touching SQLite on the hot path; the
# background flush persists it. Key channel as str(channel) — see _access_updates.
logger.info(f"Found cached media file: {cache_path}")
try:
await asyncio.to_thread(
update_media_file_access_sync,
DB_PATH, str(channel), post_id, file_unique_id,
datetime.now().timestamp()
)
except Exception as e:
logger.error(f"Failed to update timestamp for {channel}/{post_id}/{file_unique_id}: {str(e)}")
_access_updates[(str(channel), post_id, file_unique_id)] = datetime.now().timestamp()
return cache_path, False
file_id = await find_file_id_in_message(message, file_unique_id)
@@ -470,46 +629,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 +695,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)
@@ -934,6 +1065,46 @@ async def get_raw_post_json(channel: str, post_id: int, request: Request, token:
raise HTTPException(status_code=500, detail=error_message) from e
@app.get("/ping")
async def ping() -> JSONResponse:
"""Lightweight liveness probe for the container healthcheck.
Reflects process/event-loop liveness (always answers in microseconds) and TG liveness
from the watchdog's last-probe data. It MUST NOT issue any Telegram RPC (no get_me,
no safe_get_messages), touch SQLite, or walk the filesystem — that is the whole point:
it stays instant and truthful even while a real TG RPC is hung. It only reads the
already-recorded watchdog timestamp and the is_connected bool.
"""
age = client.watchdog_last_ok_age() # seconds since last OK probe, None if never
# is_connected is None before client.start() and a bool afterwards; coerce so the JSON
# "connected" field is always a bool (never null) and the pre-start window reports false.
connected = bool(client.client.is_connected)
threshold = Config["tg_ping_unhealthy_after"]
# age is None right after boot: the watchdog hasn't run its first probe yet. Treat that
# as healthy (gate on connected only) so a freshly-started container is not killed before
# its first probe cycle — otherwise start_period would have to cover a full watchdog interval.
#
# The staleness branch (age >= threshold => degraded) is only meaningful while the watchdog
# is running to refresh age. With the watchdog DISABLED (TG_WATCHDOG_ENABLED=false) nothing
# refreshes age — yet a disconnect-flap restart can still stamp it once (see _restart_client,
# which runs before the watchdog-enabled gate), after which age only grows. Letting that
# stale age drive /ping to 503 would spuriously fail the container healthcheck on a live
# connection and trigger an autoheal restart. So gate staleness on the watchdog being on;
# with it off, /ping is a pure connectivity check (no zombie-session detection — that
# TG-liveness signal only exists while the watchdog runs).
healthy = connected and (
not Config["tg_watchdog_enabled"] or age is None or age < threshold
)
return JSONResponse(
{
"status": "ok" if healthy else "degraded",
"connected": connected,
"last_probe_age_s": round(age, 1) if age is not None else None,
"threshold_s": threshold,
},
status_code=200 if healthy else 503,
)
@app.get("/health")
@app.get("/health/{token}")
async def health_check(request: Request, token: str | None = None) -> Response:
@@ -979,7 +1150,7 @@ async def health_check(request: Request, token: str | None = None) -> Response:
@app.get("/media/{channel}/{post_id}/{file_unique_id}/{digest}", response_model=None)
@app.get("/media/{channel}/{post_id}/{file_unique_id}", response_model=None)
async def get_media(channel: str, post_id: int, file_unique_id: str, request: Request, digest: str | None = None) -> StreamingResponse|Response:
async def get_media(channel: str, post_id: int, file_unique_id: str, request: Request, digest: str | None = None) -> Response:
try:
url = f"{channel}/{post_id}/{file_unique_id}"
if not verify_media_digest(url, digest):
@@ -1005,29 +1176,46 @@ async def get_media(channel: str, post_id: int, file_unique_id: str, request: Re
if os.path.exists(cache_path) and os.path.getsize(cache_path) > 0:
# File is already in cache — skip semaphore and serve directly
logger.info(f"pre_semaphore_cache_hit: {channel}/{post_id}/{file_unique_id}")
# Fire-and-forget timestamp update with error handling to avoid silent failures
async def _update_access(_ch, _pid, _fid):
try:
await asyncio.to_thread(
update_media_file_access_sync,
DB_PATH, str(_ch), _pid, _fid,
datetime.now().timestamp()
)
except Exception as _e:
logger.warning(f"Failed to update access time for {_ch}/{_pid}/{_fid}: {_e}")
asyncio.create_task(_update_access(channel, post_id, file_unique_id))
# Record the access time into the accumulator instead of firing a per-hit
# SQLite write. Key channel as str(channel) — see _access_updates.
_access_updates[(str(channel), post_id, file_unique_id)] = datetime.now().timestamp()
return await prepare_file_response(cache_path, request=request,
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 +1231,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
config.py
+21
View File
@@ -123,4 +123,25 @@ 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),
# /ping reports TG as unhealthy once the last successful watchdog probe is older than
# this many seconds. Default is derived from the watchdog cadence: it is roughly how
# long the watchdog itself would take to give up and trigger a restart —
# interval * (failures + 1) + timeout. With the defaults (60,3,10) that is 250s, so a
# transient slow probe never flaps /ping, but a genuinely stuck session (no successful
# probe for ~4 min) surfaces as 503 before/around the time the watchdog restarts.
"tg_ping_unhealthy_after": _parse_int_env(
"TG_PING_UNHEALTHY_AFTER",
_parse_int_env("TG_WATCHDOG_INTERVAL", 60) * (_parse_int_env("TG_WATCHDOG_FAILURES", 3) + 1)
+ _parse_int_env("TG_WATCHDOG_TIMEOUT", 10),
),
# 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),
# Size of the asyncio default ThreadPoolExecutor. SQLite/python-magic/pickle/os.walk
# all run via asyncio.to_thread; the interpreter default (min(32, cpu+4)) is only 5-6
# on a 1-2 CPU container, which starves those under load. 32 gives ample headroom.
"io_thread_pool_size": _parse_int_env("IO_THREAD_POOL_SIZE", 32),
}
dockercompose.yml
+12 -3
View File
@@ -25,6 +25,10 @@ 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)
# IO_THREAD_POOL_SIZE: 32 # Size of the asyncio default threadpool for blocking I/O (SQLite/python-magic/pickle/os.walk); raise on a busy 1-2 CPU box (default: 32)
PYROGRAM_BRIDGE_URL: https://pgbridge.example.com
API_PORT: 80
TOKEN: ХХХ
@@ -49,10 +53,15 @@ services:
com.centurylinklabs.watchtower.enable: "true"
autoheal: true
healthcheck:
test: ["CMD", "curl", "-s", "-f", "-o", "/dev/null", "http://127.0.0.1:80/rss/vvzvlad_lytdybr?limit=1"]
interval: 30m
# Lightweight process/loop liveness probe. /ping never touches Telegram or the
# filesystem, so it answers instantly even while a TG RPC is hung — unlike the old
# /rss?limit=1 check, which could exceed the 5s timeout on a cold cache and get the
# container restarted mid-download (corrupted temp files). TG liveness is now judged
# by the in-process watchdog, which /ping reports via its last-probe age.
test: ["CMD", "curl", "-sf", "http://127.0.0.1:80/ping"]
interval: 5m
timeout: 5s
retries: 2
retries: 3
start_period: 30s
start_interval: 5s
docs/stability-verification.md
+148
View File
@@ -0,0 +1,148 @@
# Стадия 7 — сквозная верификация плана стабилизации
Дата: 2026-07-05. Ветка: `fix/stage-7-verification` (на базе `fix/stage-6-healthcheck`,
кумулятивно содержит стадии 1–6). Эта стадия **ничего не меняет в поведении** прод-кода —
только автоматизирует проверки и фиксирует, что должен наблюдать оператор после деплоя.
## Итог автотестов
Каноничный прогон (изоляция `config` через `sys.modules`): `python -m pytest tests/`.
```
260 passed
```
Раскладка по стадиям (все зелёные):
| Стадия | Файл тестов | Кол-во | Что покрывает |
|--------|-------------|--------|----------------|
| 1 — анти-зависания | `tests/test_stage1_hangs.py` | 6 | таймаут RPC-гейта без утечки пермита; воркер переживает Exception/FloodWait, `task_done` сбалансирован; отмена в spacing-ожидании не теряет пермит |
| 2 — статика/большие видео | `tests/test_stage2_static.py` | 15 | атомарный `_download_atomic` (publish-on-rename, чистка `.part.` при таймауте/zero-size/гонке); дедуп конкурентных скачиваний; FloodWait→429; touch mtime у `temp_*`; свипер чистит `.part.`+legacy `.tmp.`; баланс HTTP-семафора |
| 3 — FileResponse | `tests/test_stage3_fileresponse.py` | 19 | матрица Range (0-499/500-/-500/за EOF→416/мусор→400/мульти-range→206 multipart); сохранены mtime-touch, delete_after-BackgroundTask, MIME-кэш; чистый ASGI-логгер |
| 4 — гигиена event loop | `tests/test_stage4_eventloop.py` | 19 | ленивый `raw_message`; вынос side-effect IO из `process_message` (bulk upsert media id); рендер-пайплайн ушёл в поток без create_task/get_running_loop; XSS вычищен во всех 4 выводах ровно одним проходом |
| 5 — батчинг SQLite | `tests/test_stage5_sqlite.py` | 8 | кэш-хит пишет в аккумулятор, а не в SQLite; flush→DB; snapshot-then-clear не теряет поздние апдейты; re-queue при сбое; str(channel)-ключи |
| 6 — healthcheck | `tests/test_stage6_healthcheck.py` | 11 | `/ping` 200/503 по connected+age; ноль TG RPC; watchdog отключён→чистая проверка коннекта; `watchdog_last_ok_age()` |
| 7 — интеграция | `tests/test_stage7_integration.py` | 8 | сквозные сценарии (см. ниже) |
| — регрессия парсера | `tests/test_postparser_*.py` | 174 | заголовки/флаги/автор (существовавшие до плана) |
## Новые интеграционные тесты стадии 7 (DoD → сквозное доказательство)
`tests/test_stage7_integration.py` драйвит **реальные точки входа** (`get_media`, `ping`,
flush), чтобы поймать регрессии, которые проявляются только при взаимодействии стадий:
1. **Range на уровне маршрута `/media`** — `test_media_route_range_prefix_0_99`,
`_range_suffix_last_100`, `_range_unsatisfiable_416`. Прогоняют `get_media` (кэш-хит) →
`prepare_file_response` → `FileResponse` через реальный ASGI: `bytes=0-99`→206 с точным
Content-Range/длиной и байтами, `bytes=-100`→206 (суффикс), `bytes=999999999-`→416 (`*/size`).
Ловит: если кэш-хит перестанет доходить до FileResponse (ре-буферизация тела, ручные
заголовки) или сломается связка digest-гейта — Range перестанет работать. (Стадия 3
пинит это на `prepare_file_response` напрямую; здесь — сшивка стадий 2+3.)
2. **`/ping` быстр и без RPC при висящей операции** — `test_ping_prompt_and_rpc_free_while_slow_op_pending`
и `_reports_degraded_promptly_while_slow_op_pending`. Пока фейковая медленная корутина
(модель зависшего hot-path) висит на `Event`, `ping()` возвращает 200/503 корректно и
**до** завершения медленной операции, при нуле вызовов `get_me`/`safe_get_messages`.
Ловит: рекаплинг `/ping` к TG RPC или к любому awaitable, который может застопориться.
3. **Дедуп + очистка при disconnect через реальный `get_media`** —
`test_get_media_concurrent_shares_one_download_and_drains_registry` (2 конкурентных
запроса → одна скачка, `_inflight` пуст) и `_request_cancel_does_not_stick_registry_or_hang_sibling`
(отмена запроса-«клиента» не оставляет застрявший ключ и не вешает соседа). Ловит:
возврат скачивания в корутину запроса (отмена убила бы download) или потерю `finally`-pop.
4. **str(channel)-ключ access-time end-to-end** — `test_media_cache_hit_flush_updates_str_keyed_db_row`.
Кэш-хит `/media` записывает str-ключ, flush обновляет ту же строку в реальной SQLite
(hit→flush→DB). Сшивает две половины: ключ, который пишет hot-path, и WHERE, по которому
апдейтит flush. Ловит любое расхождение ключа аккумулятора и WHERE bulk-UPDATE (напр.
перестановку колонок в SQL — проверено мутацией: тест краснеет, `added` остаётся stale).
(int/str-хазард самого канала живёт в `download_media_file` и закрыт стадией 5 —
`test_cache_hit_keys_channel_as_str`; здесь покрыта связка get_media+flush.)
## Соответствие DoD стадий → доказательство
- **DoD 1** «ни один путь не ждёт TG без таймаута; воркер не умирает молча» →
`test_stage1_hangs.py` (гейт-таймаут, живучесть воркера). Живая проверка supervision
под нагрузкой — **наблюдение оператора** (лог `supervised_task_crashed/…_exited`).
- **DoD 2** «клиенту никогда не отдаётся неполный файл; флуд→429; обрезки не живут >часа» →
`test_stage2_static.py` целиком + интеграционный дедуп-тест стадии 7.
- **DoD 3** «Range-тесты зелёные; поведение эндпоинта эквивалентно (± RFC-допущения)» →
`test_stage3_fileresponse.py` + Range на уровне `/media` (стадия 7). «Нет потока-на-чанк»
— **наблюдение оператора** (нагрузочный запрос большого файла, число io-потоков).
- **DoD 4** «генерация 100-сообщ. фида не блокирует луп (параллельный /ping <100 мс); XSS
зелёный; media id сохраняются» → `test_stage4_eventloop.py` (рендер в потоке, XSS,
bulk upsert) + `/ping`-decoupling стадии 7. Живой замер «<100 мс на проде» —
**наблюдение оператора**.
- **DoD 5** «на кэш-хит /media ноль обращений к SQLite; фоновая запись раз в минуту» →
`test_stage5_sqlite.py` + str-ключ end-to-end стадии 7.
- **DoD 6** «во время зависшего TG RPC /ping отвечает мгновенно (503); контейнер не
рестартится от медленного фида» → `test_stage6_healthcheck.py` + `/ping`-under-slow-op
стадии 7. Отсутствие autoheal-рестартов на холодном кэше — **наблюдение оператора**.
## Ручные curl-сценарии для оператора (пост-деплой, локально в контейнере)
Подставить реальный `{channel}/{post_id}/{fid}/{digest}` (валидная подпись обязательна).
```bash
# 1. Range-тройка (ожидания: 206 / 206 / 416)
curl -s -D- -o /dev/null -H "Range: bytes=0-99" "http://127.0.0.1:80/media/{ch}/{pid}/{fid}/{digest}"
curl -s -D- -o /dev/null -H "Range: bytes=-100" "http://127.0.0.1:80/media/{ch}/{pid}/{fid}/{digest}"
curl -s -D- -o /dev/null -H "Range: bytes=999999999-" "http://127.0.0.1:80/media/{ch}/{pid}/{fid}/{digest}"
# 2. Параллельные запросы одного БОЛЬШОГО видео (>100 MB), пока не в кэше:
# оба должны получить ПОЛНЫЙ файл (одинаковый размер), без частичной отдачи.
URL="http://127.0.0.1:80/media/{ch}/{pid}/{bigfid}/{digest}"
curl -s -o /tmp/a "$URL" & curl -s -o /tmp/b "$URL" & wait
ls -l /tmp/a /tmp/b # размеры совпадают и равны полному файлу; на диске нет *.part.*
# 3. Обрыв на середине стрима — нет утечки тасков/фд:
# считать fd до/после и убедиться, что не растут монотонно.
lsof -p $(pgrep -f api_server) | wc -l # baseline
timeout 2 curl -s -o /dev/null "$URL"; sleep 5 # оборвать скачку на середине (несколько раз)
lsof -p $(pgrep -f api_server) | wc -l # не выросло относительно baseline
# 4. Генерация большого фида + параллельный /ping (<100 мс во время генерации):
curl -s -o /dev/null "http://127.0.0.1:80/rss/{big_channel}" &
for i in $(seq 1 20); do curl -s -o /dev/null -w "%{time_total}\n" "http://127.0.0.1:80/ping"; done
wait # все замеры /ping должны быть заметно < 0.100 s
```
> Сценарии 2 и 3 (реальная скачка большого видео + реальный подсчёт fd через `lsof`) в
> headless-тестах **намеренно не подделаны** — им нужен живой сервер, реальные загрузки и
> реальные файловые дескрипторы. Дедуп-инвариант и disconnect-очистка доказаны на уровне
> реестра/`get_media` (тесты стадии 7 №3), но «нет утечки fd на проде» проверяет оператор.
## Diag-логи для наблюдения после деплоя
Ожидаемая динамика (сравнить с до-деплойным baseline):
- `diag_semaphore_wait` — ожидание HTTP-семафора должно **упасть** (реже/меньше секунд).
- `diag_download_timing` — время скачивания стабилизируется; нет длинных «зависаний».
- `diag_sanitize_slow` — почти **исчезнуть** (один sanitize на выходную границу, стадия 4).
- `watchdog: heartbeat` — **продолжаются** штатно (живость TG-сессии).
- На момент FloodWait — **нет всплеска 404** на `/media`; вместо этого `media_flood_wait`
и ответы **429** с `Retry-After` (стадия 2.3).
- Признаки supervision: любые `supervised_task_crashed` / `supervised_task_exited` на
CRITICAL — сигнал разобраться, но задача при этом перезапускается (не тихая смерть).
## План отката (rollback)
Единица отката — **стадия целиком**, не отдельный коммит. Каждая стадия живёт на своей
ветке/PR (`fix/stage-1-hangs` … `fix/stage-7-verification`) и, как правило, состоит из
**двух коммитов**: feature-коммит + фикс review-раунда (последний нередко чинит реальный
баг — напр. fail-closed XSS в стадии 4, watchdog-gate в стадии 6). Поэтому `git revert`
одного коммита осиротит фикс review-раунда и даст несогласованный откат. Откатывать нужно
на гранулярности стадии; порядок стадий = порядок деплоя.
- **Как откатывать стадию** — зависит от того, как ветки влиты в `fix/stability`/`main`:
- если стадия влита **squash-merge** (один коммит на стадию) — `git revert <squash-commit>`
откатывает её целиком, однозначно;
- если стадия влита **merge-коммитом** — `git revert -m 1 <merge-commit>` откатывает весь
вклад ветки (оба коммита) разом;
- если история линейная (rebase/fast-forward) — реверт **всех** коммитов стадии
(`git revert <feature>..<review-fix>` включительно), а не одного.
- Стадии 1 и 2 — низкорисковые, деплоятся первыми; откатываются независимо от остальных.
- Стадия 3 (FileResponse) независима — реверт возвращает прежний ручной стриминг.
- Стадия 4 требует, чтобы 4.2 откатывалась не позже 4.3 (иначе рендер-в-потоке остался бы
без безопасного пути сохранения media id) — откатывать стадию 4 целиком.
- Стадии 5 и 6 независимы, откатываются по отдельности.
- Стадия 7 — только тесты и этот документ: реверт ничего не меняет в проде.
Прод-деплой, живое наблюдение diag-логов и обновление прод-compose (healthcheck→`/ping`,
вне репозитория) — **зона ответственности оператора** (пункты 3–4 плана стадии 7).
file_io.py
+39
View File
@@ -69,6 +69,24 @@ def upsert_media_file_id_sync(db_path: str, channel: str, post_id: int, file_uni
)
def upsert_media_file_ids_bulk_sync(db_path: str, entries: List[tuple]) -> None:
"""Insert or replace multiple media file ID records in a single transaction.
entries: iterable of (channel, post_id, file_unique_id, added) tuples.
Uses executemany for batched upserts (one connection, one commit).
"""
if not entries:
return
with _db_connection(db_path) as conn:
conn.executemany(
"""INSERT INTO media_file_ids (channel, post_id, file_unique_id, added)
VALUES (?, ?, ?, ?)
ON CONFLICT(channel, post_id, file_unique_id)
DO UPDATE SET added = excluded.added""",
entries,
)
def update_media_file_access_sync(db_path: str, channel: str, post_id: int, file_unique_id: str, added: float) -> None:
"""Update the access timestamp for an existing media file ID record."""
with _db_connection(db_path) as conn:
@@ -78,6 +96,27 @@ def update_media_file_access_sync(db_path: str, channel: str, post_id: int, file
)
def update_media_file_access_bulk_sync(db_path: str, entries: List[tuple]) -> None:
"""Update the access timestamp for multiple existing media file ID records.
entries: iterable of (channel, post_id, file_unique_id, added) tuples.
Uses executemany (one connection, one commit) so a batch of cache-hit access
updates costs a single SQLite transaction instead of one connect+UPDATE per hit.
Rows that do not exist are simply not matched by the WHERE clause (no-op), mirroring
the single-row update_media_file_access_sync. An empty batch is a no-op.
"""
if not entries:
return
with _db_connection(db_path) as conn:
conn.executemany(
"UPDATE media_file_ids SET added = ? WHERE channel = ? AND post_id = ? AND file_unique_id = ?",
# Reorder each (channel, post_id, file_unique_id, added) tuple to match the
# UPDATE's placeholder order (added first, then the WHERE key columns).
[(added, channel, post_id, file_unique_id)
for (channel, post_id, file_unique_id, added) in entries],
)
def get_all_media_file_ids_sync(db_path: str) -> List[dict]:
"""Return all rows from media_file_ids as a list of dicts."""
with _db_connection(db_path) as conn:
post_parser.py
+87 -42
View File
@@ -21,16 +21,13 @@ from pyrogram.enums import MessageMediaType
from bleach.css_sanitizer import CSSSanitizer
from bleach import clean as HTMLSanitizer
from config import get_settings
from file_io import upsert_media_file_id_sync, DB_PATH
from file_io import upsert_media_file_ids_bulk_sync, DB_PATH
from url_signer import generate_media_digest
Config = get_settings()
logger = logging.getLogger(__name__)
# Module-level counter for in-flight persist tasks (used only in diagnostics)
_persist_pending_count = 0
#tests
#http://127.0.0.1:8000/post/html/DragorWW_space/114 — video
@@ -65,6 +62,12 @@ _persist_pending_count = 0
class PostParser:
def __init__(self, client):
self.client = client
# Media file-id records collected during rendering. _save_media_file_ids
# appends (channel, post_id, file_unique_id, ts) tuples here instead of
# touching asyncio/DB directly (see 4.2). A fresh PostParser is created per
# request and rendering runs in a single thread, so this list is thread-safe.
# The caller flushes it once via upsert_media_file_ids_bulk_sync after render.
self._pending_media_ids: List[tuple] = []
@staticmethod
def get_all_possible_flags() -> List[str]:
@@ -106,7 +109,14 @@ class PostParser:
logger.error(f"post_not_found_or_empty: channel {prepared_channel_id}, post_id {post_id}")
return None
processed_message = self.process_message(message)
# Single-post outputs need sanitized body/footer (no whole-feed pass exists here).
# raw_message is only needed for JSON output and for debug HTML.
include_raw = (output_type == 'json') or debug
processed_message = self.process_message(message, include_raw=include_raw, sanitize=True)
# Flush media file-id records collected during rendering with a single bulk upsert.
await self._flush_pending_media_ids()
if output_type == 'html':
return self._format_html(processed_message, debug)
elif output_type == 'json':
@@ -480,11 +490,15 @@ class PostParser:
html_content.append(f'<div class="message-body">{data["html"]["body"]}</div>')
html_content.append(f'<div class="message-footer">{data["html"]["footer"]}</div>')
# Add raw JSON debug output if debug is enabled
# Add raw JSON debug output if debug is enabled.
# raw_message is the full str(message) serialization and may contain user-controlled
# text with HTML/JS — escape it before dropping it into the <pre> (bleach can't run
# here because <pre> is not in the allowed-tags whitelist and would be stripped).
if debug:
raw_escaped = html.escape(str(data.get("raw_message", "")))
html_content.append(f'<pre class="debug-json" style="background: #f5f5f5;'
f'padding: 10px; margin-top: 20px; overflow-x: auto; font-size: 10px;'
f'white-space: pre-wrap;">{data["raw_message"]}</pre>')
f'white-space: pre-wrap;">{raw_escaped}</pre>')
html_data = '\n'.join(html_content)
return html_data
@@ -500,11 +514,39 @@ class PostParser:
return return_html
def process_message(self, message: Message) -> Dict[Any, Any]:
# Compute html body once — avoids triple _generate_html_body calls
def process_message(self, message: Message, include_raw: bool = False, sanitize: bool = True) -> Dict[Any, Any]:
"""Build the processed representation of a message.
Args:
message: the Pyrogram message.
include_raw: when True, compute the (expensive) full ``str(message)``
serialization into ``result['raw_message']``. Only JSON output and
debug HTML need it; feed generation must pass False to avoid
serializing every post.
sanitize: when True, run the html body and footer through a single
bleach pass each. Single-post HTML and JSON need this (there is no
whole-feed pass on those paths). Feed generation passes False and
relies on the final sanitize in rss_generator (per-post for RSS,
whole-feed for HTML), so no
fragment is sanitized more than once.
"""
# Compute html body once — avoids triple _generate_html_body calls.
# The internal per-fragment sanitize passes were removed (4.4); sanitize
# exactly once per output boundary here when requested.
html_body = self._generate_html_body(message)
# NOTE (stage-4 4.4 consequence): flags are now extracted from the PRE-sanitize
# body (the per-fragment sanitize that used to run inside _generate_html_body was
# removed). Legitimate links are unaffected — bleach keeps whitelisted
# <a href="http(s)://…"> — so link/foreign_channel/mention flags are identical for
# normal content. They can differ ONLY for URL-like text that bleach would strip
# (e.g. a URL inside a disallowed attribute); flags are non-security (used for
# exclude_flags filtering / display), so this edge divergence is accepted rather
# than re-adding a per-message sanitize pass that 4.4 deliberately eliminated.
flags = self._extract_flags(message, html_body=html_body)
footer = self.generate_html_footer(message, flags_list=flags)
if sanitize:
html_body = self._sanitize_html(html_body)
footer = self._sanitize_html(footer)
result = {
'channel': self.get_channel_username(message),
'message_id': message.id,
@@ -523,8 +565,10 @@ class PostParser:
'media_group_id': message.media_group_id,
'service': getattr(message, "service", None),
'show_caption_above_media': getattr(message, 'show_caption_above_media', False),
'raw_message': str(message)
}
if include_raw:
# Full serialization of the message — expensive; only for JSON/debug.
result['raw_message'] = str(message)
# Add webpage data if available
if webpage := getattr(message, "web_page", None):
@@ -671,8 +715,10 @@ class PostParser:
if message.forward_origin: content_body.append(f"--- Forwarded post end ---") # Forward info end
content_body.append(f'</div><br>')
# NOTE: sanitize is NOT applied here. Sanitization happens exactly once per
# output boundary (process_message for single-post/JSON; in rss_generator the
# per-post pass for RSS and the whole-feed pass for HTML). See the map (4.4).
html_body = '\n'.join(content_body)
html_body = self._sanitize_html(html_body)
return html_body
def _generate_html_media(self, message: Message) -> str:
@@ -751,8 +797,9 @@ class PostParser:
content_media.append(webpage_html)
content_media.append('</div>')
# Not sanitized here — this fragment is embedded in the body and sanitized
# once at the output boundary (see the 4.4 sanitize coverage map).
html_media = '\n'.join(content_media)
html_media = self._sanitize_html(html_media)
return html_media
@@ -832,8 +879,9 @@ class PostParser:
flags_html = self._format_flags(current_flags)
content_footer.append('<br>' + flags_html)
# Not sanitized here — sanitized once at the output boundary (4.4 coverage map):
# process_message for single-post/JSON; per-post (RSS) / whole-feed (HTML) for feeds.
html_footer = '\n'.join(content_footer)
html_footer = self._sanitize_html(html_footer)
return html_footer
@@ -919,8 +967,10 @@ class PostParser:
if links:
parts.append('&nbsp;|&nbsp;'.join(links))
# Raw fragment — embedded in the footer, sanitized once at the output
# boundary (see the 4.4 sanitize coverage map). Not sanitized here.
result_html = '<br>'.join(parts) if parts else None
return self._sanitize_html(result_html) if result_html else None
return result_html if result_html else None
except Exception as e:
logger.error(f"reactions_views_links_error: message_id {message.id}, error {str(e)}")
@@ -969,15 +1019,31 @@ class PostParser:
logger.error(f"file_id_extraction_error: media_type {message.media}, error {str(e)}")
return None
async def _persist_media_file_id_async(self, channel: str, post_id: int, file_unique_id: str, added: float) -> None:
"""Persist a single media file ID record to SQLite (fire-and-forget)."""
async def _flush_pending_media_ids(self) -> None:
"""Persist media file-id records collected during rendering with ONE bulk upsert.
Called by the caller (get_post / rss_generator) after rendering completes.
Runs the blocking SQLite write in a thread. No-op when nothing was collected.
"""
entries = self._pending_media_ids
if not entries:
return
try:
await asyncio.to_thread(upsert_media_file_id_sync, DB_PATH, channel, post_id, file_unique_id, added)
logger.debug(f"persist_media_file_id: upserted {channel}/{post_id}/{file_unique_id}")
await asyncio.to_thread(upsert_media_file_ids_bulk_sync, DB_PATH, entries)
logger.debug(f"persist_media_file_ids_bulk: upserted {len(entries)} records")
except Exception as e:
logger.error(f"file_id_save_error: error upserting {channel}/{post_id}/{file_unique_id}, error {str(e)}")
logger.error(f"file_id_bulk_save_error: error bulk-upserting {len(entries)} records, error {str(e)}")
finally:
# Clear regardless of outcome so a retry does not double-persist a stale batch.
self._pending_media_ids = []
def _save_media_file_ids(self, message: Message) -> None:
"""Collect a media file-id record for later bulk persistence.
IMPORTANT (4.2): this runs inside the render thread, so it must NOT touch
asyncio or the DB. It only appends to self._pending_media_ids; the caller
flushes the batch via _flush_pending_media_ids() after rendering.
"""
try:
channel_username = self.get_channel_username(message)
if not channel_username:
@@ -1003,29 +1069,8 @@ class PostParser:
if file_unique_id:
added_ts = datetime.now().timestamp()
try:
loop = asyncio.get_running_loop()
if loop.is_running():
global _persist_pending_count
task = loop.create_task(
self._persist_media_file_id_async(channel_username, message.id, file_unique_id, added_ts)
)
_persist_pending_count += 1
def _on_task_done(t):
global _persist_pending_count
_persist_pending_count -= 1
task.add_done_callback(_on_task_done)
if _persist_pending_count > 5:
logger.warning(f"persist_task_queue: {_persist_pending_count} pending _persist_media_file_id_async tasks (msg {message.id})")
logger.debug(f"persist_task_created: queued for {channel_username}/{message.id}/{file_unique_id}, total pending: {_persist_pending_count}")
else:
# Fallback sync path (should not occur during normal FastAPI runtime)
upsert_media_file_id_sync(DB_PATH, channel_username, message.id, file_unique_id, added_ts)
except Exception as e:
logger.error(f"file_id_save_error: error saving {channel_username}/{message.id}/{file_unique_id}, error {str(e)}")
# Thread-safe: just append; the caller persists the batch.
self._pending_media_ids.append((channel_username, message.id, file_unique_id, added_ts))
except Exception as e:
logger.error(f"file_id_collection_error: message_id {message.id}, error {str(e)}")
requirements.txt
+4
View File
@@ -1,4 +1,7 @@
fastapi==0.115.8
# Pinned explicitly (fastapi only requires a range): the stage-3 FileResponse Range
# behaviour and the multipart-byteranges test are tied to this Starlette version.
starlette==0.45.3
uvicorn==0.34.0
python-multipart==0.0.20
Kurigram==2.2.22
@@ -12,3 +15,4 @@ python-magic==0.4.27
bleach[css]==6.1.0
types-bleach
pytest-asyncio==1.4.0
httpx==0.28.1
rss_generator.py
+81 -25
View File
@@ -14,6 +14,7 @@ import logging
import asyncio
import re
import time
from html import escape as html_escape
from datetime import datetime, timezone
from types import SimpleNamespace
from typing import Optional
@@ -30,9 +31,12 @@ Config = get_settings()
logger = logging.getLogger(__name__)
async def _create_time_based_media_groups(messages: list[Message], merge_seconds: int = 5) -> list[Message]:
def _create_time_based_media_groups(messages: list[Message], merge_seconds: int = 5) -> list[Message]:
"""
Create media groups based on time difference between messages
Plain synchronous function (contains no await): runs inside the render thread
via _render_pipeline. Must not touch asyncio.
"""
# Deep-copy the input list to avoid mutating cached Message objects (the cache may
# reuse the same objects across calls with different merge_seconds values)
@@ -92,9 +96,11 @@ async def _create_time_based_media_groups(messages: list[Message], merge_seconds
return messages_sorted
async def _create_messages_groups(messages: list[Message]) -> list[list[Message]]:
def _create_messages_groups(messages: list[Message]) -> list[list[Message]]:
"""
Process messages into formatted posts, handling media groups
Plain synchronous function (contains no await): runs inside the render thread.
"""
processing_groups: list[list[Message]] = []
media_groups: dict[str | int, list[Message]] = {}
@@ -136,9 +142,11 @@ async def _create_messages_groups(messages: list[Message]) -> list[list[Message]
return processing_groups
async def _trim_messages_groups(messages_groups: list[list[Message]], limit: int):
def _trim_messages_groups(messages_groups: list[list[Message]], limit: int):
"""
Trim messages groups to limit
Plain synchronous function (contains no await): runs inside the render thread.
"""
if len(messages_groups) > limit: # Trim groups if they exceed the specified limit
messages_groups = messages_groups[:limit]
@@ -193,12 +201,13 @@ def processed_message_to_tg_message(processed_message: dict) -> Message:
return tg_message_mock # type: ignore
async def _render_messages_groups(messages_groups: list[list[Message]],
post_parser: PostParser,
exclude_flags: str | None = None,
def _render_messages_groups(messages_groups: list[list[Message]],
post_parser: PostParser,
exclude_flags: str | None = None,
exclude_text: str | None = None):
"""
Render message groups into HTML format
Plain synchronous function (contains no await): runs inside the render thread.
Args:
messages_groups: List of message groups (each group is a list of messages)
post_parser: PostParser instance
@@ -213,7 +222,9 @@ async def _render_messages_groups(messages_groups: list[list[Message]],
try:
if len(group) == 1: # Single message - simple case
one_message = group[0]
message_data = post_parser.process_message(one_message)
# Feed path: raw_message not needed and sanitize deferred to the final
# whole-feed pass, so each fragment is sanitized exactly once.
message_data = post_parser.process_message(one_message, include_raw=False, sanitize=False)
html_parts = [
f'<div class="message-body">{message_data["html"]["body"]}</div>',
f'<div class="message-footer">{message_data["html"]["footer"]}</div>'
@@ -228,7 +239,7 @@ async def _render_messages_groups(messages_groups: list[list[Message]],
'flags': message_data['flags']
})
else: # Multiple messages in group - merge text and html body
processed_messages = [post_parser.process_message(msg) for msg in group]
processed_messages = [post_parser.process_message(msg, include_raw=False, sanitize=False) for msg in group]
# Determine main message for header/footer/title
main_message = next(
@@ -309,7 +320,31 @@ async def _render_messages_groups(messages_groups: list[list[Message]],
rendered_posts.sort(key=lambda x: x['date'] if x['date'] is not None else 0.0, reverse=True)
return rendered_posts
async def generate_channel_rss(channel: str | int,
def _render_pipeline(messages: list[Message],
post_parser: PostParser,
limit: int,
exclude_flags: str | None,
exclude_text: str | None,
merge_seconds: int,
time_based_merge: bool):
"""
Full synchronous feed render pipeline (grouping + trimming + rendering).
Runs entirely in a worker thread via a single asyncio.to_thread call. It contains
NO await, NO asyncio primitives, NO create_task/get_running_loop — all the CPU-heavy
work (deepcopy, grouping, bleach-free rendering) happens here off the event loop.
Media file-id records are accumulated on post_parser._pending_media_ids and flushed
by the caller after this returns.
"""
if time_based_merge:
messages = _create_time_based_media_groups(messages, merge_seconds)
message_groups = _create_messages_groups(messages)
message_groups = _trim_messages_groups(message_groups, limit)
return _render_messages_groups(message_groups, post_parser, exclude_flags, exclude_text)
async def generate_channel_rss(channel: str | int,
client: Client,
limit: int = 20,
exclude_flags: str | None = None,
@@ -389,14 +424,23 @@ async def generate_channel_rss(channel: str | int,
messages_elapsed = time.time() - messages_start_time
logger.debug(f"rss_messages_retrieval_timing: channel {channel}, {len(messages)} messages retrieved in {messages_elapsed:.3f} seconds")
# Process messages into groups and render them
# Process messages into groups and render them.
# The whole grouping/trimming/rendering pipeline is CPU-heavy (deepcopy +
# per-message rendering) and contains no await, so run it in ONE worker
# thread to keep the event loop responsive.
processing_start_time = time.time()
if Config['time_based_merge']:
messages = await _create_time_based_media_groups(messages, merge_seconds)
message_groups = await _create_messages_groups(messages)
message_groups = await _trim_messages_groups(message_groups, limit)
final_posts = await _render_messages_groups(message_groups, post_parser, exclude_flags, exclude_text)
try:
final_posts = await asyncio.to_thread(
_render_pipeline, messages, post_parser, limit,
exclude_flags, exclude_text, merge_seconds, Config['time_based_merge'],
)
finally:
# Persist media file-ids collected during rendering with a single bulk
# upsert — in a finally so a partial render still records what it collected
# (the flush is best-effort and swallows its own errors, so it cannot mask a
# render exception).
await post_parser._flush_pending_media_ids()
processing_elapsed = time.time() - processing_start_time
logger.debug(f"rss_messages_processing_timing: channel {channel}, {len(final_posts)} posts processed in {processing_elapsed:.3f} seconds")
@@ -443,8 +487,13 @@ async def generate_channel_rss(channel: str | int,
)
sanitized_html = await asyncio.to_thread(_sanitize_sync, post['html'])
except Exception as e:
# FAIL CLOSED: since 4.4 removed the per-fragment passes, post['html'] is
# raw channel-controlled HTML, and this is its ONLY sanitize. If bleach
# itself throws (e.g. RecursionError on pathological nested HTML), do NOT
# emit the raw payload (that would be stored XSS) — html.escape it so the
# content survives as inert text.
logger.error(f"rss_html_sanitization_error: channel {channel}, message_id {post['message_id']}, error {str(e)}")
sanitized_html = post['html']
sanitized_html = html_escape(post['html'])
fe.content(content=sanitized_html, type='CDATA')
if post['date'] is not None:
@@ -601,16 +650,19 @@ async def generate_channel_html(channel: str | int,
enrichment_elapsed = time.time() - enrichment_start_time
logger.debug(f"html_reply_enrichment_timing: channel {channel}, replies enriched in {enrichment_elapsed:.3f} seconds")
# Process messages into groups and render them
# Process messages into groups and render them in ONE worker thread
# (CPU-heavy, no await) to keep the event loop responsive.
processing_start_time = time.time()
if Config['time_based_merge']:
messages = await _create_time_based_media_groups(messages, merge_seconds)
try:
final_posts = await asyncio.to_thread(
_render_pipeline, messages, post_parser, limit,
exclude_flags, exclude_text, merge_seconds, Config['time_based_merge'],
)
finally:
# Persist media file-ids collected during rendering with a single bulk
# upsert — in a finally so a partial render still records what it collected.
await post_parser._flush_pending_media_ids()
# Process messages into groups and render them
message_groups = await _create_messages_groups(messages)
message_groups = await _trim_messages_groups(message_groups, limit)
final_posts = await _render_messages_groups(message_groups, post_parser, exclude_flags, exclude_text)
processing_elapsed = time.time() - processing_start_time
logger.debug(f"html_messages_processing_timing: channel {channel}, {len(final_posts)} posts processed in {processing_elapsed:.3f} seconds")
@@ -646,7 +698,11 @@ async def generate_channel_html(channel: str | int,
)
html = await asyncio.to_thread(_sanitize_sync, html)
except Exception as e:
# FAIL CLOSED (see the RSS path): `html` is now the raw concatenated feed
# (4.4 removed the per-fragment passes). If bleach throws, html.escape the
# whole feed rather than returning raw channel HTML/JS to the client.
logger.error(f"html_final_sanitization_error: channel {channel}, error {str(e)}")
html = html_escape(html)
html_gen_elapsed = time.time() - html_gen_start_time
logger.debug(f"html_generation_timing: channel {channel}, HTML generated in {html_gen_elapsed:.3f} seconds")
telegram_client.py
+17 -3
View File
@@ -259,6 +259,17 @@ class TelegramClient:
# Emergency termination
os._exit(1)
def watchdog_last_ok_age(self) -> float | None:
"""Seconds since the last successful watchdog probe, or None if none succeeded yet.
Reads only the already-recorded monotonic timestamp set by the watchdog loop; it
never issues a Telegram RPC, so it is safe to call from the hot /ping path even
while a real RPC is hung.
"""
if self._wd_last_ok_monotonic is None:
return None
return time.monotonic() - self._wd_last_ok_monotonic
async def safe_get_messages(self, channel_id, post_id, max_retries=2):
"""Wrapper with retry logic for auth errors"""
for attempt in range(max_retries):
@@ -274,13 +285,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:
tests/mock_config.py
+5
View File
@@ -33,4 +33,9 @@ def get_settings():
"tg_watchdog_heartbeat_every": 30,
"tg_disconnect_flap_limit": 3,
"tg_disconnect_flap_window": 120,
"tg_ping_unhealthy_after": 250,
"media_download_timeout_min": 120,
"media_download_timeout_max": 1800,
"media_download_min_speed": 256 * 1024,
"io_thread_pool_size": 32,
}
tests/test_stage2_static.py
+403
View File
@@ -0,0 +1,403 @@
# 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)
# --------------------------------------------------------------------------- #
# 2.5 — HTTP_DOWNLOAD_SEMAPHORE balance: released exactly once on success/error,
# never released when the acquire itself timed out (503). It is a plain
# asyncio.Semaphore, so an over-release would SILENTLY inflate the permit count
# and disable the limiter — assert the count, mirroring the stage-1 gate test.
# --------------------------------------------------------------------------- #
async def test_media_semaphore_released_on_download_error(monkeypatch, tmp_path):
api_server._inflight.clear()
monkeypatch.chdir(tmp_path) # pre-semaphore cache miss -> the acquire path runs
monkeypatch.setattr(api_server, "verify_media_digest", lambda url, digest: True)
permits_before = api_server.HTTP_DOWNLOAD_SEMAPHORE._value
async def boom(channel_id, post_id, fid):
raise RuntimeError("download blew up")
monkeypatch.setattr(api_server, "_download_deduped", boom)
req = SimpleNamespace(headers={})
# get_media swallows the error into a 4xx/5xx response; the point is the permit.
try:
await api_server.get_media("chan", 5, "errfid", request=req, digest="x")
except Exception:
pass
# Permit released exactly once (back to baseline) — not leaked, not double-released.
assert api_server.HTTP_DOWNLOAD_SEMAPHORE._value == permits_before
async def test_media_semaphore_not_released_on_acquire_timeout(monkeypatch, tmp_path):
api_server._inflight.clear()
monkeypatch.chdir(tmp_path)
monkeypatch.setattr(api_server, "verify_media_digest", lambda url, digest: True)
permits_before = api_server.HTTP_DOWNLOAD_SEMAPHORE._value
# Make the bounded acquire time out: the first wait_for (wrapping the acquire)
# raises TimeoutError; we close the passed coroutine to avoid a "never awaited"
# warning. The 503 short-circuits before _download_deduped, so only this one
# wait_for is hit.
real_wait_for = asyncio.wait_for
async def fake_wait_for(coro, timeout=None):
if hasattr(coro, "close"):
coro.close()
raise asyncio.TimeoutError()
monkeypatch.setattr(api_server.asyncio, "wait_for", fake_wait_for)
req = SimpleNamespace(headers={})
resp = await api_server.get_media("chan", 5, "busyfid", request=req, digest="x")
assert resp.status_code == 503
assert resp.headers.get("retry-after") == "30"
# A timed-out acquire never held a permit, so nothing must have been released:
# the count is UNCHANGED (an erroneous release would inflate it above baseline).
assert api_server.HTTP_DOWNLOAD_SEMAPHORE._value == permits_before
monkeypatch.setattr(api_server.asyncio, "wait_for", real_wait_for)
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
tests/test_stage3_fileresponse.py
+383
View File
@@ -0,0 +1,383 @@
# 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 3 (FileResponse + HTTP-layer cleanup) regression tests.
Covers:
- The HTTP Range matrix asserted against the FINAL (Starlette FileResponse) behavior,
each case documented, including the RFC-7233-permitted deltas vs the old hand-rolled
streaming (garbage header -> 400 not 416; multi-range -> proper multipart 206 not 416;
416 Content-Range now `*/size`; ETag/Last-Modified now present; ASCII filename no longer
gets a redundant filename*= form).
- Stage-2 behavior preserved through the rewrite: a served temp_* file still has its mtime
refreshed; delete_after still schedules the temp-file BackgroundTask (and it actually
runs); the media_key MIME cache is still consulted and populated.
- The pure-ASGI RequestLoggingMiddleware: a normal request still returns (body intact, not
buffered/truncated) and is logged.
Before/after Range matrix (2048-byte file), old = hand-rolled streaming, new = FileResponse:
request | old | new (FileResponse)
-----------------------+-----------------------------+-----------------------------------
no Range | 200 full | 200 full (+ETag/Last-Modified)
bytes=0-499 | 206 bytes 0-499/2048 | 206 bytes 0-499/2048 (same bytes)
bytes=500- | 206 bytes 500-2047/2048 | 206 bytes 500-2047/2048 (same)
bytes=-500 | 206 bytes 1548-2047/2048 | 206 bytes 1548-2047/2048 (same)
bytes=999999- (>EOF) | 416 `bytes */2048` | 416 `*/2048` (Starlette omits the
| | `bytes ` prefix; RFC-acceptable)
garbage header | 416 | 400 Bad Request (RFC 7233 lets a
| | server reject/ignore a malformed
| | Range; Starlette returns 400)
bytes=0-9,20-29 | 416 (old parser bug) | 206 multipart/byteranges (correct)
All 206 responses that carry data return byte-for-byte identical slices old vs new, so no
real regression is hidden behind an "accepted difference".
"""
import os
import sys
import time
import logging
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'])
from fastapi import FastAPI, Request
from fastapi.testclient import TestClient
import api_server
# --------------------------------------------------------------------------- #
# Helpers
# --------------------------------------------------------------------------- #
BODY = bytes(range(256)) * 8 # 2048 deterministic bytes
SIZE = len(BODY)
def _make_client(file_path, delete_after=False, media_key=None):
"""Mount prepare_file_response on a tiny app so it is driven through real ASGI
(FileResponse computes Range/206/416 at send time, so it must be exercised via a client)."""
app = FastAPI()
@app.get("/f")
async def _serve(request: Request):
return await api_server.prepare_file_response(
file_path, request=request, delete_after=delete_after, media_key=media_key
)
return TestClient(app)
@pytest.fixture
def sample_file(tmp_path):
fp = tmp_path / "myfile.bin"
fp.write_bytes(BODY)
return str(fp)
# --------------------------------------------------------------------------- #
# 3.1 — Range matrix against the FINAL FileResponse behavior.
# --------------------------------------------------------------------------- #
def test_no_range_returns_full_200(sample_file):
c = _make_client(sample_file)
r = c.get("/f")
assert r.status_code == 200
assert r.content == BODY
assert r.headers["content-length"] == str(SIZE)
assert r.headers["accept-ranges"] == "bytes"
assert r.headers["cache-control"] == "public, max-age=86400, immutable"
# Content-Disposition: inline, filename set by FileResponse itself (no manual double-set).
assert r.headers["content-disposition"].startswith("inline")
assert 'filename="myfile.bin"' in r.headers["content-disposition"]
# NEW vs old: FileResponse adds validators. Old streaming had neither.
assert r.headers.get("etag")
assert r.headers.get("last-modified")
def test_range_prefix_0_499(sample_file):
c = _make_client(sample_file)
r = c.get("/f", headers={"Range": "bytes=0-499"})
assert r.status_code == 206
assert r.headers["content-range"] == f"bytes 0-499/{SIZE}"
assert r.headers["content-length"] == "500"
assert r.content == BODY[:500] # exact slice, identical to old behavior
def test_range_open_ended_500_to_eof(sample_file):
c = _make_client(sample_file)
r = c.get("/f", headers={"Range": "bytes=500-"})
assert r.status_code == 206
assert r.headers["content-range"] == f"bytes 500-{SIZE - 1}/{SIZE}"
assert r.headers["content-length"] == str(SIZE - 500)
assert r.content == BODY[500:]
def test_range_suffix_last_500(sample_file):
c = _make_client(sample_file)
r = c.get("/f", headers={"Range": "bytes=-500"})
assert r.status_code == 206
assert r.headers["content-range"] == f"bytes {SIZE - 500}-{SIZE - 1}/{SIZE}"
assert r.headers["content-length"] == "500"
assert r.content == BODY[-500:]
def test_range_start_past_eof_returns_416(sample_file):
c = _make_client(sample_file)
r = c.get("/f", headers={"Range": "bytes=999999-"})
assert r.status_code == 416
# DELTA (RFC-acceptable): Starlette's 416 Content-Range is `*/size` (it omits the
# `bytes ` unit prefix the old code emitted). Same information, permitted by RFC 7233.
assert r.headers["content-range"] == f"*/{SIZE}"
def test_garbage_range_header_returns_400(sample_file):
c = _make_client(sample_file)
r = c.get("/f", headers={"Range": "somethinggarbage"})
# DELTA (RFC-acceptable): the old hand-rolled parser returned 416 on an unparseable
# header; Starlette rejects a malformed Range with 400 Bad Request. RFC 7233 permits a
# server to reject/ignore a bad Range; this is a conscious accepted difference, not a
# data regression (no bytes are served either way).
assert r.status_code == 400
def test_multi_range_returns_multipart_206(sample_file):
c = _make_client(sample_file)
r = c.get("/f", headers={"Range": "bytes=0-9,20-29"})
# DELTA (improvement): the old parser mis-split multi-ranges and returned 416.
# FileResponse serves a proper multipart/byteranges 206 containing both slices.
# Starlette 0.45.3 quirk: it advertises the multipart boundary via the `content-range`
# header (rather than content-type, which stays the file's own media_type); the body is
# a real multipart/byteranges document. We assert on the boundary marker + both slices.
assert r.status_code == 206
assert r.headers["content-range"].startswith("multipart/byteranges")
body = r.content
assert BODY[0:10] in body
assert BODY[20:30] in body
# --------------------------------------------------------------------------- #
# 3.1 (stage-2 preserved) — temp_* mtime touch survives the FileResponse swap.
# --------------------------------------------------------------------------- #
def test_temp_file_mtime_refreshed_when_stale(tmp_path):
# A temp_* file whose mtime is OLDER than the refresh interval gets touched, so the
# 1h sweeper won't delete an actively-viewed video.
temp_file = tmp_path / "temp_bigvid"
temp_file.write_bytes(b"videodata")
stale = time.time() - 10000 # >> TEMP_MTIME_REFRESH_INTERVAL (300s)
os.utime(temp_file, (stale, stale))
c = _make_client(str(temp_file))
r = c.get("/f")
assert r.status_code == 200
assert r.content == b"videodata"
assert os.path.getmtime(temp_file) > time.time() - 100 # refreshed by the serve
def test_temp_file_mtime_stable_when_fresh(tmp_path):
# DEBOUNCE: a temp_* file touched RECENTLY (within the refresh interval) is NOT
# re-touched — so FileResponse's mtime-derived ETag stays stable across the rapid
# requests of a single resume/seek session and `If-Range` resume keeps working.
temp_file = tmp_path / "temp_freshvid"
temp_file.write_bytes(b"videodata")
recent = time.time() - 30 # << TEMP_MTIME_REFRESH_INTERVAL (300s)
os.utime(temp_file, (recent, recent))
mtime_before = os.path.getmtime(temp_file)
c = _make_client(str(temp_file))
r1 = c.get("/f")
r2 = c.get("/f")
assert r1.status_code == 200 and r2.status_code == 200
# mtime unchanged -> the ETag is identical across the two serves.
assert os.path.getmtime(temp_file) == mtime_before
assert r1.headers.get("etag") == r2.headers.get("etag")
def test_non_temp_file_mtime_untouched(tmp_path):
normal = tmp_path / "regularfile"
normal.write_bytes(b"data")
stale = time.time() - 10000
os.utime(normal, (stale, stale))
c = _make_client(str(normal))
c.get("/f")
assert os.path.getmtime(normal) < time.time() - 5000 # left alone
# --------------------------------------------------------------------------- #
# 3.1 (stage-2 preserved) — delete_after still schedules the temp-file BackgroundTask.
# --------------------------------------------------------------------------- #
async def test_delete_after_attaches_background_task(tmp_path):
fp = tmp_path / "temp_todelete"
fp.write_bytes(b"x")
class _Req:
headers = {}
resp = await api_server.prepare_file_response(str(fp), request=_Req(), delete_after=True)
# FileResponse must carry a non-None background that deletes exactly this file.
assert resp.background is not None
assert resp.background.func is api_server.delayed_delete_file
assert resp.background.args == (str(fp),)
async def test_no_delete_after_has_no_background(tmp_path):
fp = tmp_path / "keepme"
fp.write_bytes(b"x")
class _Req:
headers = {}
resp = await api_server.prepare_file_response(str(fp), request=_Req(), delete_after=False)
assert resp.background is None
def test_delete_after_background_runs_and_removes_file(tmp_path, monkeypatch):
fp = tmp_path / "temp_gone"
fp.write_bytes(BODY)
# Patch the module-level deleter to remove immediately (real one sleeps 300s); the
# BackgroundTask picks up this reference at prepare_file_response call time.
async def _delete_now(path, delay=300):
os.remove(path)
monkeypatch.setattr(api_server, "delayed_delete_file", _delete_now)
c = _make_client(str(fp), delete_after=True)
r = c.get("/f")
assert r.status_code == 200
assert r.content == BODY
# TestClient blocks until the background task has run.
assert not os.path.exists(fp)
# --------------------------------------------------------------------------- #
# 3.1 (stage-2 preserved) — media_key MIME cache is consulted and populated.
# --------------------------------------------------------------------------- #
def test_media_key_mime_cache_hit_used(sample_file, monkeypatch):
calls = {"get": 0, "set": 0, "magic": 0}
def fake_get(db, ch, pid, fid):
calls["get"] += 1
return "video/mp4" # cache HIT
def fake_set(*a, **k):
calls["set"] += 1
def fake_magic(_path):
calls["magic"] += 1
return "application/octet-stream"
monkeypatch.setattr(api_server, "get_mime_type_sync", fake_get)
monkeypatch.setattr(api_server, "set_mime_type_sync", fake_set)
monkeypatch.setattr(api_server.magic_mime, "from_file", fake_magic)
c = _make_client(sample_file, media_key=("chan", 42, "fid"))
r = c.get("/f")
assert r.status_code == 200
# Cache hit: python-magic never invoked, nothing re-written to the cache, MIME applied.
assert calls["get"] == 1
assert calls["magic"] == 0
assert calls["set"] == 0
assert r.headers["content-type"].startswith("video/mp4")
def test_media_key_mime_cache_miss_populates(sample_file, monkeypatch):
written = {}
def fake_get(db, ch, pid, fid):
return None # cache MISS
def fake_set(db, ch, pid, fid, mime):
written["mime"] = mime
def fake_magic(_path):
return "image/png"
monkeypatch.setattr(api_server, "get_mime_type_sync", fake_get)
monkeypatch.setattr(api_server, "set_mime_type_sync", fake_set)
monkeypatch.setattr(api_server.magic_mime, "from_file", fake_magic)
c = _make_client(sample_file, media_key=("chan", 42, "fid"))
r = c.get("/f")
assert r.status_code == 200
# Miss -> python-magic result is both applied and persisted for next time.
assert written.get("mime") == "image/png"
assert r.headers["content-type"].startswith("image/png")
# --------------------------------------------------------------------------- #
# 3.1 — 404 pre-check kept.
# --------------------------------------------------------------------------- #
def test_missing_file_returns_404(tmp_path):
c = _make_client(str(tmp_path / "does_not_exist"))
r = c.get("/f")
assert r.status_code == 404
# --------------------------------------------------------------------------- #
# 3.2 — pure-ASGI RequestLoggingMiddleware: request returns intact and is logged.
# --------------------------------------------------------------------------- #
def test_asgi_logging_middleware_passes_body_and_logs(tmp_path, caplog):
fp = tmp_path / "streamed.bin"
fp.write_bytes(BODY)
app = FastAPI()
app.add_middleware(api_server.RequestLoggingMiddleware)
@app.get("/f")
async def _serve(request: Request):
return await api_server.prepare_file_response(str(fp), request=request)
client = TestClient(app)
with caplog.at_level(logging.DEBUG, logger="api_server"):
r = client.get("/f")
# Body flows straight through the middleware — not buffered/truncated.
assert r.status_code == 200
assert r.content == BODY
messages = " ".join(rec.getMessage() for rec in caplog.records)
assert "Request: GET /f" in messages
assert "Response status: 200" in messages
def test_asgi_logging_middleware_range_still_works(tmp_path):
fp = tmp_path / "streamed.bin"
fp.write_bytes(BODY)
app = FastAPI()
app.add_middleware(api_server.RequestLoggingMiddleware)
@app.get("/f")
async def _serve(request: Request):
return await api_server.prepare_file_response(str(fp), request=request)
client = TestClient(app)
r = client.get("/f", headers={"Range": "bytes=0-99"})
# 206 streaming still works through the pure-ASGI middleware (send not buffered).
assert r.status_code == 206
assert r.content == BODY[:100]
async def test_asgi_logging_middleware_passes_non_http_scope_through(caplog):
# The `scope["type"] != "http"` branch (lifespan/websocket) only runs in a real
# deploy, never through TestClient's plain request path — so pin it directly:
# a non-http scope must delegate to the inner app untouched and log nothing.
called = {}
async def inner(scope, receive, send):
called["scope_type"] = scope["type"]
mw = api_server.RequestLoggingMiddleware(inner)
with caplog.at_level(logging.DEBUG, logger="api_server"):
await mw({"type": "lifespan"}, None, None)
assert called["scope_type"] == "lifespan" # delegated to the inner app
# Nothing request/response-ish was logged for a non-http scope.
messages = " ".join(rec.getMessage() for rec in caplog.records)
assert "Request:" not in messages
assert "Response status:" not in messages
tests/test_stage4_eventloop.py
+471
View File
@@ -0,0 +1,471 @@
# 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 4 (event-loop hygiene for feed generation) tests.
Covers:
- 4.1 raw_message laziness: feeds do NOT compute str(message); JSON/debug HTML do.
- 4.2 side-effect IO removed from process_message: _save_media_file_ids only appends to
self._pending_media_ids; the caller flushes once via upsert_media_file_ids_bulk_sync.
Also: the render path contains NO create_task / get_running_loop / to_thread.
- 4.3 render pipeline moved into a single thread: the four render functions are now plain
sync functions and actually execute off the main thread; deepcopy of a pickled Message
does not crash; a 100-message feed generates correctly.
- 4.4 sanitize coverage (XSS): a <script> / onerror= / javascript: payload is stripped in
ALL outputs — rss, html-feed, single-post html, and json — each with exactly one pass.
"""
import os
import re
import sys
import copy
import pickle
import asyncio
import threading
from datetime import datetime, timezone
from types import SimpleNamespace
import pytest
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
sys.modules['config'] = __import__('tests.mock_config', fromlist=['get_settings'])
from pyrogram.enums import MessageMediaType
import post_parser as pp_module
import rss_generator as rss_module
from post_parser import PostParser
from rss_generator import (
generate_channel_rss,
generate_channel_html,
_render_pipeline,
_create_time_based_media_groups,
_create_messages_groups,
_trim_messages_groups,
_render_messages_groups,
)
XSS_PAYLOAD = "<script>alert('xss')</script><img src=x onerror=\"alert(1)\"><a href=\"javascript:alert(2)\">click</a>"
class _Str(str):
"""Minimal stand-in for Pyrogram's Str: .html returns the raw string unchanged,
so a malicious payload reaches the pre-sanitization html body just like real text
carrying entities would."""
@property
def html(self):
return str(self)
def make_message(mid, text=None, media=None, photo_uid=None, username="testchan",
date=None):
m = SimpleNamespace()
m.id = mid
m.date = date or datetime(2024, 1, 1, 12, 0, mid % 60, tzinfo=timezone.utc)
m.text = _Str(text) if text is not None else None
m.caption = None
m.media = media
m.web_page = None
m.poll = None
m.service = None
m.forward_origin = None
m.reply_to_message = None
m.reply_to_message_id = None
m.sender_chat = None
m.from_user = None
m.reactions = None
m.views = 100
m.media_group_id = None
m.show_caption_above_media = False
m.chat = SimpleNamespace(id=-1001234567890, username=username)
# media sub-objects default to None
for attr in ("photo", "video", "document", "audio", "voice",
"video_note", "animation", "sticker"):
setattr(m, attr, None)
if media == MessageMediaType.PHOTO and photo_uid:
m.photo = SimpleNamespace(file_unique_id=photo_uid)
return m
def _co_names(func):
return set(func.__code__.co_names)
# ---------------------------------------------------------------------------
# 4.3 — render functions are plain sync and run off the main thread
# ---------------------------------------------------------------------------
def test_render_functions_are_sync():
for fn in (_create_time_based_media_groups, _create_messages_groups,
_trim_messages_groups, _render_messages_groups, _render_pipeline):
assert not asyncio.iscoroutinefunction(fn), f"{fn.__name__} must be a plain sync function"
@pytest.mark.asyncio
async def test_pipeline_runs_in_worker_thread(monkeypatch):
main_ident = threading.get_ident()
seen = {}
real_render = rss_module._render_messages_groups
def spy(*args, **kwargs):
seen["ident"] = threading.get_ident()
return real_render(*args, **kwargs)
monkeypatch.setattr(rss_module, "_render_messages_groups", spy)
async def fake_get_chat(client, channel):
return SimpleNamespace(title="Test", username="testchan", id=-1001234567890)
async def fake_get_history(client, channel, limit=20):
return [make_message(i, text=f"post {i}") for i in range(5)]
monkeypatch.setattr("tg_cache.cached_get_chat", fake_get_chat, raising=False)
monkeypatch.setattr("tg_cache.cached_get_chat_history", fake_get_history, raising=False)
await generate_channel_rss("testchan", client=SimpleNamespace(), limit=5)
assert "ident" in seen
assert seen["ident"] != main_ident, "render pipeline must run in a worker thread, not the loop thread"
def test_deepcopy_of_pickled_message_does_not_crash():
from pyrogram.types import Message, Chat
from pyrogram.enums import ChatType
m = Message(id=7, date=datetime.now(timezone.utc), text="hello",
chat=Chat(id=-1001, type=ChatType.CHANNEL, username="testchan"))
roundtripped = pickle.loads(pickle.dumps(m)) # mimics the pickle cache
clone = copy.deepcopy(roundtripped) # what _create_time_based_media_groups does
assert clone.id == 7
assert clone.chat.username == "testchan"
# ---------------------------------------------------------------------------
# 4.2 — no asyncio in the render path; bulk upsert after render
# ---------------------------------------------------------------------------
def test_render_path_has_no_asyncio_side_effects():
banned = {"create_task", "get_running_loop", "to_thread", "ensure_future"}
funcs = [
_render_pipeline, _create_time_based_media_groups, _create_messages_groups,
_trim_messages_groups, _render_messages_groups,
PostParser.process_message, PostParser._generate_html_body,
PostParser._generate_html_media, PostParser.generate_html_footer,
PostParser._reactions_views_links, PostParser._save_media_file_ids,
PostParser._sanitize_html,
]
for fn in funcs:
offenders = _co_names(fn) & banned
assert not offenders, f"{fn.__qualname__} references forbidden asyncio names: {offenders}"
@pytest.mark.asyncio
async def test_media_ids_persisted_via_bulk_upsert(monkeypatch):
calls = []
def fake_bulk(db_path, entries):
calls.append(list(entries))
monkeypatch.setattr(pp_module, "upsert_media_file_ids_bulk_sync", fake_bulk)
async def fake_get_chat(client, channel):
return SimpleNamespace(title="Test", username="testchan", id=-1001234567890)
async def fake_get_history(client, channel, limit=20):
return [
make_message(1, text="just text"),
make_message(2, media=MessageMediaType.PHOTO, photo_uid="uid_abc"),
]
monkeypatch.setattr("tg_cache.cached_get_chat", fake_get_chat, raising=False)
monkeypatch.setattr("tg_cache.cached_get_chat_history", fake_get_history, raising=False)
await generate_channel_rss("testchan", client=SimpleNamespace(), limit=10)
assert len(calls) == 1, "bulk upsert must be called exactly once after render"
entries = calls[0]
assert len(entries) == 1, "only the photo message contributes a media id"
channel, post_id, fid, _ts = entries[0]
assert (channel, post_id, fid) == ("testchan", 2, "uid_abc")
@pytest.mark.asyncio
async def test_save_media_file_ids_only_appends(monkeypatch):
# Even with a running loop, _save_media_file_ids must not create tasks — just append.
parser = PostParser(SimpleNamespace())
monkeypatch.setattr(pp_module, "upsert_media_file_ids_bulk_sync",
lambda *a, **k: (_ for _ in ()).throw(AssertionError("must not be called directly")))
msg = make_message(3, media=MessageMediaType.PHOTO, photo_uid="uid_x")
parser._save_media_file_ids(msg)
assert parser._pending_media_ids == [("testchan", 3, "uid_x", parser._pending_media_ids[0][3])]
# ---------------------------------------------------------------------------
# 4.1 — raw_message laziness
# ---------------------------------------------------------------------------
def test_raw_message_lazy_for_feed():
parser = PostParser(SimpleNamespace())
result = parser.process_message(make_message(10, text="hi"), include_raw=False, sanitize=False)
assert "raw_message" not in result
def test_raw_message_present_for_json_and_debug():
parser = PostParser(SimpleNamespace())
result = parser.process_message(make_message(11, text="hi"), include_raw=True)
assert "raw_message" in result
assert isinstance(result["raw_message"], str)
@pytest.mark.asyncio
async def test_100_message_feed_generates(monkeypatch):
async def fake_get_chat(client, channel):
return SimpleNamespace(title="Test", username="testchan", id=-1001234567890)
async def fake_get_history(client, channel, limit=20):
return [make_message(i, text=f"post number {i}") for i in range(1, 101)]
monkeypatch.setattr("tg_cache.cached_get_chat", fake_get_chat, raising=False)
monkeypatch.setattr("tg_cache.cached_get_chat_history", fake_get_history, raising=False)
rss = await generate_channel_rss("testchan", client=SimpleNamespace(), limit=100)
assert rss.count("<item>") == 100
assert "post number 50" in rss
# ---------------------------------------------------------------------------
# 4.4 — XSS: payload stripped in all four outputs
# ---------------------------------------------------------------------------
def _assert_clean(html_str, where):
assert "<script>" not in html_str, f"{where}: <script> survived"
assert "onerror" not in html_str, f"{where}: onerror= survived"
assert "javascript:" not in html_str, f"{where}: javascript: survived"
def _make_client_returning(msg):
client = SimpleNamespace()
async def get_messages(channel, post_id):
return msg
client.get_messages = get_messages
return client
@pytest.mark.asyncio
async def test_xss_stripped_in_json():
msg = make_message(20, text=XSS_PAYLOAD)
parser = PostParser(_make_client_returning(msg))
data = await parser.get_post("testchan", 20, "json")
_assert_clean(data["html"]["body"], "json body")
_assert_clean(data["html"]["footer"], "json footer")
@pytest.mark.asyncio
async def test_xss_stripped_in_single_post_html():
msg = make_message(21, text=XSS_PAYLOAD)
parser = PostParser(_make_client_returning(msg))
html_out = await parser.get_post("testchan", 21, "html")
_assert_clean(html_out, "single-post html")
@pytest.mark.asyncio
async def test_xss_stripped_in_single_post_html_debug():
# debug HTML embeds raw_message (str(message)) into a <pre>; it must be html-escaped so
# no live tag from the payload survives. The escaped dump legitimately still contains the
# inert words "onerror"/"javascript:" as text — what matters is that they are NOT live.
msg = make_message(22, text=XSS_PAYLOAD)
parser = PostParser(_make_client_returning(msg))
html_out = await parser.get_post("testchan", 22, "html", debug=True)
# 1) The rendered display area (everything before the raw <pre>) is fully sanitized.
display = html_out.split('<pre', 1)[0]
_assert_clean(display, "single-post debug display")
# 2) The raw <pre> dump is html-escaped: no live <script> tag anywhere, and the payload
# appears only in escaped form (proving html.escape ran).
assert "<script>" not in html_out, "debug raw dump left a live <script> tag"
assert "&lt;script&gt;" in html_out, "debug raw dump was not html-escaped"
@pytest.mark.asyncio
async def test_xss_stripped_in_rss_feed(monkeypatch):
async def fake_get_chat(client, channel):
return SimpleNamespace(title="Test", username="testchan", id=-1001234567890)
async def fake_get_history(client, channel, limit=20):
return [make_message(23, text=XSS_PAYLOAD)]
monkeypatch.setattr("tg_cache.cached_get_chat", fake_get_chat, raising=False)
monkeypatch.setattr("tg_cache.cached_get_chat_history", fake_get_history, raising=False)
rss = await generate_channel_rss("testchan", client=SimpleNamespace(), limit=5)
# The sanitized HTML lives in <content:encoded><![CDATA[...]]></content:encoded>.
cdata = re.findall(r"<!\[CDATA\[(.*?)\]\]>", rss, re.DOTALL)
assert cdata, "expected CDATA content in RSS"
for chunk in cdata:
_assert_clean(chunk, "rss content")
@pytest.mark.asyncio
async def test_xss_stripped_in_html_feed(monkeypatch):
async def fake_get_chat(client, channel):
return SimpleNamespace(title="Test", username="testchan", id=-1001234567890)
async def fake_get_history(client, channel, limit=20):
return [make_message(24, text=XSS_PAYLOAD)]
monkeypatch.setattr("tg_cache.cached_get_chat", fake_get_chat, raising=False)
monkeypatch.setattr("tg_cache.cached_get_chat_history", fake_get_history, raising=False)
html_feed = await generate_channel_html("testchan", client=SimpleNamespace(), limit=5)
_assert_clean(html_feed, "html feed")
def _media_msg_with_payload_caption(mid):
# A photo message whose CAPTION carries the payload — this exercises the media
# fragment path (_generate_html_media / caption rendering) whose internal per-fragment
# sanitize pass was removed in 4.4. The covering pass must still strip it.
m = make_message(mid, media=MessageMediaType.PHOTO, photo_uid="pic123")
m.caption = _Str(XSS_PAYLOAD)
return m
@pytest.mark.asyncio
async def test_xss_in_media_caption_stripped_direct_paths():
# json + single-post html go through process_message(sanitize=True) directly.
parser = PostParser(_make_client_returning(_media_msg_with_payload_caption(30)))
data = await parser.get_post("testchan", 30, "json")
_assert_clean(data["html"]["body"], "json body (media caption)")
_assert_clean(data["html"]["footer"], "json footer (media caption)")
parser2 = PostParser(_make_client_returning(_media_msg_with_payload_caption(31)))
html_out = await parser2.get_post("testchan", 31, "html")
_assert_clean(html_out, "single-post html (media caption)")
@pytest.mark.asyncio
async def test_xss_in_media_caption_stripped_in_feeds(monkeypatch):
async def fake_get_chat(client, channel):
return SimpleNamespace(title="Test", username="testchan", id=-1001234567890)
async def fake_get_history(client, channel, limit=20):
return [_media_msg_with_payload_caption(32)]
monkeypatch.setattr("tg_cache.cached_get_chat", fake_get_chat, raising=False)
monkeypatch.setattr("tg_cache.cached_get_chat_history", fake_get_history, raising=False)
rss = await generate_channel_rss("testchan", client=SimpleNamespace(), limit=5)
for chunk in re.findall(r"<!\[CDATA\[(.*?)\]\]>", rss, re.DOTALL):
_assert_clean(chunk, "rss content (media caption)")
async def fake_get_history2(client, channel, limit=20):
return [_media_msg_with_payload_caption(33)]
monkeypatch.setattr("tg_cache.cached_get_chat_history", fake_get_history2, raising=False)
html_feed = await generate_channel_html("testchan", client=SimpleNamespace(), limit=5)
_assert_clean(html_feed, "html feed (media caption)")
# --------------------------------------------------------------------------- #
# Review round-1: the new bulk-upsert SQL executed for real (not mocked).
# --------------------------------------------------------------------------- #
def test_bulk_upsert_media_file_ids_real_sql(tmp_path):
import sqlite3
from file_io import upsert_media_file_ids_bulk_sync, init_db_sync
db = str(tmp_path / "t.db")
init_db_sync(db)
# Empty list is a no-op (no crash).
upsert_media_file_ids_bulk_sync(db, [])
# Multi-row insert.
upsert_media_file_ids_bulk_sync(db, [
("chA", 1, "fidA", 100.0),
("chB", 2, "fidB", 200.0),
])
conn = sqlite3.connect(db)
rows = dict(((c, p, f), a) for c, p, f, a in
conn.execute("SELECT channel, post_id, file_unique_id, added FROM media_file_ids"))
assert rows[("chA", 1, "fidA")] == 100.0
assert rows[("chB", 2, "fidB")] == 200.0
# Re-upsert the SAME key updates `added` (ON CONFLICT ... DO UPDATE SET added=excluded.added).
upsert_media_file_ids_bulk_sync(db, [("chA", 1, "fidA", 999.0)])
a = conn.execute(
"SELECT added FROM media_file_ids WHERE channel='chA' AND post_id=1 AND file_unique_id='fidA'"
).fetchone()[0]
assert a == 999.0
conn.close()
# --------------------------------------------------------------------------- #
# Review round-1: media ids collected before a render exception are still
# flushed (the flush is in a finally). Removing the finally must break this.
# --------------------------------------------------------------------------- #
@pytest.mark.asyncio
async def test_pending_media_flushed_on_render_exception(monkeypatch):
async def fake_get_chat(client, channel):
return SimpleNamespace(title="Test", username="testchan", id=-1001234567890)
async def fake_get_history(client, channel, limit=20):
return [make_message(40, text="hi")]
monkeypatch.setattr("tg_cache.cached_get_chat", fake_get_chat, raising=False)
monkeypatch.setattr("tg_cache.cached_get_chat_history", fake_get_history, raising=False)
# A render pipeline that collects a pending id then raises mid-render.
def boom_pipeline(messages, post_parser, *a, **k):
post_parser._pending_media_ids.append(("chZ", 9, "fidZ", 1.0))
raise RuntimeError("render blew up")
monkeypatch.setattr(rss_module, "_render_pipeline", boom_pipeline)
flushed = {}
async def fake_bulk(db, entries):
flushed["entries"] = list(entries)
monkeypatch.setattr("post_parser.upsert_media_file_ids_bulk_sync",
lambda db, entries: flushed.__setitem__("entries", list(entries)),
raising=False)
with pytest.raises(Exception):
await generate_channel_rss("testchan", client=SimpleNamespace(), limit=5)
# The collected id was persisted despite the render exception (flush in finally).
assert flushed.get("entries") == [("chZ", 9, "fidZ", 1.0)]
# --------------------------------------------------------------------------- #
# Review round-1 [security]: if the ONLY sanitize pass throws, the feed must
# FAIL CLOSED (html.escape the raw content), never emit the raw XSS payload.
# --------------------------------------------------------------------------- #
@pytest.mark.asyncio
async def test_rss_fails_closed_when_sanitizer_raises(monkeypatch):
async def fake_get_chat(client, channel):
return SimpleNamespace(title="Test", username="testchan", id=-1001234567890)
async def fake_get_history(client, channel, limit=20):
return [make_message(41, text=XSS_PAYLOAD)]
monkeypatch.setattr("tg_cache.cached_get_chat", fake_get_chat, raising=False)
monkeypatch.setattr("tg_cache.cached_get_chat_history", fake_get_history, raising=False)
# Force bleach to blow up (e.g. the RecursionError class already seen in prod).
# rss_generator imports it as `from bleach import clean as HTMLSanitizer`.
def boom(*a, **k):
raise RecursionError("bleach exploded")
monkeypatch.setattr(rss_module, "HTMLSanitizer", boom, raising=True)
rss = await generate_channel_rss("testchan", client=SimpleNamespace(), limit=5)
chunks = re.findall(r"<!\[CDATA\[(.*?)\]\]>", rss, re.DOTALL)
assert chunks, "expected CDATA content"
for chunk in chunks:
# Fail-closed: the raw payload was html.escaped, so NO live tag survived — every
# `<` became `&lt;`. (The letters "javascript:" still appear, but as inert text
# inside an escaped &quot;…&quot;, not a live href.)
assert "<script" not in chunk, "RSS fail-open: raw <script> reached the feed"
assert "<img" not in chunk, "RSS fail-open: raw <img onerror> reached the feed"
assert "<a " not in chunk, "RSS fail-open: raw <a href> reached the feed"
# ...and the escaping actually ran (payload present as escaped text, not dropped).
assert "&lt;script&gt;" in chunk, "expected the payload html-escaped, not dropped"
tests/test_stage5_sqlite.py
+263
View File
@@ -0,0 +1,263 @@
# 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 5 (SQLite access-time batching) tests.
Covers:
- 5.1 accumulator: a /media cache hit records into api_server._access_updates and does NOT
call update_media_file_access_sync (no synchronous SQLite write on the hot path).
- 5.1 flush: seeding _access_updates + running the flush once writes the accumulated
timestamps to a real temp DB (hit -> flush -> value updated in DB), and clears the dict.
- 5.1 bulk fn update_media_file_access_bulk_sync: empty no-op, multi-row, and updating an
EXISTING row's `added` (WHERE matches on the str channel key).
- 5.1 snapshot-then-clear: an update arriving DURING the flush lands in the fresh dict and
is not lost.
- gotcha: str(channel) key discipline — an int-ish channel on the hot path keys the
accumulator (and thus the UPDATE) by the string form.
"""
import os
import sys
import sqlite3
import asyncio
from types import SimpleNamespace
import pytest
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 file_io import init_db_sync, update_media_file_access_bulk_sync
def _fake_plain_message():
"""A non-poll, non-video message so download_media_file takes the normal cache flow."""
return SimpleNamespace(media=None, video=None, empty=False)
@pytest.fixture(autouse=True)
def _clean_accumulator():
"""Each test starts with an empty accumulator and restores it afterwards."""
api_server._access_updates = {}
yield
api_server._access_updates = {}
# --------------------------------------------------------------------------- #
# 5.1 hot path: cache hit records into the accumulator, no synchronous SQLite.
# --------------------------------------------------------------------------- #
@pytest.mark.asyncio
async def test_cache_hit_records_accumulator_no_sqlite(tmp_path, monkeypatch):
monkeypatch.chdir(tmp_path)
channel, post_id, fid = "testchan", 5, "fidHIT"
cache_dir = tmp_path / "data" / "cache" / channel / str(post_id)
cache_dir.mkdir(parents=True)
(cache_dir / fid).write_bytes(b"cached-bytes")
async def fake_get(_cid, _pid):
return _fake_plain_message()
monkeypatch.setattr(api_server.client, "safe_get_messages", fake_get)
# Spy: the single-row synchronous updater must NOT be called on the hot path.
called = []
monkeypatch.setattr(api_server, "update_media_file_access_sync",
lambda *a, **k: called.append(a))
path, delete_after = await api_server.download_media_file(channel, post_id, fid)
assert path == str(cache_dir / fid)
assert delete_after is False
assert called == [] # no synchronous SQLite access-write happened
assert (channel, post_id, fid) in api_server._access_updates
# --------------------------------------------------------------------------- #
# 5.1 hot path (DoD guard): get_media's pre-semaphore cache-hit — the hottest
# changed site, the one this PR exists for — records into the accumulator and
# does NO synchronous SQLite access-write. Mirrors the download_media_file spy
# so a regression re-introducing a per-hit write into THIS branch goes red.
# --------------------------------------------------------------------------- #
@pytest.mark.asyncio
async def test_get_media_pre_semaphore_cache_hit_no_sqlite(tmp_path, monkeypatch):
monkeypatch.chdir(tmp_path)
channel, post_id, fid = "gmchan", 11, "fidGM"
cache_dir = tmp_path / "data" / "cache" / channel / str(post_id)
cache_dir.mkdir(parents=True)
(cache_dir / fid).write_bytes(b"cached-bytes")
# Bypass the HMAC digest gate and the FileResponse machinery — the test targets
# only the access-time write on the pre-semaphore cache-hit branch.
monkeypatch.setattr(api_server, "verify_media_digest", lambda *a, **k: True)
sentinel = object()
async def fake_prepare(cache_path, request=None, media_key=None):
return sentinel
monkeypatch.setattr(api_server, "prepare_file_response", fake_prepare)
# Spy: the single-row synchronous updater must NOT be called on the hot path.
called = []
monkeypatch.setattr(api_server, "update_media_file_access_sync",
lambda *a, **k: called.append(a))
resp = await api_server.get_media(channel, post_id, fid, request=object(), digest="x")
assert resp is sentinel
assert called == [] # no synchronous SQLite access-write on the hottest path
assert (channel, post_id, fid) in api_server._access_updates
# --------------------------------------------------------------------------- #
# gotcha: str(channel) key discipline on the hot path (int-ish channel).
# --------------------------------------------------------------------------- #
@pytest.mark.asyncio
async def test_cache_hit_keys_channel_as_str(tmp_path, monkeypatch):
monkeypatch.chdir(tmp_path)
channel_int, post_id, fid = 12345, 7, "fidINT"
cache_dir = tmp_path / "data" / "cache" / str(channel_int) / str(post_id)
cache_dir.mkdir(parents=True)
(cache_dir / fid).write_bytes(b"x")
async def fake_get(_cid, _pid):
return _fake_plain_message()
monkeypatch.setattr(api_server.client, "safe_get_messages", fake_get)
await api_server.download_media_file(channel_int, post_id, fid)
assert ("12345", post_id, fid) in api_server._access_updates # string form
assert (channel_int, post_id, fid) not in api_server._access_updates # never the raw int
# --------------------------------------------------------------------------- #
# 5.1 flush: hit -> flush -> the accumulated timestamp is written to the DB.
# --------------------------------------------------------------------------- #
@pytest.mark.asyncio
async def test_flush_writes_accumulated_timestamps(tmp_path, monkeypatch):
db = str(tmp_path / "flush.db")
init_db_sync(db)
# Seed two existing rows with an old timestamp.
conn = sqlite3.connect(db)
conn.executemany(
"INSERT INTO media_file_ids (channel, post_id, file_unique_id, added) VALUES (?,?,?,?)",
[("chA", 1, "fA", 1.0), ("chB", 2, "fB", 2.0)],
)
conn.commit()
conn.close()
monkeypatch.setattr(api_server, "DB_PATH", db)
api_server._access_updates = {
("chA", 1, "fA"): 111.0,
("chB", 2, "fB"): 222.0,
}
await api_server._flush_access_updates()
# Dict cleared after flush.
assert api_server._access_updates == {}
conn = sqlite3.connect(db)
rows = dict(((c, p, f), a) for c, p, f, a in
conn.execute("SELECT channel, post_id, file_unique_id, added FROM media_file_ids"))
conn.close()
assert rows[("chA", 1, "fA")] == 111.0
assert rows[("chB", 2, "fB")] == 222.0
@pytest.mark.asyncio
async def test_flush_empty_is_noop(monkeypatch):
api_server._access_updates = {}
# Must not raise and must not touch the threadpool/DB.
monkeypatch.setattr(api_server, "update_media_file_access_bulk_sync",
lambda *a, **k: pytest.fail("bulk sync should not run for an empty batch"))
await api_server._flush_access_updates()
# --------------------------------------------------------------------------- #
# 5.1 snapshot-then-clear: an update arriving DURING the flush is not lost.
# --------------------------------------------------------------------------- #
@pytest.mark.asyncio
async def test_snapshot_then_clear_does_not_lose_late_update(tmp_path, monkeypatch):
db = str(tmp_path / "race.db")
init_db_sync(db)
monkeypatch.setattr(api_server, "DB_PATH", db)
late_key = ("chLate", 9, "fLate")
def fake_bulk(_db, entries):
# Simulates a cache hit landing WHILE the flush's to_thread runs: because the flush
# already replaced the module dict with a fresh one, this write goes into the NEW dict.
api_server._access_updates[late_key] = 999.0
monkeypatch.setattr(api_server, "update_media_file_access_bulk_sync", fake_bulk)
api_server._access_updates = {("chA", 1, "fA"): 111.0}
await api_server._flush_access_updates()
# The snapshot (chA) was flushed and cleared; the late update survives in the new dict.
assert api_server._access_updates == {late_key: 999.0}
# --------------------------------------------------------------------------- #
# 5.1 re-queue on failure: a failed bulk write is not lost; setdefault keeps a
# fresher write that arrived during the flush.
# --------------------------------------------------------------------------- #
@pytest.mark.asyncio
async def test_failed_flush_requeues_batch_without_clobbering_fresh(tmp_path, monkeypatch):
db = str(tmp_path / "fail.db")
init_db_sync(db)
monkeypatch.setattr(api_server, "DB_PATH", db)
stale_key = ("chStale", 1, "fS") # only in the failing snapshot
fresh_key = ("chFresh", 2, "fF") # re-written FRESHER during the flush
def fake_bulk(_db, _entries):
# A newer cache hit for fresh_key lands in the fresh dict WHILE the bulk write runs,
# then the write fails. The re-queue must restore stale_key but must NOT overwrite
# the newer fresh_key value (setdefault, not assignment).
api_server._access_updates[fresh_key] = 999.0
raise sqlite3.OperationalError("database is locked")
monkeypatch.setattr(api_server, "update_media_file_access_bulk_sync", fake_bulk)
api_server._access_updates = {stale_key: 1.0, fresh_key: 2.0}
with pytest.raises(sqlite3.OperationalError):
await api_server._flush_access_updates()
# stale_key restored with its snapshot value; fresh_key keeps the NEWER value (not 2.0).
assert api_server._access_updates == {stale_key: 1.0, fresh_key: 999.0}
# --------------------------------------------------------------------------- #
# 5.1 bulk fn: empty no-op, multi-row, and existing-row update via str channel key.
# --------------------------------------------------------------------------- #
def test_bulk_access_update_real_sql(tmp_path):
db = str(tmp_path / "bulk.db")
init_db_sync(db)
# Empty batch is a no-op (no crash).
update_media_file_access_bulk_sync(db, [])
# Seed existing rows.
conn = sqlite3.connect(db)
conn.executemany(
"INSERT INTO media_file_ids (channel, post_id, file_unique_id, added) VALUES (?,?,?,?)",
[("chA", 1, "fA", 1.0), ("chB", 2, "fB", 2.0)],
)
conn.commit()
conn.close()
# Multi-row update of EXISTING rows: WHERE matches on the str channel key.
update_media_file_access_bulk_sync(db, [
("chA", 1, "fA", 500.0),
("chB", 2, "fB", 600.0),
])
conn = sqlite3.connect(db)
rows = dict(((c, p, f), a) for c, p, f, a in
conn.execute("SELECT channel, post_id, file_unique_id, added FROM media_file_ids"))
# A row keyed by an int channel does NOT match the string "chA" WHERE (documents the gotcha).
n = conn.execute("SELECT COUNT(*) FROM media_file_ids WHERE channel = 1").fetchone()[0]
conn.close()
assert rows[("chA", 1, "fA")] == 500.0
assert rows[("chB", 2, "fB")] == 600.0
assert n == 0
tests/test_stage6_healthcheck.py
+188
View File
@@ -0,0 +1,188 @@
# 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 6 (lightweight /ping healthcheck) regression tests.
Covers:
- /ping returns 200 "ok" when connected and the last probe is recent (age < threshold).
- /ping returns 503 "degraded" when connected but the last probe is stale (age > threshold).
- /ping returns 503 "degraded" when disconnected, regardless of probe age.
- /ping returns 200 "ok" on a fresh boot (age is None) while connected — a freshly-started
container must NOT be killed before the watchdog's first probe.
- ANTI-REGRESSION (the critical invariant): /ping issues ZERO Telegram RPC. A spy on the
fake client's get_me / safe_get_messages proves neither is ever called.
- TelegramClient.watchdog_last_ok_age(): None when never probed; a positive float afterwards.
"""
import os
import sys
import time
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'])
from fastapi.testclient import TestClient
import api_server
from telegram_client import TelegramClient
# --------------------------------------------------------------------------- #
# Fakes
# --------------------------------------------------------------------------- #
class _FakeKurigramClient:
"""Stands in for TelegramClient.client — exposes is_connected and RPC spies."""
def __init__(self, is_connected=True):
self.is_connected = is_connected
self.get_me_calls = 0
async def get_me(self):
# If /ping ever touches this, the whole point of the endpoint is defeated.
self.get_me_calls += 1
raise AssertionError("/ping must never call get_me()")
class _FakeTelegramClient:
"""Stands in for api_server.client with a controllable probe age + RPC spies."""
def __init__(self, age, is_connected=True):
self._age = age
self.client = _FakeKurigramClient(is_connected=is_connected)
self.safe_get_messages_calls = 0
def watchdog_last_ok_age(self):
return self._age
async def safe_get_messages(self, *args, **kwargs):
self.safe_get_messages_calls += 1
raise AssertionError("/ping must never call safe_get_messages()")
@pytest.fixture
def patch_client(monkeypatch):
"""Return a factory that installs a fake api_server.client and yields a TestClient."""
def _install(age, is_connected=True):
fake = _FakeTelegramClient(age=age, is_connected=is_connected)
monkeypatch.setattr(api_server, "client", fake)
return fake, TestClient(api_server.app)
return _install
THRESHOLD = api_server.Config["tg_ping_unhealthy_after"] # 250 in mock_config
# --------------------------------------------------------------------------- #
# /ping endpoint behavior
# --------------------------------------------------------------------------- #
def test_ping_healthy_connected_recent(patch_client):
fake, tc = patch_client(age=THRESHOLD - 10, is_connected=True)
r = tc.get("/ping")
assert r.status_code == 200
body = r.json()
assert body["status"] == "ok"
assert body["connected"] is True
assert body["last_probe_age_s"] == round(THRESHOLD - 10, 1)
assert body["threshold_s"] == THRESHOLD
assert fake.client.get_me_calls == 0
def test_ping_degraded_stale_probe(patch_client):
fake, tc = patch_client(age=THRESHOLD + 100, is_connected=True)
r = tc.get("/ping")
assert r.status_code == 503
body = r.json()
assert body["status"] == "degraded"
assert body["connected"] is True
assert fake.client.get_me_calls == 0
def test_ping_watchdog_disabled_stale_age_still_healthy(patch_client, monkeypatch):
# With the watchdog OFF, nothing refreshes age (a disconnect-flap restart can stamp it
# once, then it only grows). A stale age must NOT drive /ping to 503 on a live connection
# — that would spuriously fail the healthcheck and trigger an autoheal restart. So with the
# watchdog disabled, /ping is a pure connectivity check: connected + stale age => healthy.
monkeypatch.setitem(api_server.Config, "tg_watchdog_enabled", False)
fake, tc = patch_client(age=THRESHOLD + 100, is_connected=True)
r = tc.get("/ping")
assert r.status_code == 200
body = r.json()
assert body["status"] == "ok"
assert body["connected"] is True
assert fake.client.get_me_calls == 0
def test_ping_degraded_disconnected(patch_client):
# Even with a fresh probe age, a disconnected client is unhealthy.
fake, tc = patch_client(age=1.0, is_connected=False)
r = tc.get("/ping")
assert r.status_code == 503
body = r.json()
assert body["status"] == "degraded"
assert body["connected"] is False
def test_ping_degraded_disconnected_even_when_age_none(patch_client):
fake, tc = patch_client(age=None, is_connected=False)
r = tc.get("/ping")
assert r.status_code == 503
assert r.json()["status"] == "degraded"
def test_ping_fresh_boot_age_none_is_healthy(patch_client):
# Right after boot the watchdog hasn't probed yet (age None); connected => healthy.
fake, tc = patch_client(age=None, is_connected=True)
r = tc.get("/ping")
assert r.status_code == 200
body = r.json()
assert body["status"] == "ok"
assert body["last_probe_age_s"] is None
assert body["threshold_s"] == THRESHOLD
def test_ping_pre_start_connected_none_is_degraded_bool(patch_client):
# Before client.start(), Kurigram's is_connected is None. /ping must not 500: it coerces
# to a bool, so "connected" is false (never null) and the endpoint reports 503 degraded.
fake, tc = patch_client(age=None, is_connected=None)
r = tc.get("/ping")
assert r.status_code == 503
body = r.json()
assert body["status"] == "degraded"
assert body["connected"] is False # bool, not null
assert fake.client.get_me_calls == 0
def test_ping_issues_no_tg_rpc(patch_client):
"""The critical invariant: /ping never issues any Telegram RPC in any branch."""
for age, connected in [(1.0, True), (THRESHOLD + 500, True), (None, True), (1.0, False)]:
fake, tc = patch_client(age=age, is_connected=connected)
tc.get("/ping")
assert fake.client.get_me_calls == 0, f"get_me called (age={age}, connected={connected})"
assert fake.safe_get_messages_calls == 0, f"safe_get_messages called (age={age}, connected={connected})"
def test_ping_route_needs_no_token(patch_client):
# /ping is unauthenticated by design (no token path variant); it just works.
fake, tc = patch_client(age=1.0, is_connected=True)
assert tc.get("/ping").status_code == 200
# --------------------------------------------------------------------------- #
# TelegramClient.watchdog_last_ok_age accessor
# --------------------------------------------------------------------------- #
def test_watchdog_last_ok_age_none_when_never_probed():
tgc = TelegramClient()
assert tgc._wd_last_ok_monotonic is None
assert tgc.watchdog_last_ok_age() is None
def test_watchdog_last_ok_age_positive_after_probe():
tgc = TelegramClient()
tgc._wd_last_ok_monotonic = time.monotonic() - 5
age = tgc.watchdog_last_ok_age()
assert age is not None
assert age >= 5.0
# Sanity: a plausible upper bound so we didn't accidentally read the wrong field.
assert age < 60.0
tests/test_stage7_integration.py
+349
View File
@@ -0,0 +1,349 @@
# 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 7 — cross-stage END-TO-END integration tests.
Unlike the per-stage suites (which pin one seam in isolation), these drive the real
public entry points (`get_media`, `ping`, the access-time flush) so a regression that
only shows up when the stages interact goes red. Every scenario here maps to one of the
plan's "Стадия 7 — сквозные ручные сценарии"; the ones that genuinely need a running
server + real downloads + lsof (fd-leak counting) are NOT faked here — they stay in
docs/stability-verification.md for the operator's prod observation.
Automated here:
- Range semantics at the /media ROUTE level (get_media -> prepare_file_response ->
FileResponse driven through real ASGI): bytes=0-99 -> 206, bytes=-100 -> 206,
bytes=999999999- -> 416. (Stage 3 pins these on prepare_file_response directly; this
adds the end-to-end assertion that get_media's cache-hit branch reaches FileResponse
with a live Range still honored — i.e. stages 2+3 wired together.)
- /ping (stage 6) stays fast + correct while a deliberately-slow op is in flight, issuing
ZERO Telegram RPC — proving the healthcheck is decoupled from the hot/blocked paths.
- In-flight dedup + disconnect cleanup (stages 1/2) through the REAL get_media entry
point: concurrent requests share one download and the _inflight registry drains; a
cancelled request (client disconnect) leaves neither a stuck key nor a hung sibling.
- str(channel) access-time key consistency (stage 5) end-to-end: a /media cache hit for
an int-ish channel records the str-keyed timestamp, and the flush UPDATE matches the
str-keyed DB row (hit -> flush -> DB), the exact affinity gotcha the plan warns about.
"""
import os
import sys
import time
import sqlite3
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'])
from fastapi import FastAPI
from fastapi.testclient import TestClient
import api_server
from pyrogram import errors
from file_io import init_db_sync
# 2048 deterministic bytes so Range slices are byte-checkable.
BODY = bytes(range(256)) * 8
SIZE = len(BODY)
def _media_app():
"""A bare app that mounts the REAL get_media (and ping) with NO lifespan, so
client.start() never runs — a pure cache hit never touches Telegram, and FileResponse
computes Range/206/416 at send time, which only happens when driven through ASGI."""
app = FastAPI()
app.add_api_route("/media/{channel}/{post_id}/{file_unique_id}/{digest}", api_server.get_media, methods=["GET"])
app.add_api_route("/media/{channel}/{post_id}/{file_unique_id}", api_server.get_media, methods=["GET"])
return app
def _seed_cache(tmp_path, channel, post_id, fid, body=BODY):
cache_dir = tmp_path / "data" / "cache" / str(channel) / str(post_id)
cache_dir.mkdir(parents=True, exist_ok=True)
(cache_dir / fid).write_bytes(body)
return cache_dir / fid
# --------------------------------------------------------------------------- #
# Range semantics at the /media ROUTE level (stages 2 + 3 wired together).
# Plan scenario: `curl -H "Range: bytes=0-99" / "bytes=-100" / "bytes=999999999-"`.
# Regression caught: any change that makes get_media's cache-hit branch stop reaching
# FileResponse (e.g. re-buffering the body, hand-rolling headers, dropping the media_key
# path) or that breaks the digest gate wiring — the Range would stop being honored.
# --------------------------------------------------------------------------- #
def test_media_route_range_prefix_0_99(monkeypatch, tmp_path):
monkeypatch.chdir(tmp_path)
_seed_cache(tmp_path, "chan", 3, "fidR")
monkeypatch.setattr(api_server, "verify_media_digest", lambda url, digest: True)
# Keep the MIME path DB-free; the FileResponse/Range machinery is what we exercise.
monkeypatch.setattr(api_server, "get_mime_type_sync", lambda *a, **k: None)
monkeypatch.setattr(api_server, "set_mime_type_sync", lambda *a, **k: None)
c = TestClient(_media_app())
r = c.get("/media/chan/3/fidR/anydigest", headers={"Range": "bytes=0-99"})
assert r.status_code == 206
assert r.headers["content-range"] == f"bytes 0-99/{SIZE}"
assert r.headers["content-length"] == "100"
assert r.content == BODY[:100]
assert r.headers["accept-ranges"] == "bytes"
def test_media_route_range_suffix_last_100(monkeypatch, tmp_path):
monkeypatch.chdir(tmp_path)
_seed_cache(tmp_path, "chan", 3, "fidS")
monkeypatch.setattr(api_server, "verify_media_digest", lambda url, digest: True)
monkeypatch.setattr(api_server, "get_mime_type_sync", lambda *a, **k: None)
monkeypatch.setattr(api_server, "set_mime_type_sync", lambda *a, **k: None)
c = TestClient(_media_app())
r = c.get("/media/chan/3/fidS/anydigest", headers={"Range": "bytes=-100"})
assert r.status_code == 206
assert r.headers["content-range"] == f"bytes {SIZE - 100}-{SIZE - 1}/{SIZE}"
assert r.headers["content-length"] == "100"
assert r.content == BODY[-100:]
def test_media_route_range_unsatisfiable_416(monkeypatch, tmp_path):
monkeypatch.chdir(tmp_path)
_seed_cache(tmp_path, "chan", 3, "fidU")
monkeypatch.setattr(api_server, "verify_media_digest", lambda url, digest: True)
monkeypatch.setattr(api_server, "get_mime_type_sync", lambda *a, **k: None)
monkeypatch.setattr(api_server, "set_mime_type_sync", lambda *a, **k: None)
c = TestClient(_media_app())
r = c.get("/media/chan/3/fidU/anydigest", headers={"Range": "bytes=999999999-"})
assert r.status_code == 416
# Starlette's 416 Content-Range is `*/size` (documented stage-3 RFC-7233 delta).
assert r.headers["content-range"] == f"*/{SIZE}"
# --------------------------------------------------------------------------- #
# /ping (stage 6) stays fast + correct while a slow op is in flight, zero TG RPC.
# Plan scenario: "Генерация фида на 100+ сообщений + параллельный /ping -> ping < 100 мс".
# We model the concurrent slow op as a coroutine parked on an Event that is NEVER set
# during the ping, and assert ping resolves while it is still pending AND touches no RPC.
# Regression caught: re-coupling /ping to a TG RPC (get_me / safe_get_messages) or to any
# awaitable that a blocked hot path could stall — the ping would no longer return promptly
# or the spy counts would go non-zero.
# --------------------------------------------------------------------------- #
class _FakeKurigram:
def __init__(self, is_connected=True):
self.is_connected = is_connected
self.get_me_calls = 0
async def get_me(self):
self.get_me_calls += 1
raise AssertionError("/ping must never call get_me()")
class _FakeTelegramClient:
def __init__(self, age, is_connected=True):
self._age = age
self.client = _FakeKurigram(is_connected=is_connected)
self.safe_get_messages_calls = 0
def watchdog_last_ok_age(self):
return self._age
async def safe_get_messages(self, *a, **k):
self.safe_get_messages_calls += 1
raise AssertionError("/ping must never call safe_get_messages()")
async def test_ping_prompt_and_rpc_free_while_slow_op_pending(monkeypatch):
threshold = api_server.Config["tg_ping_unhealthy_after"]
fake = _FakeTelegramClient(age=threshold - 10, is_connected=True)
monkeypatch.setattr(api_server, "client", fake)
# A concurrent slow operation (a stand-in for a hung feed/RPC hot path) parked on an
# Event we deliberately never set for the duration of the ping.
gate = asyncio.Event()
async def slow_op():
await gate.wait()
slow = asyncio.create_task(slow_op())
await asyncio.sleep(0) # let slow_op start and park on the gate
# The real proof of decoupling: ping() returns under a tight deadline while the slow op
# is parked, AND issues zero TG RPC. wait_for reds if ping ever blocks; the spies (which
# raise if touched) red if ping recouples to any RPC. `not slow.done()` is only a sanity
# check that ping did not somehow drive the parked op — the gate keeps it pending anyway.
resp = await asyncio.wait_for(api_server.ping(), timeout=1.0)
assert resp.status_code == 200 # correct health while connected + fresh
assert not slow.done() # sanity: slow op still parked, ping did not await it
assert fake.client.get_me_calls == 0 # decoupled: zero TG RPC
assert fake.safe_get_messages_calls == 0
gate.set()
await slow
async def test_ping_reports_degraded_promptly_while_slow_op_pending(monkeypatch):
threshold = api_server.Config["tg_ping_unhealthy_after"]
fake = _FakeTelegramClient(age=threshold + 100, is_connected=True) # stale probe
monkeypatch.setattr(api_server, "client", fake)
gate = asyncio.Event()
async def slow_op():
await gate.wait()
slow = asyncio.create_task(slow_op())
await asyncio.sleep(0)
resp = await asyncio.wait_for(api_server.ping(), timeout=1.0)
assert resp.status_code == 503 # stale watchdog probe -> degraded, still instant
assert not slow.done() # sanity: slow op still parked, ping did not await it
assert fake.client.get_me_calls == 0
assert fake.safe_get_messages_calls == 0
gate.set()
await slow
# --------------------------------------------------------------------------- #
# In-flight dedup + disconnect cleanup (stages 1/2) through the REAL get_media entry.
# Plan scenario: "Параллельные запросы одного большого видео -> нет частичной отдачи" and
# "Отключение клиента на середине стрима -> нет утечки тасков/фд".
# The pure-unit stage-2 tests pin _download_deduped directly; these drive get_media so the
# HTTP semaphore + dedup registry + serve path are proven wired together.
# Regression caught: moving the download back into the request coroutine (so a disconnect
# cancels it), or dropping the finally-pop, would leave a stuck _inflight key / hung sibling.
# --------------------------------------------------------------------------- #
async def test_get_media_concurrent_shares_one_download_and_drains_registry(monkeypatch, tmp_path):
monkeypatch.chdir(tmp_path)
api_server._inflight.clear()
monkeypatch.setattr(api_server, "verify_media_digest", lambda url, digest: True)
# Serve step is not under test here; keep it to a sentinel so we assert on dedup + registry.
sentinel = object()
async def fake_prepare(file_path, request=None, delete_after=False, media_key=None):
return sentinel
monkeypatch.setattr(api_server, "prepare_file_response", fake_prepare)
calls = []
async def slow_dl(channel, post_id, fid):
calls.append(fid)
await asyncio.sleep(0.05) # real overlap window for the two requests
return (f"/final/{fid}", False)
monkeypatch.setattr(api_server, "download_media_file", slow_dl)
req = SimpleNamespace(headers={})
t1 = asyncio.create_task(api_server.get_media("chan", 9, "vfid", request=req, digest="x"))
t2 = asyncio.create_task(api_server.get_media("chan", 9, "vfid", request=req, digest="x"))
r1, r2 = await asyncio.gather(t1, t2)
assert r1 is sentinel and r2 is sentinel
assert calls == ["vfid"] # exactly ONE real download, shared by both requests
assert not api_server._inflight # registry drained — no forever-busy key
async def test_get_media_request_cancel_does_not_stick_registry_or_hang_sibling(monkeypatch, tmp_path):
monkeypatch.chdir(tmp_path)
api_server._inflight.clear()
monkeypatch.setattr(api_server, "verify_media_digest", lambda url, digest: True)
sentinel = object()
async def fake_prepare(file_path, request=None, delete_after=False, media_key=None):
return sentinel
monkeypatch.setattr(api_server, "prepare_file_response", fake_prepare)
gate = asyncio.Event()
calls = []
async def held_dl(channel, post_id, fid):
calls.append(fid)
await gate.wait() # hold the shared download open until we release it
return (f"/final/{fid}", False)
monkeypatch.setattr(api_server, "download_media_file", held_dl)
req = SimpleNamespace(headers={})
t1 = asyncio.create_task(api_server.get_media("chan", 9, "cfid", request=req, digest="x"))
await asyncio.sleep(0.02) # t1 registers the future + starts the detached download
t2 = asyncio.create_task(api_server.get_media("chan", 9, "cfid", request=req, digest="x"))
await asyncio.sleep(0.02)
# First client disconnects: its request coroutine is cancelled mid-wait.
t1.cancel()
with pytest.raises(asyncio.CancelledError):
await t1
# The detached download is unaffected; releasing it resolves the surviving request.
gate.set()
r2 = await asyncio.wait_for(t2, timeout=2.0)
assert r2 is sentinel
assert calls == ["cfid"] # download ran exactly once (not restarted)
assert not api_server._inflight # registry drained despite the disconnect
# --------------------------------------------------------------------------- #
# str(channel) access-time key consistency (stage 5) END-TO-END: hit -> flush -> DB.
# Plan gotcha #9: "Ключи SQLite: channel всегда str(...)"; if the key the hot path RECORDS
# and the key the flush UPDATEs by ever diverge, the timestamp silently stops updating and
# the file eventually falls out of the cache. Stage 5 pins hit and flush separately; this
# stitches them: the /media cache-hit records a (str-channel) key, and the flush's bulk
# UPDATE must land on that exact DB row. (get_media's route channel is already a str, so the
# str() there is a defensive no-op; the live int/str hazard is in download_media_file, pinned
# by stage-5's test_cache_hit_keys_channel_as_str — this test covers the get_media+flush leg.)
# Regression caught: any accumulator-key vs UPDATE-WHERE inconsistency (e.g. a transposed
# key-column order in the bulk SQL, verified to turn this test red) leaves `added` stale.
# --------------------------------------------------------------------------- #
async def test_media_cache_hit_flush_updates_str_keyed_db_row(monkeypatch, tmp_path):
monkeypatch.chdir(tmp_path)
api_server._access_updates = {}
channel, post_id, fid = "12345", 7, "fidINT" # int-ish channel, passed as the route str
_seed_cache(tmp_path, channel, post_id, fid)
db = str(tmp_path / "access.db")
init_db_sync(db)
# Seed the row keyed by the STRING channel with an old access time.
conn = sqlite3.connect(db)
conn.execute(
"INSERT INTO media_file_ids (channel, post_id, file_unique_id, added) VALUES (?,?,?,?)",
(channel, post_id, fid, 1.0),
)
conn.commit()
conn.close()
monkeypatch.setattr(api_server, "DB_PATH", db)
monkeypatch.setattr(api_server, "verify_media_digest", lambda url, digest: True)
sentinel = object()
async def fake_prepare(file_path, request=None, media_key=None):
return sentinel
monkeypatch.setattr(api_server, "prepare_file_response", fake_prepare)
before = time.time()
resp = await api_server.get_media(channel, post_id, fid, request=SimpleNamespace(headers={}), digest="x")
assert resp is sentinel
# Hot path recorded the str-keyed access time (never the raw int form).
assert (channel, post_id, fid) in api_server._access_updates
# Flush the accumulator; the bulk UPDATE must match the str-keyed row and refresh `added`.
await api_server._flush_access_updates()
assert api_server._access_updates == {}
conn = sqlite3.connect(db)
added = conn.execute(
"SELECT added FROM media_file_ids WHERE channel = ? AND post_id = ? AND file_unique_id = ?",
(channel, post_id, fid),
).fetchone()[0]
conn.close()
assert added >= before # refreshed from the stale 1.0 to ~now via the str-keyed WHERE