feat: real-time download progress, cancel support, admin auto-poll

This commit is contained in:
mini
2026-02-19 00:27:25 +08:00
parent 62a51305c3
commit 5741945531
6 changed files with 162 additions and 50 deletions

View File

@@ -5,6 +5,7 @@ import uuid
import json
import asyncio
import logging
import threading
import urllib.request
from pathlib import Path
from typing import Optional
@@ -12,6 +13,45 @@ import yt_dlp
logger = logging.getLogger(__name__)
# ── In-memory progress / cancel store (thread-safe via GIL) ─────────────────
_download_progress: dict[str, int] = {} # task_id → 0-100
_cancel_flags: dict[str, threading.Event] = {} # task_id → Event
def register_task(task_id: str):
_cancel_flags[task_id] = threading.Event()
_download_progress[task_id] = 0
def get_progress(task_id: str) -> int:
return _download_progress.get(task_id, 0)
def request_cancel(task_id: str):
flag = _cancel_flags.get(task_id)
if flag:
flag.set()
def cleanup_task(task_id: str):
_cancel_flags.pop(task_id, None)
_download_progress.pop(task_id, None)
def _make_hook(task_id: str):
"""yt-dlp progress hook with real-time tracking and cancel support."""
def hook(d):
flag = _cancel_flags.get(task_id)
if flag and flag.is_set():
raise yt_dlp.utils.DownloadCancelled("Cancelled by user")
if d["status"] == "downloading":
total = d.get("total_bytes") or d.get("total_bytes_estimate") or 0
done = d.get("downloaded_bytes", 0)
_download_progress[task_id] = int(done * 100 / total) if total else 0
elif d["status"] == "finished":
_download_progress[task_id] = 99 # merging, not 100 yet
return hook
VIDEO_BASE_PATH = os.getenv("VIDEO_BASE_PATH", "/home/xdl/xdl_videos")
X_VIDEOS_PATH = os.path.join(VIDEO_BASE_PATH, "x_videos")
YOUTUBE_VIDEOS_PATH = os.path.join(VIDEO_BASE_PATH, "youtube_videos")
@@ -149,7 +189,7 @@ def _parse_twitter_video(url: str) -> dict:
}
def _download_twitter_video(url: str, format_id: str = "best", progress_callback=None) -> dict:
def _download_twitter_video(url: str, format_id: str = "best", progress_callback=None, task_id: str = None) -> dict:
"""Download Twitter video using syndication API."""
tweet_id = _extract_tweet_id(url)
if not tweet_id:
@@ -198,14 +238,22 @@ def _download_twitter_video(url: str, format_id: str = "best", progress_callback
with open(filename, 'wb') as f:
while True:
# Check cancel flag
if task_id and _cancel_flags.get(task_id, threading.Event()).is_set():
raise yt_dlp.utils.DownloadCancelled("Cancelled by user")
chunk = resp.read(65536)
if not chunk:
break
f.write(chunk)
downloaded += len(chunk)
pct = int(downloaded * 100 / total) if total > 0 else 0
if task_id:
_download_progress[task_id] = pct
if progress_callback and total > 0:
progress_callback(int(downloaded * 100 / total))
progress_callback(pct)
if task_id:
_download_progress[task_id] = 99
if progress_callback:
progress_callback(100)
@@ -276,7 +324,7 @@ def _parse_youtube_video(url: str) -> dict:
}
def _download_youtube_video(url: str, format_id: str = "best", progress_callback=None) -> dict:
def _download_youtube_video(url: str, format_id: str = "best", progress_callback=None, task_id: str = None) -> dict:
"""Download YouTube video using yt-dlp."""
task_id = str(uuid.uuid4())[:8]
output_template = os.path.join(YOUTUBE_VIDEOS_PATH, f"%(id)s_{task_id}.%(ext)s")
@@ -286,14 +334,7 @@ def _download_youtube_video(url: str, format_id: str = "best", progress_callback
else:
format_spec = f"{format_id}+bestaudio/best"
def hook(d):
if d["status"] == "downloading" and progress_callback:
total = d.get("total_bytes") or d.get("total_bytes_estimate") or 0
downloaded = d.get("downloaded_bytes", 0)
pct = int(downloaded * 100 / total) if total > 0 else 0
progress_callback(pct)
elif d["status"] == "finished" and progress_callback:
progress_callback(100)
hooks = [_make_hook(task_id)] if task_id else []
ydl_opts = {
"format": format_spec,
@@ -301,7 +342,7 @@ def _download_youtube_video(url: str, format_id: str = "best", progress_callback
"merge_output_format": "mp4",
"quiet": True,
"no_warnings": True,
"progress_hooks": [hook],
"progress_hooks": hooks,
}
with yt_dlp.YoutubeDL(ydl_opts) as ydl:
@@ -384,7 +425,7 @@ def _parse_pornhub_video(url: str) -> dict:
}
def _download_pornhub_video(url: str, format_id: str = "best", progress_callback=None) -> dict:
def _download_pornhub_video(url: str, format_id: str = "best", progress_callback=None, task_id: str = None) -> dict:
"""Download Pornhub video using yt-dlp."""
task_id = str(uuid.uuid4())[:8]
output_template = os.path.join(PH_VIDEOS_PATH, f"%(id)s_{task_id}.%(ext)s")
@@ -396,14 +437,7 @@ def _download_pornhub_video(url: str, format_id: str = "best", progress_callback
# The format may already contain audio (merged); try with audio fallback gracefully
format_spec = f"{format_id}+bestaudio/{format_id}/best"
def hook(d):
if d["status"] == "downloading" and progress_callback:
total = d.get("total_bytes") or d.get("total_bytes_estimate") or 0
downloaded = d.get("downloaded_bytes", 0)
pct = int(downloaded * 100 / total) if total > 0 else 0
progress_callback(pct)
elif d["status"] == "finished" and progress_callback:
progress_callback(100)
hooks = [_make_hook(task_id)] if task_id else []
ydl_opts = {
"format": format_spec,
@@ -412,7 +446,7 @@ def _download_pornhub_video(url: str, format_id: str = "best", progress_callback
"quiet": True,
"no_warnings": True,
"http_headers": _PH_HEADERS,
"progress_hooks": [hook],
"progress_hooks": hooks,
}
with yt_dlp.YoutubeDL(ydl_opts) as ydl:
@@ -507,39 +541,32 @@ def parse_video_url(url: str) -> dict:
}
def download_video(url: str, format_id: str = "best", progress_callback=None) -> dict:
def download_video(url: str, format_id: str = "best", progress_callback=None, task_id: str = None) -> dict:
"""Download video and return file info."""
# Use syndication API for Twitter/X URLs
if _is_twitter_url(url):
logger.info(f"Using Twitter syndication API for download: {url}")
try:
return _download_twitter_video(url, format_id, progress_callback)
return _download_twitter_video(url, format_id, progress_callback, task_id=task_id)
except Exception as e:
logger.warning(f"Twitter syndication download failed, falling back to yt-dlp: {e}")
# YouTube URLs
if _is_youtube_url(url):
logger.info(f"Downloading YouTube video: {url}")
return _download_youtube_video(url, format_id, progress_callback)
return _download_youtube_video(url, format_id, progress_callback, task_id=task_id)
# Pornhub URLs
if _is_pornhub_url(url):
logger.info(f"Downloading Pornhub video: {url}")
return _download_pornhub_video(url, format_id, progress_callback)
return _download_pornhub_video(url, format_id, progress_callback, task_id=task_id)
task_id = str(uuid.uuid4())[:8]
output_template = os.path.join(X_VIDEOS_PATH, f"%(id)s_{task_id}.%(ext)s")
format_spec = "bestvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best" if format_id == "best" else f"{format_id}+bestaudio/best"
def hook(d):
if d["status"] == "downloading" and progress_callback:
total = d.get("total_bytes") or d.get("total_bytes_estimate") or 0
downloaded = d.get("downloaded_bytes", 0)
pct = int(downloaded * 100 / total) if total > 0 else 0
progress_callback(pct)
elif d["status"] == "finished" and progress_callback:
progress_callback(100)
hooks = [_make_hook(task_id)] if task_id else []
ydl_opts = {
"format": format_spec,
@@ -547,7 +574,7 @@ def download_video(url: str, format_id: str = "best", progress_callback=None) ->
"merge_output_format": "mp4",
"quiet": True,
"no_warnings": True,
"progress_hooks": [hook],
"progress_hooks": hooks,
}
with yt_dlp.YoutubeDL(ydl_opts) as ydl: