Files
redairsoft_vk_gt_max_sender/src/media_processor.py
T
exostring d01a17146e Retry oversized video downloads instead of giving up permanently
VK can briefly serve a not-yet-transcoded master stream on a capped-height
URL right after a video is published, tripping the size check on the very
first attempt. Treat that as transient (retry with a longer backoff) instead
of failing permanently, and merge the fallback note into a single "Смотреть
видео в VK" link instead of two separate lines.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-09-04 23:24:06 +05:00

322 lines
13 KiB
Python

from __future__ import annotations
from loguru import logger
import asyncio
import os
import sys
import tempfile
import time
import uuid
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Optional, Set
import aiohttp
try:
from .config import settings
except (ImportError, ValueError):
from config import settings
# Global registry of paths currently being written / downloaded / processed
_ACTIVE_LOCK = asyncio.Lock()
_ACTIVE_PATHS: Set[str] = set()
async def register_active_path(path: Path) -> None:
async with _ACTIVE_LOCK:
_ACTIVE_PATHS.add(str(path.resolve()))
async def unregister_active_path(path: Path) -> None:
async with _ACTIVE_LOCK:
_ACTIVE_PATHS.discard(str(path.resolve()))
async def is_path_active(path: Path) -> bool:
async with _ACTIVE_LOCK:
return str(path.resolve()) in _ACTIVE_PATHS
@dataclass
class ProcessedMedia:
media_type: str # "photo" or "video"
local_path: Optional[Path]
original_url: str
attachment_id: str
size_bytes: int = 0
duration_sec: Optional[int] = None
width: Optional[int] = None
height: Optional[int] = None
error: Optional[str] = None
is_link_only: bool = False
thumbnail_url: Optional[str] = None
class MediaProcessor:
def __init__(self, is_local_tg_api: bool = False) -> None:
self.is_local_tg_api = is_local_tg_api
self.cache_dir = settings.cache_path
@property
def max_video_bytes(self) -> int:
max_mb = (
settings.video_max_size_mb_local
if self.is_local_tg_api
else settings.video_max_size_mb_cloud
)
return max_mb * 1024 * 1024
async def download_photo(
self, session: aiohttp.ClientSession, url: str, prefix: str = "photo_"
) -> Optional[Path]:
if not url.startswith(("http://", "https://")):
return None
suffix = Path(url.split("?", 1)[0]).suffix[:6] or ".jpg"
tmp = tempfile.NamedTemporaryFile(
dir=self.cache_dir, prefix=prefix, suffix=suffix, delete=False
)
tmp_path = Path(tmp.name)
tmp.close()
await register_active_path(tmp_path)
try:
async with session.get(
url, timeout=aiohttp.ClientTimeout(total=settings.media_download_timeout_sec)
) as resp:
if resp.status != 200:
logger.warning("Photo download failed with status {}: {}", resp.status, url)
tmp_path.unlink(missing_ok=True)
await unregister_active_path(tmp_path)
return None
with open(tmp_path, "wb") as f:
async for chunk in resp.content.iter_chunked(64 * 1024):
f.write(chunk)
return tmp_path
except Exception as exc:
logger.warning("Error downloading photo {}: {}", url, exc)
tmp_path.unlink(missing_ok=True)
await unregister_active_path(tmp_path)
return None
async def download_video_ytdlp(
self, url: str, prefix: str = "video_"
) -> tuple[Optional[Path], Optional[str], bool]:
"""Retries transient failures (stalls, timeouts, one-off yt-dlp/network hiccups)
a couple of times before giving up - a single flaky attempt used to fall straight
through to the link-only fallback even though the video was perfectly downloadable."""
last_result: tuple[Optional[Path], Optional[str], bool] = (None, "video download failed", False)
for attempt in range(1, settings.yt_dlp_retry_attempts + 2):
path, err, is_permanent = await self._download_video_ytdlp_attempt(url, prefix)
if path or is_permanent:
return path, err, is_permanent
last_result = (path, err, is_permanent)
if attempt <= settings.yt_dlp_retry_attempts:
# A fresh upload's lower-quality renditions can still be mid-transcode on
# VK's side - the "720p" URL then briefly serves the full-res master until
# the ladder catches up, which trips our size cap. That's not a permanent
# rejection, just needs more time than a stalled connection does.
delay = 60.0 if "exceeds size limit" in (err or "") else 3.0
logger.warning(
"yt-dlp attempt {}/{} failed transiently for {}: {} - retrying in {}s",
attempt, settings.yt_dlp_retry_attempts + 1, url, err, delay,
)
await asyncio.sleep(delay)
return last_result
async def _download_video_ytdlp_attempt(
self, url: str, prefix: str = "video_"
) -> tuple[Optional[Path], Optional[str], bool]:
"""Returns (path, error_message, is_permanent_error)"""
output_path = self.cache_dir / f"{prefix}{uuid.uuid4().hex[:12]}.mp4"
await register_active_path(output_path)
max_size = self.max_video_bytes
netrc_path = None
cmd = [sys.executable, "-m", "yt_dlp"]
if "vk.com" in url and settings.vk_access_token:
netrc_path = f"{output_path}.netrc"
try:
with open(netrc_path, "w", encoding="utf-8") as fh:
fh.write(f"machine vk.com login vk_token password {settings.vk_access_token}\n")
cmd.extend(["--netrc-location", netrc_path])
except Exception as exc:
logger.warning("Failed to write netrc for yt-dlp: {}", exc)
cmd.extend([
url,
"-o", str(output_path),
"--no-playlist",
"--match-filter", f"duration <= {settings.video_max_duration_sec}",
"--merge-output-format", "mp4",
# Every fallback keeps the height cap: without it yt-dlp can pull an
# arbitrarily large/high-res stream only to have it discarded afterwards
# by the size check below, burning bandwidth and the whole timeout budget.
"-f", (
f"best[height<={settings.video_max_height}][filesize<{max_size}]"
f"/best[height<={settings.video_max_height}]"
f"/bestvideo[height<={settings.video_max_height}][filesize<{max_size}]+bestaudio"
f"/bestvideo[height<={settings.video_max_height}]+bestaudio"
),
# No --quiet: we need yt-dlp's own progress lines to tell a slow-but-alive
# download (fine, however long it takes) apart from a genuinely hung one -
# a flat wall-clock timeout can't tell those apart and was killing large
# videos that just needed more time (same class of bug as the TG upload
# timeout). --newline makes each progress update its own line instead of
# overwriting via \r, so we can read it with readline().
"--newline", "--no-warnings",
])
proc = None
output_lines: list[str] = []
deadline = asyncio.get_running_loop().time() + settings.yt_dlp_timeout_sec
stall_reason: Optional[str] = None
try:
proc = await asyncio.create_subprocess_exec(
*cmd,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.STDOUT,
)
while True:
remaining_total = deadline - asyncio.get_running_loop().time()
if remaining_total <= 0:
stall_reason = "video download exceeded overall timeout"
break
wait_for = min(settings.yt_dlp_stall_timeout_sec, remaining_total)
try:
line = await asyncio.wait_for(proc.stdout.readline(), timeout=wait_for)
except asyncio.TimeoutError:
stall_reason = "video download stalled (no progress from yt-dlp)"
break
if not line:
break # stdout closed - process is finishing up
output_lines.append(line.decode("utf-8", errors="ignore"))
if stall_reason:
proc.kill()
await proc.communicate()
output_path.unlink(missing_ok=True)
await unregister_active_path(output_path)
return None, stall_reason, False
await proc.wait()
finally:
if netrc_path:
Path(netrc_path).unlink(missing_ok=True)
if proc.returncode != 0 or not output_path.exists() or output_path.stat().st_size == 0:
err_msg = "".join(output_lines).lower()
is_permanent = any(
m in err_msg for m in (
"removed", "unavailable", "private", "access denied", "does not pass filter", "sign in"
)
)
output_path.unlink(missing_ok=True)
await unregister_active_path(output_path)
return None, f"yt-dlp failed: {err_msg[:200]}", is_permanent
size = output_path.stat().st_size
if size > max_size:
logger.warning(
"Downloaded video size {} MB exceeds limit {} MB",
round(size / (1024 * 1024), 2),
round(max_size / (1024 * 1024), 2),
)
output_path.unlink(missing_ok=True)
await unregister_active_path(output_path)
return None, "video exceeds size limit", False
return output_path, None, False
async def process_media_items(
self, items: list[Any]
) -> list[ProcessedMedia]:
results: list[ProcessedMedia] = []
async with aiohttp.ClientSession() as session:
for item in items:
media_type = item.media_type
url = item.url
att_id = item.attachment_id
if media_type == "photo":
p = await self.download_photo(session, url, prefix=f"p_{att_id}_")
if p:
results.append(
ProcessedMedia(
media_type="photo",
local_path=p,
original_url=url,
attachment_id=att_id,
size_bytes=p.stat().st_size,
width=item.width,
height=item.height,
)
)
else:
results.append(
ProcessedMedia(
media_type="photo",
local_path=None,
original_url=url,
attachment_id=att_id,
error="failed to download photo",
is_link_only=True,
)
)
elif media_type == "video":
# Check duration
if item.duration_sec and item.duration_sec > settings.video_max_duration_sec:
results.append(
ProcessedMedia(
media_type="video",
local_path=None,
original_url=url,
attachment_id=att_id,
duration_sec=item.duration_sec,
error="video duration exceeds maximum",
is_link_only=True,
thumbnail_url=item.thumbnail_url,
)
)
continue
p, err, perm = await self.download_video_ytdlp(url, prefix=f"v_{att_id}_")
if p:
results.append(
ProcessedMedia(
media_type="video",
local_path=p,
original_url=url,
attachment_id=att_id,
size_bytes=p.stat().st_size,
duration_sec=item.duration_sec,
width=item.width,
height=item.height,
)
)
else:
results.append(
ProcessedMedia(
media_type="video",
local_path=None,
original_url=url,
attachment_id=att_id,
duration_sec=item.duration_sec,
error=err or "failed to download video",
is_link_only=True,
thumbnail_url=item.thumbnail_url,
)
)
return results
async def cleanup(self, items: list[ProcessedMedia]) -> None:
"""Immediately unregisters and removes downloaded local files."""
for item in items:
if item.local_path:
try:
await unregister_active_path(item.local_path)
item.local_path.unlink(missing_ok=True)
except Exception as exc:
logger.warning("Error cleaning up {}: {}", item.local_path, exc)
item.local_path = None