1004 lines
44 KiB
Python
1004 lines
44 KiB
Python
"""Discovery and the strictly sequential processing pipeline.
|
|
|
|
Exactly one worker task exists per process, and InstanceLock guarantees one
|
|
process, so at most one video is downloading or uploading at any time.
|
|
|
|
Per item: activate -> reconcile with Drive -> resolve folder -> download ->
|
|
choose name -> upload (resumable) -> verify -> COMPLETED -> clean up.
|
|
|
|
Every step is idempotent and driven by the persisted record, so the same code
|
|
path handles first runs, retries and recovery after a crash or dyno restart.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import random
|
|
import threading
|
|
import time
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timedelta
|
|
|
|
from caption_parser import parse_caption
|
|
from database import iso, retry_time, utcnow
|
|
from drive_client import SOURCE_PROPERTY
|
|
from errors import (
|
|
BotError,
|
|
Interrupted,
|
|
PermanentError,
|
|
QueueHaltError,
|
|
SkipItem,
|
|
TransientError,
|
|
UploadSessionExpired,
|
|
)
|
|
from models import (
|
|
FINAL_TASK_STATUSES,
|
|
SETUP_TASK_STATUSES,
|
|
DriveFile,
|
|
LocalFile,
|
|
SourceMessage,
|
|
State,
|
|
TaskStatus,
|
|
UploadTask,
|
|
VideoRecord,
|
|
)
|
|
from naming import add_suffix, build_filename, clean_component, detect_video_extension, mime_for_name, normalize_key, truncate_name
|
|
from progress import HtmlReply, ProgressReporter, esc, format_duration
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
CURSOR_KEY = "last_discovered_message_id"
|
|
COMPLETED_CURSOR_KEY = "last_completed_message_id"
|
|
GLOBAL_RUN_KEY = "queue_global_run" # "1" after /run, "0" after /stop
|
|
PAUSED_KEY = "queue_paused" # "1" after /pause, "0" after /resume
|
|
CANCELLED_PREFIX = "cancelled:"
|
|
SESSION_MAX_AGE = timedelta(days=6) # Drive resumable sessions expire after about a week
|
|
MAX_DB_FAILURES = 5
|
|
UNFINISHED = ("PENDING", "RETRY_WAIT", "DOWNLOADING", "DOWNLOADED", "UPLOADING")
|
|
|
|
|
|
@dataclass
|
|
class RunStats:
|
|
uploaded: int = 0
|
|
reused: int = 0
|
|
skipped: int = 0
|
|
failed: int = 0
|
|
retries: int = 0
|
|
|
|
|
|
@dataclass
|
|
class SyncResult:
|
|
new: int = 0
|
|
known: int = 0
|
|
ignored: int = 0
|
|
skipped: int = 0
|
|
note: str | None = None
|
|
|
|
def add(self, outcome: str) -> None:
|
|
setattr(self, outcome, getattr(self, outcome) + 1)
|
|
|
|
def summary(self) -> str:
|
|
text = (
|
|
f"New videos queued: {self.new}\nAlready known: {self.known}\n"
|
|
f"Skipped by filters/duplicates: {self.skipped}\nNon-video messages ignored: {self.ignored}"
|
|
)
|
|
return f"{text}\nNote: {self.note}" if self.note else text
|
|
|
|
|
|
def public_error(err: BaseException) -> str:
|
|
text = getattr(err, "public", None) or type(err).__name__
|
|
return text if len(text) <= 300 else text[:297] + "..."
|
|
|
|
|
|
class QueueManager:
|
|
def __init__(
|
|
self,
|
|
config,
|
|
db,
|
|
drive,
|
|
downloader,
|
|
folders,
|
|
*,
|
|
source=None,
|
|
reporter: ProgressReporter | None = None,
|
|
db_error_delay: float = 30.0,
|
|
):
|
|
self.config = config
|
|
self.db = db
|
|
self.drive = drive
|
|
self.downloader = downloader
|
|
self.folders = folders
|
|
self.source = source
|
|
self.reporter = reporter
|
|
self.db_error_delay = db_error_delay
|
|
self.listening = False # True while live channel updates are being received
|
|
self.halt_reason: str | None = None
|
|
self.stats = RunStats()
|
|
self.current: VideoRecord | None = None
|
|
self.active_uploads = 0
|
|
self.max_concurrent_uploads = 0
|
|
self._task: asyncio.Task | None = None
|
|
self._paused = False
|
|
self._stop_requested = False
|
|
self._resume_event = asyncio.Event()
|
|
self._wake = asyncio.Event()
|
|
self._stop_event = asyncio.Event()
|
|
self._interrupt = threading.Event() # read from the upload thread
|
|
self._sync_lock = asyncio.Lock()
|
|
self._item_lock = asyncio.Lock() # held while one item is processed
|
|
self._db_failures = 0
|
|
self._last_cleanup = time.monotonic()
|
|
downloader.reclaim = lambda: downloader.cleanup_orphans(db, keep_failed=config.keep_failed_downloads)
|
|
|
|
# -- discovery ------------------------------------------------------------------
|
|
|
|
async def sync(self, *, rescan: bool = False) -> SyncResult:
|
|
"""Discover channel messages and queue unseen videos. Safe to repeat."""
|
|
if self.source is None:
|
|
raise RuntimeError("no Telegram source configured")
|
|
if self._sync_lock.locked():
|
|
return SyncResult(note="a sync is already running")
|
|
async with self._sync_lock:
|
|
cursor = self.db.get_int(CURSOR_KEY, 0)
|
|
if self.config.process_mode == "new" and cursor == 0 and not rescan:
|
|
return SyncResult(note="PROCESS_MODE=new and no cursor yet; only new posts will be queued")
|
|
start = self.config.start_message_id if rescan else max(self.config.start_message_id, cursor + 1)
|
|
result = SyncResult()
|
|
async for msg in self.source.iter_history(min_id=start, start_date=self.config.start_date):
|
|
result.add(self.register(msg))
|
|
self._wake.set()
|
|
log.info("Sync finished: %s", result.summary().replace("\n", "; "))
|
|
return result
|
|
|
|
def register(self, msg: SourceMessage, *, task_id: int | None = None) -> str:
|
|
"""Record one discovered message. Returns new / known / skipped / ignored."""
|
|
if msg.media is None:
|
|
self.db.raise_int(CURSOR_KEY, msg.message_id)
|
|
return "ignored"
|
|
if self.db.exists(msg.channel_id, msg.message_id):
|
|
return "known"
|
|
parsed = parse_caption(msg.caption)
|
|
reason = self._filter_reason(parsed)
|
|
if reason is None and self.config.duplicate_media_policy == "skip" and msg.media.media_id is not None:
|
|
for other in self.db.find_by_media_id(msg.media.media_id):
|
|
if other.message_id != msg.message_id and other.state != State.SKIPPED:
|
|
reason = f"duplicate of message {other.message_id} (same Telegram media)"
|
|
break
|
|
state = State.SKIPPED if reason else State.PENDING
|
|
sequence = None if reason else self.db.next_sequence()
|
|
inserted = self.db.insert_video(
|
|
msg, parsed, state=state, skip_reason=reason, sequence_no=sequence, task_id=task_id
|
|
)
|
|
self.db.raise_int(CURSOR_KEY, msg.message_id)
|
|
if not inserted:
|
|
return "known"
|
|
if parsed.missing or parsed.warnings:
|
|
log.warning("Message %s caption issues: %s", msg.message_id, "; ".join(parsed.warnings))
|
|
self._wake.set()
|
|
return "skipped" if reason else "new"
|
|
|
|
def _filter_reason(self, parsed) -> str | None:
|
|
if self.config.filter_topics and normalize_key(parsed.topic or "") not in self.config.filter_topics:
|
|
return "filtered: topic not in FILTER_TOPICS"
|
|
if self.config.filter_batches and normalize_key(parsed.batch or "") not in self.config.filter_batches:
|
|
return "filtered: batch not in FILTER_BATCHES"
|
|
return None
|
|
|
|
# -- control ----------------------------------------------------------------------
|
|
#
|
|
# Two persisted switches (kv table) describe what the operator asked for:
|
|
# queue_global_run = 1 -> /run: process every discovered video
|
|
# queue_paused = 1 -> /pause: nothing runs until /resume
|
|
# /upload tasks run whenever the queue is not paused, even without /run.
|
|
|
|
@property
|
|
def running(self) -> bool:
|
|
return self._task is not None and not self._task.done()
|
|
|
|
@property
|
|
def status(self) -> str:
|
|
if self.running:
|
|
if self._stop_requested:
|
|
return "stopping"
|
|
return "paused" if self._paused else "running"
|
|
return "halted" if self.halt_reason else "idle"
|
|
|
|
@property
|
|
def global_enabled(self) -> bool:
|
|
return self.db.get_kv(GLOBAL_RUN_KEY, "0") == "1"
|
|
|
|
@property
|
|
def paused_flag(self) -> bool:
|
|
return self.db.get_kv(PAUSED_KEY, "0") == "1"
|
|
|
|
def desired_state(self) -> str:
|
|
if self.paused_flag:
|
|
return "paused"
|
|
return "running" if self.global_enabled else "stopped"
|
|
|
|
def _persist_desired(self, value: str) -> None:
|
|
"""Back-compat helper: 'running' | 'paused' | 'stopped'."""
|
|
if value == "paused":
|
|
self.db.set_kv(PAUSED_KEY, "1")
|
|
else:
|
|
self.db.set_kv(PAUSED_KEY, "0")
|
|
self.db.set_kv(GLOBAL_RUN_KEY, "1" if value == "running" else "0")
|
|
|
|
def start(self) -> bool:
|
|
"""/run: process the whole queue. Returns False if the worker was already running."""
|
|
self._persist_desired("running")
|
|
return self.ensure_worker()
|
|
|
|
def ensure_worker(self) -> bool:
|
|
"""Start the worker if needed (used by /run, /upload tasks and startup)."""
|
|
if self.running:
|
|
if self._paused:
|
|
return self.resume()
|
|
self._wake.set()
|
|
return False
|
|
self.halt_reason = None
|
|
self._paused = False
|
|
self._stop_requested = False
|
|
self._interrupt.clear()
|
|
self._stop_event.clear()
|
|
self._resume_event.set()
|
|
self._task = asyncio.create_task(self._worker(), name="queue-worker")
|
|
return True
|
|
|
|
def has_runnable_work(self) -> bool:
|
|
return self.global_enabled or self.db.next_runnable_task() is not None
|
|
|
|
def boot(self) -> str:
|
|
"""Decide what to do after a (re)start, from durable state only."""
|
|
if self.paused_flag:
|
|
return "paused"
|
|
if self.has_runnable_work():
|
|
self.ensure_worker()
|
|
return "running"
|
|
return "idle"
|
|
|
|
def pause(self) -> bool:
|
|
if not self.running or self._paused:
|
|
return False
|
|
self.db.set_kv(PAUSED_KEY, "1")
|
|
self.db.set_task_status(TaskStatus.RUNNING, TaskStatus.PAUSED)
|
|
self._paused = True
|
|
self._resume_event.clear()
|
|
self._interrupt.set()
|
|
self._stop_event.set()
|
|
return True
|
|
|
|
def resume(self) -> bool:
|
|
self.db.set_kv(PAUSED_KEY, "0")
|
|
resumed_tasks = self.db.set_task_status(TaskStatus.PAUSED, TaskStatus.RUNNING)
|
|
if not self.running:
|
|
if self.has_runnable_work():
|
|
return self.ensure_worker()
|
|
return False
|
|
if not self._paused:
|
|
self._wake.set()
|
|
return resumed_tasks > 0
|
|
self._paused = False
|
|
self._interrupt.clear()
|
|
self._stop_event.clear()
|
|
self._resume_event.set()
|
|
return True
|
|
|
|
async def stop(self, *, persist: bool = True, timeout: float | None = None) -> None:
|
|
"""Stop at the next resumable checkpoint and end the worker.
|
|
|
|
With ``persist`` (the /stop command) the whole-queue run is switched off and a
|
|
running /upload task is parked as PAUSED. Without it (SIGTERM / dyno restart)
|
|
nothing durable changes, so processing resumes automatically after restart.
|
|
"""
|
|
if persist:
|
|
self._persist_desired("stopped")
|
|
self.db.set_task_status(TaskStatus.RUNNING, TaskStatus.PAUSED)
|
|
if not self.running:
|
|
return
|
|
self._stop_requested = True
|
|
self._interrupt.set()
|
|
self._stop_event.set()
|
|
self._resume_event.set()
|
|
self._wake.set()
|
|
try:
|
|
await asyncio.wait_for(asyncio.shield(self._task), timeout)
|
|
except asyncio.TimeoutError:
|
|
log.warning("Worker did not reach a checkpoint within %ss; its state remains resumable", timeout)
|
|
|
|
async def shutdown(self, timeout: float = 20.0) -> None:
|
|
"""Process termination (SIGTERM): checkpoint without changing the operator's choice."""
|
|
await self.stop(persist=False, timeout=timeout)
|
|
|
|
async def wait(self) -> None:
|
|
if self._task:
|
|
await self._task
|
|
|
|
def wake(self) -> None:
|
|
self._wake.set()
|
|
|
|
def retry_failed(self) -> int:
|
|
count = self.db.reset_failed(has_local=self.downloader.has_complete_local)
|
|
self._wake.set()
|
|
return count
|
|
|
|
def remaining(self) -> int:
|
|
counts = self.db.counts()
|
|
return sum(counts[s] for s in UNFINISHED)
|
|
|
|
# -- /upload tasks --------------------------------------------------------------------
|
|
|
|
def submit_task(self, task_id: int) -> UploadTask:
|
|
task = self.db.update_task(task_id, status=TaskStatus.QUEUED, confirmed_at=iso(utcnow()))
|
|
log.info("Upload task #%s queued (range %s, channel %s)", task.id, task.range_text, task.channel_id)
|
|
if not self.paused_flag:
|
|
self.ensure_worker()
|
|
self._wake.set()
|
|
return task
|
|
|
|
async def cancel_task(self, task_id: int, *, reason: str = "cancelled by administrator") -> str:
|
|
"""Cancel a setup or an active task. Completed uploads are always kept.
|
|
|
|
Returns one of: unknown, finished, setup_cancelled, cancelled.
|
|
"""
|
|
task = self.db.get_task(task_id)
|
|
if task is None:
|
|
return "unknown"
|
|
if task.status in SETUP_TASK_STATUSES:
|
|
self.db.update_task(task_id, status=TaskStatus.CANCELLED, completed_at=iso(utcnow()), last_error=reason)
|
|
return "setup_cancelled"
|
|
if task.status in FINAL_TASK_STATUSES:
|
|
return "finished"
|
|
if self.current is not None and self.current.task_id == task_id:
|
|
self._interrupt.set() # current transfer stops at its next checkpoint
|
|
async with self._item_lock:
|
|
await self._finalize_task(self.db.get_task(task_id), TaskStatus.CANCELLED, reason=reason)
|
|
if not self._paused and not self._stop_requested:
|
|
self._interrupt.clear()
|
|
self._wake.set()
|
|
return "cancelled"
|
|
|
|
async def _scan_task(self, task: UploadTask) -> None:
|
|
"""Examine every message ID of the task's inclusive range (idempotent)."""
|
|
self.db.update_task(task.id, status=TaskStatus.RUNNING, started_at=task.started_at or iso(utcnow()), last_error=None)
|
|
async for message_id, msg in self.source.iter_range(task.channel_id, task.start_message_id, task.end_message_id):
|
|
self._attach(task, message_id, msg)
|
|
task = self.db.update_task(task.id, scanned_at=iso(utcnow()))
|
|
counts = self._sync_task_counts(task.id)
|
|
log.info("Upload task #%s scanned: %s", task.id, counts)
|
|
if self.reporter:
|
|
self.reporter.new_message()
|
|
await self.reporter.notify(
|
|
HtmlReply(
|
|
f"🚀 <b>Upload Task #{task.id} started</b>\nRange: {task.range_text}\n"
|
|
f"Eligible videos: {counts['eligible_count']} "
|
|
f"(already uploaded: {counts['duplicate_count']}, non-video messages: {counts['non_video_count']})"
|
|
),
|
|
extra_chat_ids=(task.admin_user_id,),
|
|
)
|
|
|
|
def _attach(self, task: UploadTask, message_id: int, msg: SourceMessage | None) -> str:
|
|
key = (task.channel_id, message_id)
|
|
if msg is None:
|
|
result = "missing"
|
|
elif msg.media is None:
|
|
result = "non_video"
|
|
else:
|
|
existing = self.db.get(key)
|
|
if existing is None:
|
|
outcome = self.register(msg, task_id=task.id)
|
|
if outcome in ("new", "known"):
|
|
result = "pending"
|
|
else:
|
|
rec = self.db.get(key)
|
|
result = "duplicate" if (rec.skip_reason or "").startswith("duplicate") else "skipped"
|
|
else:
|
|
result = self._claim(existing, task)
|
|
self.db.add_task_item(task.id, task.channel_id, message_id, result)
|
|
return result
|
|
|
|
def _claim(self, rec: VideoRecord, task: UploadTask) -> str:
|
|
"""Attach an already-known video to ``task``; never re-uploads completed work."""
|
|
note = f"claimed by upload task #{task.id}"
|
|
if rec.state == State.COMPLETED:
|
|
return "duplicate"
|
|
if rec.state == State.SKIPPED:
|
|
reason = rec.skip_reason or ""
|
|
if reason.startswith(CANCELLED_PREFIX):
|
|
self.db.transition(rec.key, State.PENDING, note=note, task_id=task.id, skip_reason=None)
|
|
return "pending"
|
|
return "duplicate" if reason.startswith("duplicate") else "skipped"
|
|
if rec.state == State.FAILED:
|
|
target = State.DOWNLOADED if self.downloader.has_complete_local(rec) else State.PENDING
|
|
self.db.transition(rec.key, target, note=note, task_id=task.id, attempts=0, next_retry_at=None)
|
|
return "pending"
|
|
self.db.update_fields(rec.key, task_id=task.id)
|
|
return "pending"
|
|
|
|
def _sync_task_counts(self, task_id: int) -> dict[str, int]:
|
|
r = self.db.task_item_results(task_id)
|
|
total = sum(r.values())
|
|
counts = {
|
|
"examined_count": total - r.get("missing", 0),
|
|
"eligible_count": sum(r.get(k, 0) for k in ("pending", "uploaded", "duplicate", "skipped", "failed", "cancelled")),
|
|
"uploaded_count": r.get("uploaded", 0),
|
|
"duplicate_count": r.get("duplicate", 0),
|
|
"non_video_count": r.get("non_video", 0),
|
|
"skipped_count": r.get("skipped", 0) + r.get("cancelled", 0),
|
|
"missing_count": r.get("missing", 0),
|
|
"failed_count": r.get("failed", 0),
|
|
}
|
|
self.db.update_task(task_id, **counts)
|
|
if self.reporter and self.reporter.board.task and self.reporter.board.task.get("id") == task_id:
|
|
self.reporter.board.task.update(
|
|
completed=counts["uploaded_count"],
|
|
skipped=counts["duplicate_count"] + counts["non_video_count"] + counts["skipped_count"] + counts["missing_count"],
|
|
failed=counts["failed_count"],
|
|
)
|
|
return counts
|
|
|
|
async def _finalize_task(self, task: UploadTask, status: TaskStatus, *, reason: str | None = None) -> None:
|
|
if status == TaskStatus.CANCELLED:
|
|
for rec in self.db.videos_for_task(task.id):
|
|
if rec.state in (State.COMPLETED, State.SKIPPED, State.FAILED):
|
|
continue
|
|
if rec.state == State.DOWNLOADING:
|
|
rec = self.db.transition(rec.key, State.PENDING, note="checkpoint before cancel")
|
|
skip = f"{CANCELLED_PREFIX} task #{task.id}"
|
|
self.db.transition(rec.key, State.SKIPPED, note=skip, skip_reason=skip, upload_session_uri=None)
|
|
self.db.set_task_item_result(task.id, rec.channel_id, rec.message_id, "cancelled")
|
|
deleted = self.downloader.cleanup(rec)
|
|
if deleted:
|
|
self.db.add_task_counter(task.id, "deleted_files_count", deleted)
|
|
counts = self._sync_task_counts(task.id)
|
|
task = self.db.update_task(
|
|
task.id, status=status, completed_at=iso(utcnow()), current_message_id=None,
|
|
last_error=reason if status != TaskStatus.COMPLETED else task.last_error,
|
|
)
|
|
log.info("Upload task #%s %s: %s", task.id, status.value, counts)
|
|
if self.reporter:
|
|
if self.reporter.board.task and self.reporter.board.task.get("id") == task.id:
|
|
self.reporter.board.task = None
|
|
await self.reporter.notify(self.task_summary(task), extra_chat_ids=(task.admin_user_id,))
|
|
self.reporter.new_message()
|
|
|
|
def task_summary(self, task: UploadTask) -> HtmlReply:
|
|
title = {
|
|
TaskStatus.COMPLETED: "✅ <b>Upload Task Completed</b>",
|
|
TaskStatus.CANCELLED: "🛑 <b>Upload Task Cancelled</b>",
|
|
TaskStatus.FAILED: "❌ <b>Upload Task Failed</b>",
|
|
}.get(task.status, f"<b>Upload Task {esc(task.status.value)}</b>")
|
|
folders = self.db.task_folders(task.id)
|
|
lines = [
|
|
title,
|
|
"",
|
|
f"Task: #{task.id}",
|
|
f"Requested range: {task.range_text}",
|
|
f"Messages examined: {task.examined_count}" + (f" (+{task.missing_count} missing/inaccessible)" if task.missing_count else ""),
|
|
f"Eligible videos found: {task.eligible_count}",
|
|
f"Uploaded: {task.uploaded_count}",
|
|
f"Duplicates skipped: {task.duplicate_count}",
|
|
f"Non-video messages skipped: {task.non_video_count}",
|
|
f"Other skipped: {task.skipped_count}",
|
|
f"Failed: {task.failed_count}",
|
|
f"Local temporary files deleted: {task.deleted_files_count}",
|
|
"Drive folders used: " + (", ".join(esc(f) for f in folders) if folders else "-"),
|
|
]
|
|
if task.last_error and task.status != TaskStatus.COMPLETED:
|
|
lines.append(f"Reason: {esc(task.last_error)}")
|
|
return HtmlReply("\n".join(lines))
|
|
|
|
async def _run_task_step(self, task: UploadTask) -> bool:
|
|
"""Advance ``task`` by one step. Returns False when it is waiting for a retry."""
|
|
if task.status == TaskStatus.QUEUED or task.scanned_at is None:
|
|
retry_in = None
|
|
async with self._item_lock: # a concurrent /cancel waits for the scan to finish
|
|
task = self.db.get_task(task.id)
|
|
if task.status not in (TaskStatus.QUEUED, TaskStatus.RUNNING):
|
|
return True
|
|
try:
|
|
await self._scan_task(task)
|
|
except PermanentError as exc:
|
|
await self._finalize_task(task, TaskStatus.FAILED, reason=public_error(exc))
|
|
except TransientError as exc:
|
|
self.db.update_task(task.id, last_error=public_error(exc))
|
|
log.warning("Upload task #%s scan failed (%s); will retry", task.id, exc)
|
|
retry_in = exc.retry_after or max(self.config.retry_base_delay, 5.0)
|
|
if retry_in:
|
|
await self._sleep(retry_in)
|
|
return True
|
|
rec = self.db.next_item(task_id=task.id)
|
|
if rec is None:
|
|
if self.db.task_unfinished_count(task.id) == 0:
|
|
await self._finalize_task(task, TaskStatus.COMPLETED)
|
|
return True
|
|
return False
|
|
if self.reporter:
|
|
board = self.reporter.board
|
|
if not board.task or board.task.get("id") != task.id:
|
|
board.task = {"id": task.id, "start": task.start_message_id, "end": task.end_message_id}
|
|
self._sync_task_counts(task.id)
|
|
board.task["current"] = rec.message_id
|
|
self.db.update_task(task.id, current_message_id=rec.message_id)
|
|
await self._process_locked(rec)
|
|
self._sync_task_counts(task.id)
|
|
return True
|
|
|
|
async def _process_locked(self, rec: VideoRecord) -> None:
|
|
async with self._item_lock:
|
|
fresh = self.db.get(rec.key)
|
|
# A cancellation may have finalized this item while we waited for the lock.
|
|
if fresh is None or fresh.state in (State.COMPLETED, State.SKIPPED, State.FAILED):
|
|
return
|
|
await self.process_item(fresh)
|
|
if self.config.processing_delay and not self._stop_requested and not self._paused:
|
|
await self._sleep(self.config.processing_delay)
|
|
|
|
def _maybe_cleanup_orphans(self) -> None:
|
|
interval = self.config.temp_cleanup_interval
|
|
now = time.monotonic()
|
|
if interval and now - self._last_cleanup >= interval:
|
|
self._last_cleanup = now
|
|
deleted = self.downloader.cleanup_orphans(self.db, keep_failed=self.config.keep_failed_downloads)
|
|
if deleted:
|
|
log.info("Periodic cleanup removed %d abandoned temporary file(s)", deleted)
|
|
|
|
# -- worker -------------------------------------------------------------------------
|
|
|
|
async def _worker(self) -> None:
|
|
self.stats = RunStats()
|
|
if self.reporter:
|
|
await self.reporter.start()
|
|
log.info("Queue worker started")
|
|
try:
|
|
while not self._stop_requested:
|
|
if self._paused:
|
|
if self.reporter:
|
|
self.reporter.board.status = "Paused"
|
|
await self._resume_event.wait()
|
|
continue
|
|
self._wake.clear()
|
|
try:
|
|
self._maybe_cleanup_orphans()
|
|
task = self.db.next_runnable_task()
|
|
if task is not None:
|
|
if not await self._run_task_step(task):
|
|
await self._wait_for_work(force=True)
|
|
continue
|
|
if not self.global_enabled:
|
|
break
|
|
rec = self.db.next_item()
|
|
except self.db.errors as exc:
|
|
await self._on_db_error(exc)
|
|
continue
|
|
if rec is None:
|
|
if not await self._wait_for_work():
|
|
break
|
|
continue
|
|
await self._process_locked(rec)
|
|
finally:
|
|
log.info("Queue worker stopped (%s)", self.halt_reason or ("stop requested" if self._stop_requested else "queue empty"))
|
|
if self.reporter:
|
|
self.reporter.board.status = "Idle"
|
|
self.reporter.board.task = None
|
|
await self.reporter.stop()
|
|
await self.reporter.notify(self._summary())
|
|
|
|
async def _wait_for_work(self, *, force: bool = False) -> bool:
|
|
next_retry = self.db.earliest_retry_at()
|
|
if next_retry is None and not self.listening and not force:
|
|
return False
|
|
timeout = 60.0 if next_retry is None else max(1.0, (next_retry - utcnow()).total_seconds())
|
|
if self.reporter:
|
|
self.reporter.board.status = "Waiting"
|
|
try:
|
|
await asyncio.wait_for(self._wake.wait(), min(timeout, 300.0))
|
|
except asyncio.TimeoutError:
|
|
pass
|
|
return True
|
|
|
|
async def _sleep(self, seconds: float) -> None:
|
|
try:
|
|
await asyncio.wait_for(self._stop_event.wait(), seconds)
|
|
except asyncio.TimeoutError:
|
|
pass
|
|
|
|
def _summary(self) -> str:
|
|
s = self.stats
|
|
lines = [
|
|
"Run finished" if not self.halt_reason else f"Processing halted: {self.halt_reason}",
|
|
f"Uploaded: {s.uploaded}",
|
|
f"Already in Drive (duplicates reused): {s.reused}",
|
|
f"Skipped: {s.skipped}",
|
|
f"Failed: {s.failed}",
|
|
f"Remaining: {self.remaining()}",
|
|
]
|
|
return "\n".join(lines)
|
|
|
|
async def _notify(self, text: str) -> None:
|
|
if self.reporter:
|
|
await self.reporter.notify(text)
|
|
|
|
async def _halt(self, reason: str) -> None:
|
|
self.halt_reason = reason
|
|
self._stop_requested = True
|
|
log.error("Queue halted: %s", reason)
|
|
|
|
async def _on_db_error(self, exc: BaseException) -> None:
|
|
self._db_failures += 1
|
|
log.error("Database error (%d/%d): %s", self._db_failures, MAX_DB_FAILURES, exc)
|
|
if self._db_failures >= MAX_DB_FAILURES:
|
|
await self._halt("database unavailable; check DATABASE_URL / database health, then /resume")
|
|
else:
|
|
await self._sleep(self.db_error_delay)
|
|
|
|
# -- per item -------------------------------------------------------------------------
|
|
|
|
def folder_path(self, rec: VideoRecord) -> list[str]:
|
|
strict = self.config.strict_filenames
|
|
topic = clean_component(rec.topic, strict=strict) or self.config.fallback_topic
|
|
if self.config.folder_structure == "batch_topic":
|
|
return [clean_component(rec.batch, strict=strict) or self.config.fallback_batch, topic]
|
|
return [topic]
|
|
|
|
async def process_item(self, rec: VideoRecord) -> None:
|
|
key = rec.key
|
|
self.current = rec
|
|
if self.reporter:
|
|
counts = self.db.counts()
|
|
total = sum(counts.values()) - counts["SKIPPED"]
|
|
self.reporter.board.begin_item(
|
|
counts["COMPLETED"] + 1, total, " / ".join(self.folder_path(rec)), rec.title or f"(message {rec.message_id})"
|
|
)
|
|
try:
|
|
try:
|
|
await self._process(rec)
|
|
self._db_failures = 0
|
|
except Interrupted as exc:
|
|
log.info("Message %s: %s", rec.message_id, exc)
|
|
except SkipItem as exc:
|
|
self._skip(key, exc.reason)
|
|
except QueueHaltError as exc:
|
|
self._park(key, exc)
|
|
await self._notify(f"Processing halted on message {rec.message_id}: {public_error(exc)}")
|
|
await self._halt(public_error(exc))
|
|
except PermanentError as exc:
|
|
await self._fail(key, exc)
|
|
except TransientError as exc:
|
|
await self._retry_or_fail(key, exc)
|
|
except self.db.errors:
|
|
raise
|
|
except Exception as exc: # unexpected bug: record it and keep the queue moving
|
|
log.exception("Unexpected error processing message %s", rec.message_id)
|
|
wrapped = TransientError(f"{type(exc).__name__}: {exc}", public=f"Internal error ({type(exc).__name__}); see logs")
|
|
await self._retry_or_fail(key, wrapped)
|
|
except self.db.errors as exc:
|
|
await self._on_db_error(exc)
|
|
finally:
|
|
self.current = None
|
|
|
|
async def _process(self, rec: VideoRecord) -> None:
|
|
key = rec.key
|
|
rec = self._activate(rec)
|
|
board = self.reporter.board if self.reporter else None
|
|
|
|
# A previous attempt may have uploaded the file and crashed before saving.
|
|
if rec.upload_started_at or rec.drive_file_id:
|
|
if board:
|
|
board.status = "Checking Google Drive"
|
|
existing = await asyncio.to_thread(self._find_existing_upload, rec, None)
|
|
if existing:
|
|
await self._complete(rec, existing, None, reused=True)
|
|
return
|
|
|
|
if board:
|
|
board.status = "Resolving folder"
|
|
folder_id = await asyncio.to_thread(self.folders.resolve, self.folder_path(rec))
|
|
|
|
if rec.state == State.PENDING:
|
|
local = await self._download(rec)
|
|
else:
|
|
local = await self.downloader.local_file(rec)
|
|
rec = self.db.get(key)
|
|
|
|
drive_file, reused = await self._upload(rec, local, folder_id)
|
|
await self._complete(self.db.get(key), drive_file, local, reused=reused)
|
|
|
|
def _activate(self, rec: VideoRecord) -> VideoRecord:
|
|
has_local = self.downloader.has_complete_local(rec)
|
|
if rec.state == State.RETRY_WAIT:
|
|
return self.db.transition(rec.key, State.DOWNLOADED if has_local else State.PENDING, note="retry")
|
|
if rec.state == State.DOWNLOADING:
|
|
return self.db.transition(rec.key, State.PENDING, note="resume interrupted download")
|
|
if rec.state in (State.DOWNLOADED, State.UPLOADING) and not has_local:
|
|
# e.g. a Heroku dyno restart wiped the ephemeral disk; the upload
|
|
# session URI is kept and reused if the re-download is byte-identical.
|
|
return self.db.transition(rec.key, State.PENDING, note="local file missing; downloading again")
|
|
return rec
|
|
|
|
async def _download(self, rec: VideoRecord) -> LocalFile:
|
|
key = rec.key
|
|
rec = self.db.transition(key, State.DOWNLOADING, download_started_at=iso(utcnow()))
|
|
board = self.reporter.board if self.reporter else None
|
|
if board:
|
|
board.start_download(rec.expected_size)
|
|
|
|
def on_progress(done: int, total: int | None) -> None:
|
|
if board and board.download:
|
|
board.download.update(done, total)
|
|
|
|
try:
|
|
local = await self.downloader.download(rec, on_progress=on_progress, should_stop=self._interrupt.is_set)
|
|
except Interrupted:
|
|
self.db.transition(key, State.PENDING, note="checkpoint during download")
|
|
raise
|
|
fields = dict(
|
|
local_path=str(local.path), local_size=local.size, md5=local.md5, sha256=local.sha256, downloaded_at=iso(utcnow())
|
|
)
|
|
if rec.upload_session_uri and rec.md5 and rec.md5 != local.md5:
|
|
log.warning("Message %s: re-downloaded file differs; discarding old upload session", rec.message_id)
|
|
fields["upload_session_uri"] = None
|
|
self.db.transition(key, State.DOWNLOADED, **fields)
|
|
return local
|
|
|
|
def _find_existing_upload(self, rec: VideoRecord, local: LocalFile | None) -> DriveFile | None:
|
|
"""Find a verified Drive file previously uploaded for this exact message."""
|
|
candidates: list[DriveFile] = []
|
|
if rec.drive_file_id:
|
|
known = self.drive.get_file(rec.drive_file_id)
|
|
if known:
|
|
candidates.append(known)
|
|
candidates += self.drive.find_by_source_key(rec.source_key)
|
|
size = local.size if local else (rec.local_size or rec.expected_size)
|
|
md5 = local.md5 if local else rec.md5
|
|
for f in candidates:
|
|
if f.trashed or (size is not None and f.size != size):
|
|
continue
|
|
if md5 and f.md5 and f.md5 != md5:
|
|
continue
|
|
log.info("Message %s already uploaded as Drive file %s; reusing it", rec.message_id, f.id)
|
|
return f
|
|
if candidates:
|
|
log.warning("Message %s has Drive files that failed verification; uploading again", rec.message_id)
|
|
return None
|
|
|
|
def _choose_name(self, rec: VideoRecord, local: LocalFile, folder_id: str) -> tuple[str, DriveFile | None]:
|
|
cfg = self.config
|
|
name, warnings = build_filename(
|
|
rec.title,
|
|
message_id=rec.message_id,
|
|
original_filename=rec.original_filename,
|
|
mime_type=rec.mime_type,
|
|
detected_ext=detect_video_extension(local.path),
|
|
missing_title_policy=cfg.missing_title_policy,
|
|
mismatch_policy=cfg.extension_mismatch_policy,
|
|
strict=cfg.strict_filenames,
|
|
)
|
|
if cfg.sequence_prefix_enabled and rec.sequence_no:
|
|
name = truncate_name(cfg.sequence_prefix_format.format(seq=rec.sequence_no) + name)
|
|
new_warnings = [w for w in warnings if w not in rec.warnings]
|
|
if new_warnings:
|
|
self.db.update_fields(rec.key, warnings=rec.warnings + new_warnings)
|
|
|
|
candidate, n = name, 1
|
|
while True:
|
|
clashes = self.drive.find_by_name(folder_id, candidate)
|
|
if not clashes:
|
|
return candidate, None
|
|
for f in clashes:
|
|
if f.app_properties.get(SOURCE_PROPERTY) == rec.source_key and f.size == local.size:
|
|
return candidate, f
|
|
if cfg.reuse_identical_files and f.size == local.size and f.md5 and f.md5 == local.md5:
|
|
log.info("Identical file %r (same size and MD5) already in folder; reusing it", candidate)
|
|
return candidate, f
|
|
policy = cfg.duplicate_name_policy
|
|
if policy == "skip":
|
|
raise SkipItem(f"a different file named {candidate!r} already exists in the folder")
|
|
if policy == "fail":
|
|
raise PermanentError(
|
|
f"name collision: {candidate!r}", public=f"Name collision: {candidate!r} already exists (DUPLICATE_POLICY=fail)"
|
|
)
|
|
n += 1
|
|
if n > 1000:
|
|
raise PermanentError(f"could not find a free name for {name!r}")
|
|
if policy == "suffix_message_id":
|
|
suffix = f"[{rec.message_id}]" if n == 2 else f"[{rec.message_id}] ({n - 1})"
|
|
else:
|
|
suffix = f"({n})"
|
|
candidate = add_suffix(name, suffix)
|
|
|
|
async def _upload(self, rec: VideoRecord, local: LocalFile, folder_id: str) -> tuple[DriveFile, bool]:
|
|
key = rec.key
|
|
session = rec.upload_session_uri
|
|
if session and rec.upload_session_created_at:
|
|
if utcnow() - datetime.fromisoformat(rec.upload_session_created_at) > SESSION_MAX_AGE:
|
|
session = None
|
|
if session and rec.drive_folder_id and rec.drive_folder_id != folder_id:
|
|
session = None
|
|
|
|
if session:
|
|
name = rec.final_name
|
|
else:
|
|
existing = await asyncio.to_thread(self._find_existing_upload, rec, local)
|
|
if existing:
|
|
return existing, True
|
|
name, reuse = await asyncio.to_thread(self._choose_name, rec, local, folder_id)
|
|
if reuse:
|
|
return reuse, True
|
|
|
|
fields = dict(final_name=name, drive_folder_id=folder_id, upload_session_uri=session)
|
|
if rec.state != State.UPLOADING:
|
|
# upload_started_at is written *before* the upload begins, so a crash at
|
|
# any later point triggers the Drive reconciliation on restart.
|
|
self.db.transition(key, State.UPLOADING, upload_started_at=rec.upload_started_at or iso(utcnow()), **fields)
|
|
else:
|
|
self.db.update_fields(key, **fields)
|
|
|
|
board = self.reporter.board if self.reporter else None
|
|
if board:
|
|
board.title = name
|
|
board.start_upload(local.size)
|
|
|
|
def on_session(uri: str) -> None:
|
|
self.db.update_fields(key, upload_session_uri=uri, upload_session_created_at=iso(utcnow()))
|
|
|
|
def on_progress(done: int, total: int) -> None:
|
|
if board and board.upload:
|
|
board.upload.update(done, total)
|
|
|
|
self.active_uploads += 1
|
|
self.max_concurrent_uploads = max(self.max_concurrent_uploads, self.active_uploads)
|
|
try:
|
|
uploaded = await asyncio.to_thread(
|
|
self.drive.upload_file,
|
|
local.path,
|
|
name=name,
|
|
parent_id=folder_id,
|
|
mime_type=mime_for_name(name, rec.mime_type),
|
|
size=local.size,
|
|
app_properties={SOURCE_PROPERTY: rec.source_key},
|
|
description=self._description(rec),
|
|
session_uri=session,
|
|
on_session=on_session,
|
|
on_progress=on_progress,
|
|
should_stop=self._interrupt.is_set,
|
|
)
|
|
except UploadSessionExpired:
|
|
self.db.update_fields(key, upload_session_uri=None)
|
|
raise
|
|
finally:
|
|
self.active_uploads -= 1
|
|
|
|
if board:
|
|
board.status = "Verifying"
|
|
try:
|
|
verified = await asyncio.to_thread(self._verify, uploaded, local)
|
|
except TransientError:
|
|
self.db.update_fields(key, upload_session_uri=None)
|
|
raise
|
|
return verified, False
|
|
|
|
def _verify(self, uploaded: DriveFile, local: LocalFile) -> DriveFile:
|
|
found = self.drive.get_file(uploaded.id)
|
|
if found is None or found.trashed:
|
|
raise TransientError(f"uploaded file {uploaded.id} not found during verification", public="Upload verification failed")
|
|
if found.size != local.size:
|
|
raise TransientError(
|
|
f"size mismatch after upload: Drive {found.size} vs local {local.size}", public="Upload verification failed (size)"
|
|
)
|
|
if found.md5 and found.md5 != local.md5:
|
|
raise TransientError("MD5 mismatch after upload", public="Upload verification failed (checksum)")
|
|
return found
|
|
|
|
def _description(self, rec: VideoRecord) -> str:
|
|
parts = [f"Telegram message {rec.message_id}"]
|
|
if rec.batch:
|
|
parts.append(f"Batch: {rec.batch}")
|
|
if rec.indexes:
|
|
parts.append(f"Index: {', '.join(rec.indexes)}")
|
|
return " | ".join(parts)
|
|
|
|
async def _complete(self, rec: VideoRecord, drive_file: DriveFile, local: LocalFile | None, *, reused: bool) -> None:
|
|
if self.config.share_mode == "anyone_with_link" and not reused:
|
|
await asyncio.to_thread(self.drive.share_with_link, drive_file.id)
|
|
fields = dict(
|
|
drive_file_id=drive_file.id,
|
|
final_name=drive_file.name,
|
|
completed_at=iso(utcnow()),
|
|
upload_session_uri=None,
|
|
last_error=None,
|
|
error_kind=None,
|
|
next_retry_at=None,
|
|
)
|
|
if drive_file.parents:
|
|
fields["drive_folder_id"] = drive_file.parents[0]
|
|
if local:
|
|
fields.update(local_size=local.size, md5=local.md5, sha256=local.sha256)
|
|
self.db.transition(rec.key, State.COMPLETED, note="reused existing Drive file" if reused else "uploaded and verified", **fields)
|
|
self.db.raise_int(COMPLETED_CURSOR_KEY, rec.message_id)
|
|
if rec.task_id is not None:
|
|
self.db.set_task_item_result(
|
|
rec.task_id, rec.channel_id, rec.message_id, "duplicate" if reused else "uploaded",
|
|
folder=" / ".join(self.folder_path(rec)),
|
|
)
|
|
# Local files are removed only after the COMPLETED state is committed.
|
|
if not self.config.keep_local_files:
|
|
self._cleanup_files(rec)
|
|
if reused:
|
|
self.stats.reused += 1
|
|
else:
|
|
self.stats.uploaded += 1
|
|
log.info("Message %s completed: %r -> Drive %s", rec.message_id, drive_file.name, drive_file.id)
|
|
|
|
def _cleanup_files(self, rec: VideoRecord) -> int:
|
|
deleted = self.downloader.cleanup(rec)
|
|
if deleted and rec.task_id is not None:
|
|
self.db.add_task_counter(rec.task_id, "deleted_files_count", deleted)
|
|
return deleted
|
|
|
|
def _skip(self, key, reason: str) -> None:
|
|
rec = self.db.transition(key, State.SKIPPED, skip_reason=reason, note=reason)
|
|
if rec.task_id is not None:
|
|
result = "duplicate" if reason.startswith("duplicate") else "skipped"
|
|
self.db.set_task_item_result(rec.task_id, rec.channel_id, rec.message_id, result)
|
|
if not self.config.keep_local_files:
|
|
self._cleanup_files(rec)
|
|
self.stats.skipped += 1
|
|
log.warning("Message %s skipped: %s", key[1], reason)
|
|
|
|
def _park(self, key, err: BotError) -> None:
|
|
"""Keep the item for a later retry without consuming an attempt."""
|
|
rec = self.db.get(key)
|
|
fields = dict(last_error=str(err)[:2000], error_kind=type(err).__name__, next_retry_at=iso(utcnow()))
|
|
if rec.state == State.RETRY_WAIT:
|
|
self.db.update_fields(key, **fields)
|
|
else:
|
|
self.db.transition(key, State.RETRY_WAIT, note="queue halted", **fields)
|
|
|
|
async def _retry_or_fail(self, key, err: TransientError) -> None:
|
|
rec = self.db.get(key)
|
|
attempts = rec.attempts + (0 if err.rate_limited else 1)
|
|
if attempts >= self.config.max_attempts:
|
|
await self._fail(key, err, attempts=attempts)
|
|
return
|
|
if err.retry_after is not None:
|
|
delay = float(err.retry_after)
|
|
else:
|
|
base = self.config.retry_base_delay * (2 ** max(attempts - 1, 0))
|
|
delay = min(base, self.config.retry_max_delay) * random.uniform(0.9, 1.1)
|
|
fields = dict(
|
|
attempts=attempts, last_error=str(err)[:2000], error_kind=type(err).__name__, next_retry_at=retry_time(delay)
|
|
)
|
|
if rec.state == State.RETRY_WAIT:
|
|
self.db.update_fields(key, **fields)
|
|
else:
|
|
self.db.transition(key, State.RETRY_WAIT, note="transient error", **fields)
|
|
self.stats.retries += 1
|
|
log.warning("Message %s: %s; retry in %.0fs (attempt %d/%d)", key[1], err, delay, attempts, self.config.max_attempts)
|
|
await self._notify(
|
|
f"Message {key[1]}: {public_error(err)}\nWill retry in {format_duration(delay)} "
|
|
f"(attempt {attempts} of {self.config.max_attempts})."
|
|
)
|
|
|
|
async def _fail(self, key, err: BotError, *, attempts: int | None = None) -> None:
|
|
rec = self.db.get(key)
|
|
fields = dict(last_error=str(err)[:2000], error_kind=type(err).__name__, next_retry_at=None)
|
|
if attempts is not None:
|
|
fields["attempts"] = attempts
|
|
if rec.state == State.FAILED:
|
|
rec = self.db.get(key)
|
|
self.db.update_fields(key, **fields)
|
|
else:
|
|
rec = self.db.transition(key, State.FAILED, note="failed", **fields)
|
|
if rec.task_id is not None:
|
|
self.db.set_task_item_result(rec.task_id, rec.channel_id, rec.message_id, "failed")
|
|
# The failure is recorded first; the partial/complete local file can always be
|
|
# downloaded again, so it is removed unless KEEP_FAILED_DOWNLOADS=true.
|
|
if not self.config.keep_failed_downloads and not self.config.keep_local_files:
|
|
self._cleanup_files(rec)
|
|
self.stats.failed += 1
|
|
log.error("Message %s failed: %s", key[1], err)
|
|
await self._notify(f"Message {key[1]} FAILED: {public_error(err)}\nNot retried automatically; use /retry.")
|
|
if self.config.strict_stop:
|
|
await self._halt(f"STRICT_STOP is enabled and message {key[1]} failed")
|