first commit
This commit is contained in:
gosangam
+246
@@ -0,0 +1,246 @@
|
||||
"""Stream Telegram media to disk with resume, disk-space checks and hashing."""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import errno
|
||||
import hashlib
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from typing import AsyncIterator, Callable, Protocol
|
||||
|
||||
from errors import DiskSpaceError, Interrupted, TransientError
|
||||
from models import LocalFile, State, VideoRecord
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
# Partial downloads are resumed from a 1 MiB boundary: Telegram's upload.getFile
|
||||
# requires offsets aligned to the requested chunk size (512 KiB).
|
||||
RESUME_ALIGNMENT = 1024 * 1024
|
||||
HASH_BLOCK = 4 * 1024 * 1024
|
||||
# Only files with exactly this shape are ever deleted by cleanup routines.
|
||||
TEMP_NAME = re.compile(r"^(\d+)_(\d+)\.(part|video)$")
|
||||
|
||||
|
||||
class MediaSource(Protocol):
|
||||
def iter_media_chunks(
|
||||
self, message_id: int, *, offset: int, expected_size: int | None, channel_id: int | None = None
|
||||
) -> AsyncIterator[bytes]: ...
|
||||
|
||||
|
||||
class _Hashes:
|
||||
def __init__(self, sha256: bool):
|
||||
self.md5 = hashlib.md5()
|
||||
self.sha256 = hashlib.sha256() if sha256 else None
|
||||
|
||||
def update(self, data: bytes) -> None:
|
||||
self.md5.update(data)
|
||||
if self.sha256:
|
||||
self.sha256.update(data)
|
||||
|
||||
def digest(self) -> tuple[str, str | None]:
|
||||
return self.md5.hexdigest(), self.sha256.hexdigest() if self.sha256 else None
|
||||
|
||||
|
||||
def hash_file(path: Path, *, sha256: bool, limit: int | None = None, hashes: _Hashes | None = None) -> _Hashes:
|
||||
hashes = hashes or _Hashes(sha256)
|
||||
remaining = limit
|
||||
with open(path, "rb") as fh:
|
||||
while remaining is None or remaining > 0:
|
||||
block = fh.read(HASH_BLOCK if remaining is None else min(HASH_BLOCK, remaining))
|
||||
if not block:
|
||||
break
|
||||
hashes.update(block)
|
||||
if remaining is not None:
|
||||
remaining -= len(block)
|
||||
return hashes
|
||||
|
||||
|
||||
class Downloader:
|
||||
def __init__(
|
||||
self,
|
||||
source: MediaSource,
|
||||
download_dir: Path,
|
||||
*,
|
||||
min_free_bytes: int = 0,
|
||||
max_temp_bytes: int = 0,
|
||||
compute_sha256: bool = True,
|
||||
disk_usage: Callable[[str], tuple] = shutil.disk_usage,
|
||||
):
|
||||
self.source = source
|
||||
self.download_dir = download_dir
|
||||
self.min_free_bytes = min_free_bytes
|
||||
self.max_temp_bytes = max_temp_bytes
|
||||
self.compute_sha256 = compute_sha256
|
||||
self._disk_usage = disk_usage
|
||||
# Optional callback that frees space by removing abandoned temp files;
|
||||
# returns the number of files deleted. Set by the queue manager.
|
||||
self.reclaim: Callable[[], int] | None = None
|
||||
download_dir.mkdir(parents=True, exist_ok=True)
|
||||
self._root = download_dir.resolve()
|
||||
|
||||
# Local files are named by message identity, never by caption text, so
|
||||
# captions cannot cause path problems.
|
||||
def complete_path(self, rec: VideoRecord) -> Path:
|
||||
return self.download_dir / f"{abs(rec.channel_id)}_{rec.message_id}.video"
|
||||
|
||||
def partial_path(self, rec: VideoRecord) -> Path:
|
||||
return self.download_dir / f"{abs(rec.channel_id)}_{rec.message_id}.part"
|
||||
|
||||
def has_complete_local(self, rec: VideoRecord) -> bool:
|
||||
path = self.complete_path(rec)
|
||||
if not path.is_file():
|
||||
return False
|
||||
expected = rec.expected_size or rec.local_size
|
||||
return expected is not None and path.stat().st_size == expected
|
||||
|
||||
async def local_file(self, rec: VideoRecord) -> LocalFile:
|
||||
"""Describe an already complete download, reusing stored hashes if valid."""
|
||||
path = self.complete_path(rec)
|
||||
size = path.stat().st_size
|
||||
if rec.md5 and rec.local_size == size:
|
||||
return LocalFile(path, size, rec.md5, rec.sha256)
|
||||
hashes = await asyncio.to_thread(hash_file, path, sha256=self.compute_sha256)
|
||||
md5, sha = hashes.digest()
|
||||
return LocalFile(path, size, md5, sha)
|
||||
|
||||
def check_space(self, needed: int, *, exclude: Path | None = None) -> None:
|
||||
free = self._disk_usage(str(self.download_dir))[2]
|
||||
if free - needed < self.min_free_bytes:
|
||||
raise DiskSpaceError(
|
||||
f"not enough disk space: need {needed} bytes plus {self.min_free_bytes} reserve, {free} free",
|
||||
public="Not enough free disk space for the next video; free space and /resume.",
|
||||
)
|
||||
if self.max_temp_bytes:
|
||||
used = sum(
|
||||
p.stat().st_size for p in self.download_dir.iterdir() if p.is_file() and p != exclude
|
||||
)
|
||||
if used + needed > self.max_temp_bytes:
|
||||
raise DiskSpaceError(
|
||||
f"MAX_TEMP_STORAGE exceeded: {used} used + {needed} needed > {self.max_temp_bytes}",
|
||||
public="Temporary storage limit (MAX_TEMP_STORAGE) reached; /retry or clean up downloads.",
|
||||
)
|
||||
|
||||
async def download(
|
||||
self,
|
||||
rec: VideoRecord,
|
||||
*,
|
||||
on_progress: Callable[[int, int | None], None],
|
||||
should_stop: Callable[[], bool],
|
||||
) -> LocalFile:
|
||||
if self.has_complete_local(rec):
|
||||
return await self.local_file(rec)
|
||||
part = self.partial_path(rec)
|
||||
final = self.complete_path(rec)
|
||||
expected = rec.expected_size
|
||||
offset = part.stat().st_size if part.exists() else 0
|
||||
offset -= offset % RESUME_ALIGNMENT
|
||||
if expected is not None and offset > expected:
|
||||
offset = 0
|
||||
needed = (expected - offset) if expected else 0
|
||||
try:
|
||||
self.check_space(needed, exclude=part)
|
||||
except DiskSpaceError:
|
||||
if not (self.reclaim and self.reclaim()):
|
||||
raise
|
||||
self.check_space(needed, exclude=part)
|
||||
|
||||
hashes = _Hashes(self.compute_sha256)
|
||||
if offset:
|
||||
log.info("Resuming download of message %s at byte %d", rec.message_id, offset)
|
||||
await asyncio.to_thread(hash_file, part, sha256=self.compute_sha256, limit=offset, hashes=hashes)
|
||||
written = offset
|
||||
chunks = self.source.iter_media_chunks(
|
||||
rec.message_id, offset=offset, expected_size=expected, channel_id=rec.channel_id
|
||||
)
|
||||
try:
|
||||
with open(part, "r+b" if part.exists() else "wb") as fh:
|
||||
fh.truncate(offset)
|
||||
fh.seek(offset)
|
||||
on_progress(written, expected)
|
||||
async for chunk in chunks:
|
||||
fh.write(chunk)
|
||||
hashes.update(chunk)
|
||||
written += len(chunk)
|
||||
on_progress(written, expected)
|
||||
if expected is not None and written > expected:
|
||||
break
|
||||
if should_stop():
|
||||
fh.flush()
|
||||
raise Interrupted(f"download paused at byte {written}")
|
||||
fh.flush()
|
||||
os.fsync(fh.fileno())
|
||||
except OSError as exc:
|
||||
if exc.errno == errno.ENOSPC:
|
||||
raise DiskSpaceError("disk full while downloading", public="Local disk is full; free space and /resume.") from exc
|
||||
raise TransientError(f"download failed: {exc}", public="Download failed (I/O or network error)") from exc
|
||||
finally:
|
||||
await chunks.aclose()
|
||||
|
||||
size = part.stat().st_size
|
||||
if expected is not None and size != expected:
|
||||
if size > expected:
|
||||
part.unlink(missing_ok=True)
|
||||
raise TransientError(
|
||||
f"incomplete download: {size} of {expected} bytes", public="Download incomplete or corrupted; will retry"
|
||||
)
|
||||
if size == 0:
|
||||
part.unlink(missing_ok=True)
|
||||
raise TransientError("downloaded file is empty", public="Downloaded file is empty; will retry")
|
||||
os.replace(part, final)
|
||||
md5, sha = hashes.digest()
|
||||
return LocalFile(final, size, md5, sha)
|
||||
|
||||
def cleanup(self, rec: VideoRecord) -> int:
|
||||
"""Delete this video's temporary files. Returns how many files were removed."""
|
||||
return sum(self.safe_unlink(path) for path in (self.complete_path(rec), self.partial_path(rec)))
|
||||
|
||||
def safe_unlink(self, path: Path) -> bool:
|
||||
"""Delete ``path`` only if it is a regular temp file directly inside the download dir."""
|
||||
try:
|
||||
if path.is_symlink() or path.parent.resolve() != self._root or not TEMP_NAME.match(path.name):
|
||||
log.error("Refusing to delete %s: not an application temp file", path.name)
|
||||
return False
|
||||
if not path.exists():
|
||||
return False
|
||||
path.unlink()
|
||||
log.info("Deleted temporary file %s", path.name)
|
||||
return True
|
||||
except OSError as exc:
|
||||
log.warning("Could not delete temporary file %s: %s", path.name, exc)
|
||||
return False
|
||||
|
||||
def cleanup_orphans(self, db, *, keep_failed: bool = False) -> int:
|
||||
"""Remove temp files that no unfinished item needs any more.
|
||||
|
||||
A file is deleted only when it matches the application's temp-name pattern,
|
||||
lives directly in the download directory (symlinks are never followed), and
|
||||
its database record is missing or finished (COMPLETED / SKIPPED, or FAILED
|
||||
unless ``keep_failed``). Leftover ``.part`` files of fully downloaded items
|
||||
are removed too. Anything else in the directory is left untouched.
|
||||
"""
|
||||
deleted = 0
|
||||
try:
|
||||
entries = list(os.scandir(self._root))
|
||||
except FileNotFoundError:
|
||||
return 0
|
||||
for entry in entries:
|
||||
match = TEMP_NAME.match(entry.name)
|
||||
if not match or entry.is_symlink() or not entry.is_file(follow_symlinks=False):
|
||||
continue
|
||||
channel_id, message_id, kind = -int(match.group(1)), int(match.group(2)), match.group(3)
|
||||
rec = db.get((channel_id, message_id))
|
||||
if rec is None or rec.state in (State.COMPLETED, State.SKIPPED):
|
||||
reason = "record finished" if rec else "no database record"
|
||||
elif rec.state == State.FAILED and not keep_failed:
|
||||
reason = "item failed (can be downloaded again)"
|
||||
elif kind == "part" and rec.state in (State.DOWNLOADED, State.UPLOADING):
|
||||
reason = "stale partial of a completed download"
|
||||
else:
|
||||
continue
|
||||
if self.safe_unlink(Path(entry.path)):
|
||||
deleted += 1
|
||||
log.info("Orphan cleanup: removed %s (%s)", entry.name, reason)
|
||||
return deleted
|
||||
Reference in New Issue
Block a user