Spaces:
Running
Running
| """Parallel (multi-connection) transfers over MTProto. | |
| The part-scheduling logic (``download_parts`` / ``upload_parts``) only talks to a | |
| tiny ``Transport`` interface, so it is tested offline with fakes. The Telethon | |
| specific glue (``open_transports`` and friends) uses a few private Telethon | |
| helpers; if anything in the fast path fails the wrappers ``download`` / | |
| ``upload`` fall back to Telethon's stock implementation, so the bot keeps | |
| working (just slower). | |
| """ | |
| from __future__ import annotations | |
| import asyncio | |
| import itertools | |
| import logging | |
| import math | |
| import os | |
| import threading | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| from typing import Any, Awaitable, Callable, List, Optional, Protocol, Sequence, Tuple | |
| from .errors import Cancelled | |
| log = logging.getLogger(__name__) | |
| DOWNLOAD_PART = 1024 * 1024 # upload.getFile allows up to 1 MiB per request | |
| UPLOAD_PART = 512 * 1024 # upload.saveBigFilePart parts are at most 512 KiB | |
| BIG_FILE_THRESHOLD = 10 * 1024 * 1024 # Telegram wants saveBigFilePart above 10 MB | |
| RELAY_BUFFER_MIN = 4 # downloaded parts a streamed rename may hold back at least (2 per upload connection otherwise) | |
| PART_TIMEOUT = 120.0 | |
| MAX_RETRIES = 5 | |
| RETRY_BASE_DELAY = 1.0 | |
| # Errors that retrying cannot fix: fail fast and let the caller fall back. | |
| _FATAL_ERROR_NAMES = { | |
| "FileReferenceExpiredError", "FileReferenceInvalidError", "FileIdInvalidError", "LocationInvalidError", | |
| "OffsetInvalidError", "LimitInvalidError", "FilePartsInvalidError", "FilePartSizeInvalidError", | |
| "FileMigrateError", "AuthKeyUnregisteredError", "AuthKeyInvalidError", "SessionRevokedError", | |
| "UserDeactivatedError", "UserDeactivatedBanError", | |
| } | |
| ProgressCb = Callable[[int, int], None] | |
| class TransferStats: | |
| """Which way the transfers since the start went: over several connections, or by Telethon itself (shown by /status).""" | |
| fast_down: int = 0 # downloads over several connections | |
| standard_down: int = 0 # downloads by Telethon itself: parallel transfer is off, or it failed | |
| fast_up: int = 0 # uploads over several connections | |
| small_up: int = 0 # uploads of files up to BIG_FILE_THRESHOLD, which Telethon's own upload handles on purpose | |
| standard_up: int = 0 # uploads by Telethon itself: parallel transfer is off, or it failed | |
| fallbacks: int = 0 # how often the parallel path failed and the standard one took over | |
| last_fallback: str = "" # what failed last, e.g. "download: RuntimeError: ..." | |
| streamed: int = 0 # renames that went straight from the download into the upload, without touching the disk | |
| stream_fallbacks: int = 0 # streamed renames that failed, so the file was saved to disk first | |
| last_stream_fallback: str = "" # what failed last, e.g. "RuntimeError: ..." | |
| def fell_back(self, what: str, exc: BaseException) -> None: | |
| self.fallbacks += 1 | |
| self.last_fallback = f"{what}: {type(exc).__name__}: {exc}" | |
| def stream_failed(self, exc: BaseException) -> None: | |
| self.stream_fallbacks += 1 | |
| self.last_stream_fallback = f"{type(exc).__name__}: {exc}" | |
| stats = TransferStats() | |
| class Transport(Protocol): | |
| async def get_file(self, location: Any, offset: int, limit: int) -> bytes: ... | |
| async def save_part(self, file_id: int, part: int, total_parts: int, data: bytes) -> None: ... | |
| async def close(self) -> None: ... | |
| # --------------------------------------------------------------------------- | |
| # positional file I/O that also works where os.pread/os.pwrite are missing | |
| # --------------------------------------------------------------------------- | |
| class _File: | |
| def __init__(self, path: Path | str, mode: str): | |
| flags = {"r": os.O_RDONLY, "w": os.O_RDWR | os.O_CREAT | os.O_TRUNC}[mode] | |
| flags |= getattr(os, "O_BINARY", 0) | |
| self.fd = os.open(str(path), flags, 0o644) | |
| self._lock = threading.Lock() | |
| self._positional = hasattr(os, "pwrite") and hasattr(os, "pread") | |
| def pwrite(self, data: bytes, offset: int) -> None: | |
| view = memoryview(data) | |
| if self._positional: | |
| written = 0 | |
| while written < len(view): | |
| written += os.pwrite(self.fd, view[written:], offset + written) | |
| return | |
| with self._lock: | |
| os.lseek(self.fd, offset, os.SEEK_SET) | |
| written = 0 | |
| while written < len(view): | |
| written += os.write(self.fd, view[written:]) | |
| def pread(self, size: int, offset: int) -> bytes: | |
| if self._positional: | |
| return os.pread(self.fd, size, offset) | |
| with self._lock: | |
| os.lseek(self.fd, offset, os.SEEK_SET) | |
| return os.read(self.fd, size) | |
| def close(self) -> None: | |
| os.close(self.fd) | |
| # --------------------------------------------------------------------------- | |
| # generic engines | |
| # --------------------------------------------------------------------------- | |
| async def _retry(op: Callable[[], Awaitable[Any]], *, cancel: Optional[asyncio.Event], what: str) -> Any: | |
| delay = RETRY_BASE_DELAY | |
| for attempt in range(MAX_RETRIES + 1): | |
| if cancel is not None and cancel.is_set(): | |
| raise Cancelled() | |
| try: | |
| return await asyncio.wait_for(op(), PART_TIMEOUT) | |
| except (asyncio.CancelledError, Cancelled): | |
| raise | |
| except Exception as exc: # noqa: BLE001 - classified below | |
| if type(exc).__name__ in _FATAL_ERROR_NAMES or attempt >= MAX_RETRIES: | |
| raise | |
| log.warning("%s failed (%s: %s) - retry %d/%d", what, type(exc).__name__, exc, attempt + 1, MAX_RETRIES) | |
| await asyncio.sleep(delay) | |
| delay = min(delay * 2, 8.0) | |
| raise RuntimeError("unreachable") # pragma: no cover | |
| async def _run_workers(workers: Sequence[Callable[[], Awaitable[None]]]) -> None: | |
| tasks = [asyncio.ensure_future(w()) for w in workers] | |
| try: | |
| pending = set(tasks) | |
| while pending: | |
| done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_EXCEPTION) | |
| for task in done: | |
| if not task.cancelled() and task.exception() is not None: | |
| raise task.exception() # type: ignore[misc] | |
| finally: | |
| for task in tasks: | |
| task.cancel() | |
| await asyncio.gather(*tasks, return_exceptions=True) | |
| async def download_parts( | |
| transports: Sequence[Transport], | |
| location: Any, | |
| size: int, | |
| dest: Path | str, | |
| *, | |
| part_size: int = DOWNLOAD_PART, | |
| progress: Optional[ProgressCb] = None, | |
| cancel: Optional[asyncio.Event] = None, | |
| max_bytes: Optional[int] = None, | |
| ) -> int: | |
| """Download ``size`` bytes (or the first ``max_bytes``) into ``dest``; returns bytes written.""" | |
| total = size if max_bytes is None else min(size, max_bytes) | |
| out = _File(dest, "w") | |
| try: | |
| if total <= 0: | |
| return 0 | |
| n_parts = math.ceil(total / part_size) | |
| counter = itertools.count() | |
| done = 0 | |
| async def worker(transport: Transport) -> None: | |
| nonlocal done | |
| while True: | |
| if cancel is not None and cancel.is_set(): | |
| raise Cancelled() | |
| idx = next(counter) | |
| if idx >= n_parts: | |
| return | |
| offset = idx * part_size | |
| want = min(part_size, total - offset) | |
| async def fetch() -> bytes: | |
| data = await transport.get_file(location, offset, part_size) | |
| if len(data) < want: | |
| raise IOError(f"short read: got {len(data)} of {want} bytes at offset {offset}") | |
| return data | |
| data = await _retry(fetch, cancel=cancel, what=f"download part {idx + 1}/{n_parts}") | |
| await asyncio.to_thread(out.pwrite, data[:want], offset) | |
| done += want | |
| if progress is not None: | |
| progress(done, total) | |
| await _run_workers([lambda t=t: worker(t) for t in transports]) | |
| if done != total: | |
| raise IOError(f"download incomplete: {done} of {total} bytes") | |
| return total | |
| finally: | |
| out.close() | |
| async def upload_parts( | |
| transports: Sequence[Transport], | |
| path: Path | str, | |
| size: int, | |
| file_id: int, | |
| *, | |
| part_size: int = UPLOAD_PART, | |
| progress: Optional[ProgressCb] = None, | |
| cancel: Optional[asyncio.Event] = None, | |
| ) -> int: | |
| """Upload ``path`` with ``saveBigFilePart``; returns the number of parts.""" | |
| n_parts = max(1, math.ceil(size / part_size)) | |
| src = _File(path, "r") | |
| try: | |
| counter = itertools.count() | |
| done = 0 | |
| async def worker(transport: Transport) -> None: | |
| nonlocal done | |
| while True: | |
| if cancel is not None and cancel.is_set(): | |
| raise Cancelled() | |
| idx = next(counter) | |
| if idx >= n_parts: | |
| return | |
| offset = idx * part_size | |
| want = min(part_size, size - offset) | |
| data = await asyncio.to_thread(src.pread, want, offset) | |
| if len(data) != want: | |
| raise IOError(f"read {len(data)} of {want} bytes at offset {offset}; was the file modified?") | |
| await _retry( | |
| lambda d=data, i=idx: transport.save_part(file_id, i, n_parts, d), | |
| cancel=cancel, | |
| what=f"upload part {idx + 1}/{n_parts}", | |
| ) | |
| done += want | |
| if progress is not None: | |
| progress(done, size) | |
| await _run_workers([lambda t=t: worker(t) for t in transports]) | |
| if done != size: | |
| raise IOError(f"upload incomplete: {done} of {size} bytes") | |
| return n_parts | |
| finally: | |
| src.close() | |
| async def relay_parts( | |
| down: Sequence[Transport], | |
| up: Sequence[Transport], | |
| location: Any, | |
| size: int, | |
| file_id: int, | |
| *, | |
| part_size: int = DOWNLOAD_PART, | |
| send_size: int = UPLOAD_PART, | |
| progress: Optional[ProgressCb] = None, | |
| cancel: Optional[asyncio.Event] = None, | |
| ) -> int: | |
| """Copy ``size`` bytes of ``location`` into an upload without saving them anywhere; returns the number of upload parts. | |
| The ``down`` connections fetch ``part_size`` pieces and hand them to the ``up`` connections through a small | |
| buffer, which cut each piece into ``send_size`` parts for ``saveBigFilePart``. Parts can be uploaded in any | |
| order, so the two sides only wait for each other when the buffer is full or empty: the file takes about as long as | |
| the slower direction, not as long as both together. Progress counts what has been uploaded. | |
| """ | |
| if size <= 0: | |
| raise ValueError("there is nothing to copy") | |
| if part_size % send_size: | |
| raise ValueError("a download part must hold a whole number of upload parts") | |
| if not down or not up: | |
| raise ValueError("a streamed copy needs connections both ways") | |
| n_fetch = math.ceil(size / part_size) | |
| n_parts = math.ceil(size / send_size) | |
| per_piece = part_size // send_size | |
| buffer: asyncio.Queue = asyncio.Queue(maxsize=max(RELAY_BUFFER_MIN, 2 * len(up))) | |
| fetch_counter, take_counter = itertools.count(), itertools.count() | |
| sent = 0 | |
| async def fetch_worker(transport: Transport) -> None: | |
| while True: | |
| if cancel is not None and cancel.is_set(): | |
| raise Cancelled() | |
| idx = next(fetch_counter) | |
| if idx >= n_fetch: | |
| return | |
| offset = idx * part_size | |
| want = min(part_size, size - offset) | |
| async def fetch() -> bytes: | |
| data = await transport.get_file(location, offset, part_size) | |
| if len(data) < want: | |
| raise IOError(f"short read: got {len(data)} of {want} bytes at offset {offset}") | |
| return data | |
| data = await _retry(fetch, cancel=cancel, what=f"relay download part {idx + 1}/{n_fetch}") | |
| await buffer.put((idx, data[:want])) | |
| async def send_worker(transport: Transport) -> None: | |
| nonlocal sent | |
| while True: | |
| if cancel is not None and cancel.is_set(): | |
| raise Cancelled() | |
| if next(take_counter) >= n_fetch: # every downloaded part already has a sender | |
| return | |
| idx, data = await buffer.get() | |
| for start in range(0, len(data), send_size): | |
| part = idx * per_piece + start // send_size | |
| piece = data[start:start + send_size] | |
| await _retry( | |
| lambda p=piece, i=part: transport.save_part(file_id, i, n_parts, p), | |
| cancel=cancel, | |
| what=f"relay upload part {part + 1}/{n_parts}", | |
| ) | |
| sent += len(piece) | |
| if progress is not None: | |
| progress(sent, size) | |
| await _run_workers([lambda t=t: fetch_worker(t) for t in down] + [lambda t=t: send_worker(t) for t in up]) | |
| if sent != size: | |
| raise IOError(f"relay incomplete: {sent} of {size} bytes uploaded") | |
| return n_parts | |
| async def close_all(transports: Sequence[Transport]) -> None: | |
| await asyncio.gather(*(t.close() for t in transports), return_exceptions=True) | |
| # --------------------------------------------------------------------------- | |
| # Telethon glue | |
| # --------------------------------------------------------------------------- | |
| class TelethonTransport: | |
| def __init__(self, client: Any, sender: Any): | |
| self._client = client | |
| self._sender = sender | |
| async def get_file(self, location: Any, offset: int, limit: int) -> bytes: | |
| from telethon.tl.functions.upload import GetFileRequest | |
| result = await self._client._call(self._sender, GetFileRequest(location=location, offset=offset, limit=limit)) | |
| if not hasattr(result, "bytes"): | |
| raise IOError(f"unexpected upload.getFile result: {type(result).__name__}") | |
| return result.bytes | |
| async def save_part(self, file_id: int, part: int, total_parts: int, data: bytes) -> None: | |
| from telethon.tl.functions.upload import SaveBigFilePartRequest | |
| ok = await self._client._call(self._sender, SaveBigFilePartRequest(file_id, part, total_parts, data)) | |
| if not ok: | |
| raise IOError("upload.saveBigFilePart returned False") | |
| async def close(self) -> None: | |
| await self._sender.disconnect() | |
| async def _open_sender(client: Any, dc_id: int, auth_key: Any) -> Any: | |
| from telethon.network import MTProtoSender | |
| dc = await client._get_dc(dc_id) | |
| sender = MTProtoSender(auth_key, loggers=client._log) | |
| await sender.connect( | |
| client._connection( | |
| dc.ip_address, dc.port, dc.id, loggers=client._log, proxy=client._proxy, local_addr=client._local_addr | |
| ) | |
| ) | |
| return sender | |
| async def open_transports(client: Any, dc_id: int, connections: int) -> List[TelethonTransport]: | |
| """Open up to ``connections`` extra MTProto connections to ``dc_id``.""" | |
| from telethon.tl.alltlobjects import LAYER | |
| from telethon.tl.functions import InvokeWithLayerRequest | |
| from telethon.tl.functions.auth import ExportAuthorizationRequest, ImportAuthorizationRequest | |
| if dc_id == client.session.dc_id: | |
| auth_key = client.session.auth_key | |
| first = await _open_sender(client, dc_id, auth_key) | |
| else: | |
| # a different DC needs its own key, plus our authorization exported/imported once | |
| first = await _open_sender(client, dc_id, None) | |
| try: | |
| exported = await client(ExportAuthorizationRequest(dc_id)) | |
| client._init_request.query = ImportAuthorizationRequest(id=exported.id, bytes=exported.bytes) | |
| await first.send(InvokeWithLayerRequest(LAYER, client._init_request)) | |
| except BaseException: | |
| await first.disconnect() | |
| raise | |
| auth_key = first.auth_key | |
| senders = [first] | |
| if connections > 1: | |
| results = await asyncio.gather( | |
| *(_open_sender(client, dc_id, auth_key) for _ in range(connections - 1)), return_exceptions=True | |
| ) | |
| for res in results: | |
| if isinstance(res, BaseException): | |
| log.warning("could not open an extra connection to DC %s: %s", dc_id, res) | |
| else: | |
| senders.append(res) | |
| log.info("opened %d connection(s) to DC %s", len(senders), dc_id) | |
| return [TelethonTransport(client, s) for s in senders] | |
| async def fast_download( | |
| client: Any, | |
| document: Any, | |
| dest: Path | str, | |
| *, | |
| connections: int = 8, | |
| progress: Optional[ProgressCb] = None, | |
| cancel: Optional[asyncio.Event] = None, | |
| max_bytes: Optional[int] = None, | |
| ) -> int: | |
| from telethon import utils | |
| dc_id, location = utils.get_input_location(document) | |
| transports = await open_transports(client, dc_id, connections) | |
| try: | |
| written = await download_parts( | |
| transports, location, int(document.size), dest, progress=progress, cancel=cancel, max_bytes=max_bytes | |
| ) | |
| finally: | |
| await close_all(transports) | |
| if max_bytes is None: # reading the start of a file is not a transfer | |
| stats.fast_down += 1 | |
| return written | |
| async def fast_upload( | |
| client: Any, | |
| path: Path | str, | |
| file_name: str, | |
| *, | |
| connections: int = 8, | |
| progress: Optional[ProgressCb] = None, | |
| cancel: Optional[asyncio.Event] = None, | |
| ) -> Any: | |
| from telethon.tl.types import InputFileBig | |
| size = os.path.getsize(path) | |
| if size <= BIG_FILE_THRESHOLD: | |
| small = await client.upload_file(str(path), file_name=file_name, progress_callback=_checked(progress, cancel)) | |
| stats.small_up += 1 | |
| return small | |
| file_id = int.from_bytes(os.urandom(8), "big", signed=True) | |
| transports = await open_transports(client, client.session.dc_id, connections) | |
| try: | |
| parts = await upload_parts(transports, path, size, file_id, progress=progress, cancel=cancel) | |
| finally: | |
| await close_all(transports) | |
| stats.fast_up += 1 | |
| return InputFileBig(id=file_id, parts=parts, name=file_name) | |
| async def open_relay_transports( | |
| client: Any, dc_id: int, connections: int | |
| ) -> Tuple[List[TelethonTransport], List[TelethonTransport]]: | |
| """Connections to fetch from ``dc_id`` and connections to upload to our own DC, opened together.""" | |
| results = await asyncio.gather( | |
| open_transports(client, dc_id, connections), | |
| open_transports(client, client.session.dc_id, connections), | |
| return_exceptions=True, | |
| ) | |
| failed = next((r for r in results if isinstance(r, BaseException)), None) | |
| if failed is not None: | |
| for r in results: | |
| if not isinstance(r, BaseException): | |
| await close_all(r) | |
| raise failed | |
| return results[0], results[1] # type: ignore[return-value] | |
| async def relay( | |
| client: Any, | |
| document: Any, | |
| file_name: str, | |
| *, | |
| connections: int = 8, | |
| progress: Optional[ProgressCb] = None, | |
| cancel: Optional[asyncio.Event] = None, | |
| ) -> Any: | |
| """Upload a copy of ``document`` named ``file_name`` while it downloads, without saving it; returns an ``InputFileBig``. | |
| Telethon has no equivalent to fall back to here: when this fails the caller downloads and uploads the usual way. | |
| """ | |
| from telethon import utils | |
| from telethon.tl.types import InputFileBig | |
| dc_id, location = utils.get_input_location(document) | |
| size = int(document.size) | |
| file_id = int.from_bytes(os.urandom(8), "big", signed=True) | |
| down, up = await open_relay_transports(client, dc_id, connections) | |
| try: | |
| parts = await relay_parts(down, up, location, size, file_id, progress=progress, cancel=cancel) | |
| finally: | |
| await close_all(down + up) | |
| return InputFileBig(id=file_id, parts=parts, name=file_name) | |
| def is_file_part_error(exc: BaseException) -> bool: | |
| """Is this Telegram refusing an upload because of its parts (one missing, or the wrong size or number of them)?""" | |
| return type(exc).__name__.startswith("FilePart") | |
| # --------------------------------------------------------------------------- | |
| # public wrappers with fallback | |
| # --------------------------------------------------------------------------- | |
| def _checked(progress: Optional[ProgressCb], cancel: Optional[asyncio.Event]) -> ProgressCb: | |
| """Progress callback for Telethon's own transfers; raising aborts them.""" | |
| def cb(current: int, total: int) -> None: | |
| if cancel is not None and cancel.is_set(): | |
| raise Cancelled() | |
| if progress is not None: | |
| progress(current, total) | |
| return cb | |
| def _unlink(path: Path | str) -> None: | |
| try: | |
| os.unlink(path) | |
| except FileNotFoundError: | |
| pass | |
| async def download( | |
| client: Any, | |
| message: Any, | |
| dest: Path | str, | |
| *, | |
| fast: bool = True, | |
| connections: int = 8, | |
| progress: Optional[ProgressCb] = None, | |
| cancel: Optional[asyncio.Event] = None, | |
| max_bytes: Optional[int] = None, | |
| ) -> None: | |
| """Download a message's media to ``dest`` (optionally only the first ``max_bytes``).""" | |
| if fast and getattr(message, "document", None) is not None: | |
| try: | |
| await fast_download( | |
| client, message.document, dest, connections=connections, progress=progress, cancel=cancel, | |
| max_bytes=max_bytes, | |
| ) | |
| return | |
| except Cancelled: | |
| raise | |
| except Exception as exc: # noqa: BLE001 | |
| log.exception("fast download failed (%s: %s); using Telethon's standard download", type(exc).__name__, exc) | |
| stats.fell_back("download", exc) | |
| _unlink(dest) | |
| if max_bytes is not None: | |
| request = 512 * 1024 | |
| with open(dest, "wb") as fh: | |
| async for chunk in client.iter_download( | |
| message.document, request_size=request, limit=max(1, math.ceil(max_bytes / request)) | |
| ): | |
| fh.write(chunk) | |
| if cancel is not None and cancel.is_set(): | |
| raise Cancelled() | |
| fh.truncate(min(fh.tell(), max_bytes)) | |
| return | |
| await client.download_media(message, file=str(dest), progress_callback=_checked(progress, cancel)) | |
| stats.standard_down += 1 | |
| async def upload( | |
| client: Any, | |
| path: Path | str, | |
| file_name: str, | |
| *, | |
| fast: bool = True, | |
| connections: int = 8, | |
| progress: Optional[ProgressCb] = None, | |
| cancel: Optional[asyncio.Event] = None, | |
| ) -> Any: | |
| """Upload ``path``; returns an ``InputFile``/``InputFileBig`` for ``send_file``.""" | |
| if fast: | |
| try: | |
| return await fast_upload(client, path, file_name, connections=connections, progress=progress, cancel=cancel) | |
| except Cancelled: | |
| raise | |
| except Exception as exc: # noqa: BLE001 | |
| log.exception("fast upload failed (%s: %s); using Telethon's standard upload", type(exc).__name__, exc) | |
| stats.fell_back("upload", exc) | |
| standard = await client.upload_file(str(path), file_name=file_name, progress_callback=_checked(progress, cancel)) | |
| stats.standard_up += 1 | |
| return standard | |