exostring
48f50d846c
Reverts the guessed long-retry-budget approach (no documented VK transcoding SLA exists to tune it against, and it would've blocked other routes in the same poll cycle for minutes). Instead, when a download comes out over the platform size limit - whether genuinely oversized or VK having briefly served an unfinished-transcode master on a height-capped URL - re-encode it locally with ffmpeg to a bitrate computed from its own real duration (ffprobe), so it fits deterministically rather than by retrying and hoping. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
395 lines
17 KiB
Python
395 lines
17 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.
|
|
An over-the-limit download is handled separately, by re-encoding down to size (see
|
|
_reencode_to_fit) rather than by retrying - VK publishes no signal for when a
|
|
fresher upload's lower-quality renditions finish transcoding, so waiting on a guess
|
|
isn't reliable, and re-encoding the file we already have is."""
|
|
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:
|
|
logger.warning(
|
|
"yt-dlp attempt {}/{} failed transiently for {}: {} - retrying",
|
|
attempt, settings.yt_dlp_retry_attempts + 1, url, err,
|
|
)
|
|
await asyncio.sleep(3.0)
|
|
return last_result
|
|
|
|
async def _ffprobe_duration_sec(self, path: Path) -> Optional[float]:
|
|
try:
|
|
proc = await asyncio.create_subprocess_exec(
|
|
"ffprobe", "-v", "error", "-show_entries", "format=duration",
|
|
"-of", "default=noprint_wrappers=1:nokey=1", str(path),
|
|
stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.DEVNULL,
|
|
)
|
|
out, _ = await proc.communicate()
|
|
return float(out.decode().strip())
|
|
except (OSError, ValueError):
|
|
return None
|
|
|
|
async def _reencode_to_fit(self, input_path: Path, max_size: int) -> Optional[Path]:
|
|
"""Re-encodes an over-the-limit download down under max_size instead of discarding
|
|
a perfectly watchable video - handles both a genuinely oversized source and VK
|
|
having briefly served an unfinished-transcode master on a height-capped URL. Target
|
|
bitrate is computed from the file's own real duration (ffprobe), not from VK's
|
|
metadata, so it's exact regardless of why the file came out this size."""
|
|
duration_sec = await self._ffprobe_duration_sec(input_path)
|
|
if not duration_sec or duration_sec <= 0:
|
|
return None
|
|
audio_bitrate = 128_000
|
|
# 10% headroom: encoders overshoot their target average bitrate somewhat, and we'd
|
|
# rather re-encode a touch smaller than land just over the limit again.
|
|
target_total_bitrate = int((max_size * 8) / duration_sec * 0.90)
|
|
target_video_bitrate = target_total_bitrate - audio_bitrate
|
|
if target_video_bitrate < 300_000:
|
|
# Duration this long vs. size this small leaves no sane quality to encode at -
|
|
# a real oversized-content case, not a transcode-timing artifact.
|
|
return None
|
|
|
|
output_path = input_path.with_name(f"{input_path.stem}_fit.mp4")
|
|
await register_active_path(output_path)
|
|
cmd = [
|
|
"ffmpeg", "-y", "-i", str(input_path),
|
|
"-vf", f"scale=-2:min({settings.video_max_height}\\,ih)",
|
|
"-c:v", "libx264", "-preset", "veryfast",
|
|
"-b:v", str(target_video_bitrate), "-maxrate", str(int(target_video_bitrate * 1.2)),
|
|
"-bufsize", str(target_video_bitrate * 2),
|
|
"-c:a", "aac", "-b:a", str(audio_bitrate),
|
|
str(output_path),
|
|
]
|
|
proc = None
|
|
try:
|
|
proc = await asyncio.create_subprocess_exec(
|
|
*cmd, stdout=asyncio.subprocess.DEVNULL, stderr=asyncio.subprocess.DEVNULL,
|
|
)
|
|
await asyncio.wait_for(proc.wait(), timeout=settings.yt_dlp_timeout_sec)
|
|
except (OSError, asyncio.TimeoutError):
|
|
if proc and proc.returncode is None:
|
|
proc.kill()
|
|
await proc.communicate()
|
|
output_path.unlink(missing_ok=True)
|
|
await unregister_active_path(output_path)
|
|
return None
|
|
|
|
if proc.returncode != 0 or not output_path.exists() or output_path.stat().st_size == 0:
|
|
output_path.unlink(missing_ok=True)
|
|
await unregister_active_path(output_path)
|
|
return None
|
|
if output_path.stat().st_size > max_size:
|
|
output_path.unlink(missing_ok=True)
|
|
await unregister_active_path(output_path)
|
|
return None
|
|
return output_path
|
|
|
|
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 - re-encoding to fit",
|
|
round(size / (1024 * 1024), 2),
|
|
round(max_size / (1024 * 1024), 2),
|
|
)
|
|
fitted = await self._reencode_to_fit(output_path, max_size)
|
|
output_path.unlink(missing_ok=True)
|
|
await unregister_active_path(output_path)
|
|
if fitted:
|
|
logger.info(
|
|
"Re-encoded to {} MB, under the {} MB limit",
|
|
round(fitted.stat().st_size / (1024 * 1024), 2),
|
|
round(max_size / (1024 * 1024), 2),
|
|
)
|
|
return fitted, None, False
|
|
return None, "video exceeds size limit even after re-encoding", True
|
|
|
|
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
|