From d27cd758e084d844819c87c84582952d19e9811f Mon Sep 17 00:00:00 2001 From: VulcanoSoftware Date: Wed, 9 Sep 2026 09:48:48 +0000 Subject: [PATCH 1/2] perf(sftp): speed up downloads via attachment caching, wider parallelism and cheaper buffering Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- dcfs/app/sftp/__init__.py | 25 +++- dcfs/app/sftp/handler.py | 121 +++++++++++------- dcfs/core/api/message/__init__.py | 23 +++- .../repository/impl/file_content/__init__.py | 2 +- dcfs/discord/impl/discord_bot.py | 115 +++++++++++++++-- tests/dcfs/app/test_sftp_buffered_read.py | 70 ++++++++++ .../dcfs/core/api/message/test_split_count.py | 26 ++++ tests/dcfs/discord/__init__.py | 0 tests/dcfs/discord/test_attachment_cache.py | 65 ++++++++++ 9 files changed, 390 insertions(+), 57 deletions(-) create mode 100644 tests/dcfs/app/test_sftp_buffered_read.py create mode 100644 tests/dcfs/core/api/message/test_split_count.py create mode 100644 tests/dcfs/discord/__init__.py create mode 100644 tests/dcfs/discord/test_attachment_cache.py diff --git a/dcfs/app/sftp/__init__.py b/dcfs/app/sftp/__init__.py index 1fb3b20..755c8be 100644 --- a/dcfs/app/sftp/__init__.py +++ b/dcfs/app/sftp/__init__.py @@ -2,6 +2,7 @@ import os import asyncssh +from asyncssh.encryption import get_encryption_algs from dcfs.config import DATA_DIR, Config from dcfs.core import Clients @@ -10,6 +11,24 @@ logger = logging.getLogger(__name__) +# AES-GCM runs on the CPU's AES instructions and encrypts roughly five times +# faster than asyncssh's ChaCha20-Poly1305, which is the default first choice. +# Every other algorithm keeps its original relative order so no client loses +# compatibility. +PREFERRED_ENCRYPTION_ALGS = ( + "aes128-gcm@openssh.com", + "aes256-gcm@openssh.com", + "aes128-ctr", + "aes256-ctr", +) + + +def _encryption_algs() -> list[str]: + supported = [alg.decode("ascii") for alg in get_encryption_algs()] + preferred = [alg for alg in PREFERRED_ENCRYPTION_ALGS if alg in supported] + return preferred + [alg for alg in supported if alg not in preferred] + + class DCFSSFTPServer(asyncssh.SSHServer): def __init__(self, clients: Clients, config: Config): self.clients = clients @@ -58,7 +77,11 @@ def sftp_factory(channel): config.dcfs.sftp.host, config.dcfs.sftp.port, server_host_keys=[host_key], - sftp_factory=sftp_factory + sftp_factory=sftp_factory, + encryption_algs=_encryption_algs(), + # File content is already compressed and/or encrypted, so compression + # only burns CPU that the transfer needs. + compression_algs=["none"], ) async def run_sftp_server(server: asyncssh.SSHListener, host: str, port: int): diff --git a/dcfs/app/sftp/handler.py b/dcfs/app/sftp/handler.py index 674e051..91794e0 100644 --- a/dcfs/app/sftp/handler.py +++ b/dcfs/app/sftp/handler.py @@ -4,6 +4,7 @@ import os import stat import time +from collections import deque from typing import Any, AsyncGenerator, AsyncIterator, Optional, cast import asyncssh @@ -333,10 +334,13 @@ def __init__(self, ops: Ops, path: str, mode: str, client_name: str): self.buffer = bytearray() self.closed = False - # Read streaming state + # Read streaming state. Buffered data is kept as the chunks handed + # over by the download stream, tagged with their absolute file + # offset, so serving a read never copies or shifts the whole buffer. self._read_stream: Optional[AsyncIterator[bytes]] = None self._buf_offset = 0 - self._read_buf = bytearray() + self._buf_len = 0 + self._chunks: deque[tuple[int, bytes]] = deque() self._read_lock = asyncio.Lock() self._cached_attrs: Optional[asyncssh.SFTPAttrs] = None self._highest_offset = 0 @@ -380,8 +384,59 @@ async def _stop_prefetch(self) -> None: self._prefetch_queue = None self._prefetch_eof = False + def _buf_end(self) -> int: + return self._buf_offset + self._buf_len + + def _append_chunk(self, chunk: bytes) -> None: + if not isinstance(chunk, bytes): + chunk = bytes(chunk) + self._chunks.append((self._buf_end(), chunk)) + self._buf_len += len(chunk) + + def _extract(self, offset: int, size: int) -> bytes: + """Copy ``size`` bytes starting at ``offset`` out of the buffer.""" + if size <= 0 or offset < self._buf_offset or offset >= self._buf_end(): + return b"" + + out: Optional[bytearray] = None + pos = offset + remaining = size + for start, chunk in self._chunks: + end = start + len(chunk) + if end <= pos: + continue + if start > pos: + break + rel = pos - start + piece = chunk[rel : rel + remaining] + if out is None and len(piece) == remaining: + # Whole request satisfied by a single chunk: no extra copy. + return piece + if out is None: + out = bytearray() + out += piece + remaining -= len(piece) + pos += len(piece) + if remaining == 0: + break + + return bytes(out) if out else b"" + + def _prune(self) -> None: + """Drop chunks that fall behind the backward-retain window.""" + keep_from = self._highest_offset - self.MAX_BACKWARD_RETAIN + while self._chunks: + start, chunk = self._chunks[0] + end = start + len(chunk) + if end > keep_from: + break + self._chunks.popleft() + self._buf_len -= len(chunk) + self._buf_offset = end + async def _start_prefetch(self, offset: int) -> None: - self._read_buf = bytearray() + self._chunks.clear() + self._buf_len = 0 self._buf_offset = offset self._highest_offset = offset self._read_stream = await self.ops.download( @@ -397,16 +452,27 @@ async def _start_prefetch(self, offset: int) -> None: self._run_prefetch(self._read_stream, self._prefetch_queue) ) + async def _fill_until(self, target_end: int) -> None: + while self._buf_end() < target_end and not self._prefetch_eof: + if self._prefetch_queue is None: + break + item = await self._prefetch_queue.get() + if item is None: + self._prefetch_eof = True + break + if isinstance(item, Exception): + self._prefetch_eof = True + raise item + self._append_chunk(item) + async def read(self, offset: int, size: int) -> bytes: if "r" not in self.mode: raise asyncssh.SFTPPermissionDenied("File not open for reading") async with self._read_lock: - buf_end = self._buf_offset + len(self._read_buf) - can_reuse_stream = ( self._read_stream is not None - and self._buf_offset <= offset <= buf_end + self.MAX_FORWARD_SKIP + and self._buf_offset <= offset <= self._buf_end() + self.MAX_FORWARD_SKIP ) if not can_reuse_stream: @@ -414,49 +480,18 @@ async def read(self, offset: int, size: int) -> bytes: await self._start_prefetch(offset) target_end = offset + size - while (self._buf_offset + len(self._read_buf) < target_end) and not self._prefetch_eof: - if self._prefetch_queue is None: - break - item = await self._prefetch_queue.get() - if item is None: - self._prefetch_eof = True - break - if isinstance(item, Exception): - self._prefetch_eof = True - raise item - self._read_buf.extend(item) + await self._fill_until(target_end) # If stream reached EOF before reaching offset, restart stream at offset - if self._prefetch_eof and (self._buf_offset + len(self._read_buf) <= offset) and size > 0: + if self._prefetch_eof and self._buf_end() <= offset and size > 0: await self._stop_prefetch() await self._start_prefetch(offset) - while (self._buf_offset + len(self._read_buf) < target_end) and not self._prefetch_eof: - if self._prefetch_queue is None: - break - item = await self._prefetch_queue.get() - if item is None: - self._prefetch_eof = True - break - if isinstance(item, Exception): - self._prefetch_eof = True - raise item - self._read_buf.extend(item) - - rel_offset = offset - self._buf_offset - if rel_offset >= 0 and rel_offset < len(self._read_buf): - data = bytes(self._read_buf[rel_offset : rel_offset + size]) - else: - data = b"" + await self._fill_until(target_end) - self._highest_offset = max(self._highest_offset, offset + len(data)) + data = self._extract(offset, size) - # Prune buffer behind prune_target to keep memory bounded - prune_target = self._highest_offset - self.MAX_BACKWARD_RETAIN - if prune_target > self._buf_offset: - discard = min(prune_target - self._buf_offset, len(self._read_buf)) - if discard > 0: - self._read_buf = self._read_buf[discard:] - self._buf_offset += discard + self._highest_offset = max(self._highest_offset, offset + len(data)) + self._prune() return data diff --git a/dcfs/core/api/message/__init__.py b/dcfs/core/api/message/__init__.py index da59b5f..c07a3b6 100644 --- a/dcfs/core/api/message/__init__.py +++ b/dcfs/core/api/message/__init__.py @@ -29,6 +29,13 @@ DELETE_BATCH_SIZE = 100 DISCORD_MSG_LIMIT = 4000 + +# Bounds on the number of concurrent CDN range requests used to fetch a single +# message attachment, and how many chunks each of them may run ahead. +MIN_PARALLEL_RANGE_REQUESTS = 4 +MAX_PARALLEL_RANGE_REQUESTS = 8 +PARALLEL_RANGE_TARGET_SIZE = 2 * 1024 * 1024 +PARALLEL_QUEUE_MAXSIZE = 32 OVERFLOW_SENTINEL = "DCFS_OVERFLOW" OVERFLOW_FILENAME = "overflow.json" @@ -175,10 +182,22 @@ def split_download_tasks( def _size(begin: int, end: int) -> int: return end - begin + 1 + @staticmethod + def _parallel_split_count(size: int) -> int: + """Number of concurrent CDN range requests for a ``size`` byte range. + + Scales with the range size so large parts use more connections, while + staying bounded to avoid per-request overhead outweighing the added + bandwidth. + """ + scaled = size // PARALLEL_RANGE_TARGET_SIZE + return max(MIN_PARALLEL_RANGE_REQUESTS, + min(MAX_PARALLEL_RANGE_REQUESTS, scaled)) + async def download_file_parallel(self, message_id: int, begin: int, end: int): # Split the range into concurrent sub-range downloads so we can # utilise CDN bandwidth better for large single-part files. - n = 4 + n = self._parallel_split_count(self._size(begin, end)) sub_ranges = list(self.split_download_tasks(begin, end, n)) resps = await asyncio.gather(*[ @@ -195,7 +214,7 @@ async def download_file_parallel(self, message_id: int, begin: int, end: int): ]) queues: list[asyncio.Queue[object]] = [ - asyncio.Queue(maxsize=32) for _ in range(n) + asyncio.Queue(maxsize=PARALLEL_QUEUE_MAXSIZE) for _ in range(n) ] async def _producer( diff --git a/dcfs/core/repository/impl/file_content/__init__.py b/dcfs/core/repository/impl/file_content/__init__.py index 88a3286..1f2272e 100644 --- a/dcfs/core/repository/impl/file_content/__init__.py +++ b/dcfs/core/repository/impl/file_content/__init__.py @@ -26,7 +26,7 @@ async def _empty_iterator() -> AsyncIterator[bytes]: class DCMsgFileContentRepository(IFileContentRepository): - def __init__(self, message_api: MessageApi, max_concurrent_parts: int = 3): + def __init__(self, message_api: MessageApi, max_concurrent_parts: int = 6): self._message_api = message_api self._download_semaphore = asyncio.Semaphore(max_concurrent_parts) self._upload_semaphore = asyncio.Semaphore(max_concurrent_parts) diff --git a/dcfs/discord/impl/discord_bot.py b/dcfs/discord/impl/discord_bot.py index e3bdb3b..7ae5f8e 100644 --- a/dcfs/discord/impl/discord_bot.py +++ b/dcfs/discord/impl/discord_bot.py @@ -1,7 +1,10 @@ import asyncio import io import logging -from typing import Any, List, Optional +import time +from dataclasses import dataclass +from typing import Any, Dict, List, Optional +from urllib.parse import parse_qs, urlparse import aiohttp import discord @@ -34,6 +37,31 @@ CHUNK_SIZE = 1024 * 1024 # 1 MB chunks for downloads +# Attachment URLs are signed by Discord and stay valid for ~24h. Caching them +# removes a REST round trip per downloaded part (and per parallel sub-range), +# which dominates the latency of multi-part downloads. +ATTACHMENT_CACHE_TTL = 15 * 60 +ATTACHMENT_EXPIRY_MARGIN = 60 +ATTACHMENT_CACHE_CAPACITY = 4096 + + +@dataclass +class _CachedAttachment: + url: str + size: int + expires_at: float + + +def _url_expiry(url: str) -> Optional[float]: + """Return the signed-URL expiry (unix seconds) encoded in ``ex=``.""" + raw = parse_qs(urlparse(url).query).get("ex", [None])[0] + if not raw: + return None + try: + return float(int(raw, 16)) + except ValueError: + return None + class DiscordBotAPI(IDiscordClient): def __init__(self, bot: discord.Client, bot_token: str): @@ -41,12 +69,82 @@ def __init__(self, bot: discord.Client, bot_token: str): self._bot = bot self._bot_token = bot_token self._http_session: Optional[aiohttp.ClientSession] = None + self._attachment_cache: Dict[int, _CachedAttachment] = {} + self._attachment_locks: Dict[int, asyncio.Lock] = {} async def _ensure_http_session(self) -> aiohttp.ClientSession: if self._http_session is None or self._http_session.closed: - self._http_session = aiohttp.ClientSession() + connector = aiohttp.TCPConnector( + limit=0, + limit_per_host=0, + ttl_dns_cache=300, + enable_cleanup_closed=True, + ) + self._http_session = aiohttp.ClientSession( + connector=connector, + read_bufsize=CHUNK_SIZE, + ) return self._http_session + async def _resolve_attachment( + self, channel: Any, message_id: int + ) -> _CachedAttachment: + """Resolve a message's attachment URL, caching the REST lookup. + + Concurrent callers for the same message (the parallel sub-range + downloads) share a single lookup instead of each issuing their own + ``fetch_message`` request. + """ + now = time.time() + cached = self._attachment_cache.get(message_id) + if cached is not None and cached.expires_at > now: + return cached + + lock = self._attachment_locks.setdefault(message_id, asyncio.Lock()) + async with lock: + cached = self._attachment_cache.get(message_id) + now = time.time() + if cached is not None and cached.expires_at > now: + return cached + + try: + msg = await channel.fetch_message(message_id) + except discord.NotFound: + raise MessageNotFound(message_id) + if not msg.attachments: + raise UnDownloadableMessage(message_id) + attachment = msg.attachments[0] + + expires_at = now + ATTACHMENT_CACHE_TTL + signed_until = _url_expiry(attachment.url) + if signed_until is not None: + expires_at = min( + expires_at, signed_until - ATTACHMENT_EXPIRY_MARGIN + ) + + entry = _CachedAttachment( + url=attachment.url, size=attachment.size, expires_at=expires_at + ) + self._prune_attachment_cache() + self._attachment_cache[message_id] = entry + return entry + + def _prune_attachment_cache(self) -> None: + if len(self._attachment_cache) < ATTACHMENT_CACHE_CAPACITY: + return + now = time.time() + stale = [ + mid + for mid, entry in self._attachment_cache.items() + if entry.expires_at <= now + ] + if not stale: + # Nothing expired yet: drop the oldest insertions instead. + stale = list(self._attachment_cache)[: ATTACHMENT_CACHE_CAPACITY // 4] + for mid in stale: + self._attachment_cache.pop(mid, None) + self._attachment_locks.pop(mid, None) + async def _get_channel(self, channel_id: int) -> Any: channel = self._bot.get_channel(channel_id) if channel is None: @@ -154,13 +252,7 @@ async def edit_message_media(self, req: EditMessageMediaReq) -> Message: async def download_file(self, req: DownloadFileReq) -> DownloadFileResp: channel_id = self._parse_channel_id(req.chat) channel = await self._get_channel(channel_id) - try: - msg = await channel.fetch_message(req.message_id) - except discord.NotFound: - raise MessageNotFound(req.message_id) - if not msg.attachments: - raise UnDownloadableMessage(req.message_id) - attachment = msg.attachments[0] + attachment = await self._resolve_attachment(channel, req.message_id) session = await self._ensure_http_session() @@ -182,9 +274,12 @@ async def download_file(self, req: DownloadFileReq) -> DownloadFileResp: # timeout a slow or hung CDN connection would cause the whole # WebDAV GET handler to hang indefinitely, making WinSCP / the # client time out with a generic "connection timed out" error. + # Timeout: connect within 15s and require progress every 60s. A total + # deadline is deliberately avoided: it would abort otherwise healthy + # long-running range downloads on slower links. timeout = aiohttp.ClientTimeout( connect=15.0, - total=120.0, + sock_read=60.0, ) t0 = asyncio.get_event_loop().time() response = await session.get(url, headers=headers, timeout=timeout) diff --git a/tests/dcfs/app/test_sftp_buffered_read.py b/tests/dcfs/app/test_sftp_buffered_read.py new file mode 100644 index 0000000..0cc9112 --- /dev/null +++ b/tests/dcfs/app/test_sftp_buffered_read.py @@ -0,0 +1,70 @@ +from unittest.mock import MagicMock + +import pytest + +from dcfs.app.sftp.handler import DCFSSFTPBufferedFile + +DATA = bytes((i * 7) % 251 for i in range(6 * 1024 * 1024)) +CHUNK = 64 * 1024 + + +def _file(data: bytes = DATA) -> tuple[DCFSSFTPBufferedFile, list[int]]: + ops = MagicMock() + downloads: list[int] = [] + + async def download(path, offset, end, name, validate=False): + downloads.append(offset) + + async def gen(): + for pos in range(offset, len(data), CHUNK): + yield data[pos : pos + CHUNK] + + return gen() + + ops.download = download + return DCFSSFTPBufferedFile(ops, "/c/file.bin", "r", "c"), downloads + + +@pytest.mark.asyncio +async def test_sequential_reads_return_the_whole_file(): + f, downloads = _file() + out = bytearray() + offset = 0 + while True: + data = await f.read(offset, 256 * 1024) + if not data: + break + out += data + offset += len(data) + + assert bytes(out) == DATA + assert downloads[0] == 0 + + +@pytest.mark.asyncio +async def test_read_spanning_chunk_boundaries(): + f, _ = _file() + await f.read(0, 10) + data = await f.read(CHUNK - 5, CHUNK + 10) + + assert data == DATA[CHUNK - 5 : 2 * CHUNK + 5] + + +@pytest.mark.asyncio +async def test_backward_seek_restarts_the_stream(): + f, downloads = _file() + far = DCFSSFTPBufferedFile.MAX_BACKWARD_RETAIN + 512 * 1024 + offset = 0 + while offset < far: + offset += len(await f.read(offset, 256 * 1024)) + + data = await f.read(0, 32) + + assert data == DATA[:32] + assert downloads == [0, 0] + + +@pytest.mark.asyncio +async def test_read_past_end_of_file_returns_empty(): + f, _ = _file(data=DATA[: 8 * 1024]) + assert await f.read(16 * 1024, 1024) == b"" diff --git a/tests/dcfs/core/api/message/test_split_count.py b/tests/dcfs/core/api/message/test_split_count.py new file mode 100644 index 0000000..bb940a8 --- /dev/null +++ b/tests/dcfs/core/api/message/test_split_count.py @@ -0,0 +1,26 @@ +from dcfs.core.api.message import ( + MAX_PARALLEL_RANGE_REQUESTS, + MIN_PARALLEL_RANGE_REQUESTS, + PARALLEL_RANGE_TARGET_SIZE, + MessageApi, +) + + +def test_small_ranges_use_the_minimum(): + assert MessageApi._parallel_split_count(0) == MIN_PARALLEL_RANGE_REQUESTS + assert MessageApi._parallel_split_count(1) == MIN_PARALLEL_RANGE_REQUESTS + assert ( + MessageApi._parallel_split_count(PARALLEL_RANGE_TARGET_SIZE) + == MIN_PARALLEL_RANGE_REQUESTS + ) + + +def test_split_count_scales_with_size(): + assert MessageApi._parallel_split_count(6 * PARALLEL_RANGE_TARGET_SIZE) == 6 + + +def test_split_count_is_capped(): + assert ( + MessageApi._parallel_split_count(1024 * PARALLEL_RANGE_TARGET_SIZE) + == MAX_PARALLEL_RANGE_REQUESTS + ) diff --git a/tests/dcfs/discord/__init__.py b/tests/dcfs/discord/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/dcfs/discord/test_attachment_cache.py b/tests/dcfs/discord/test_attachment_cache.py new file mode 100644 index 0000000..b6f9003 --- /dev/null +++ b/tests/dcfs/discord/test_attachment_cache.py @@ -0,0 +1,65 @@ +import asyncio +import time +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from dcfs.discord.impl.discord_bot import DiscordBotAPI, _url_expiry + + +def _channel(url: str = "https://cdn.example/file.bin", size: int = 10): + attachment = MagicMock() + attachment.url = url + attachment.size = size + message = MagicMock() + message.attachments = [attachment] + + channel = MagicMock() + channel.fetch_message = AsyncMock(return_value=message) + return channel + + +def test_url_expiry_parses_signed_urls(): + assert _url_expiry("https://cdn.example/f.bin?ex=68b0a1c0&is=1") == float( + 0x68B0A1C0 + ) + assert _url_expiry("https://cdn.example/f.bin") is None + assert _url_expiry("https://cdn.example/f.bin?ex=zzz") is None + + +@pytest.mark.asyncio +async def test_concurrent_lookups_share_a_single_fetch(): + api = DiscordBotAPI(MagicMock(), "token") + channel = _channel() + + resolved = await asyncio.gather( + *[api._resolve_attachment(channel, 42) for _ in range(8)] + ) + + assert channel.fetch_message.await_count == 1 + assert {r.url for r in resolved} == {"https://cdn.example/file.bin"} + + +@pytest.mark.asyncio +async def test_cached_url_is_reused_until_it_expires(): + api = DiscordBotAPI(MagicMock(), "token") + channel = _channel() + + await api._resolve_attachment(channel, 42) + await api._resolve_attachment(channel, 42) + assert channel.fetch_message.await_count == 1 + + api._attachment_cache[42].expires_at = time.time() - 1 + await api._resolve_attachment(channel, 42) + assert channel.fetch_message.await_count == 2 + + +@pytest.mark.asyncio +async def test_cache_never_outlives_the_signed_url(): + signed_until = time.time() + 30 + api = DiscordBotAPI(MagicMock(), "token") + channel = _channel(url=f"https://cdn.example/f.bin?ex={int(signed_until):x}") + + entry = await api._resolve_attachment(channel, 7) + + assert entry.expires_at < signed_until From e86fb5edf428fa75360bc0e3e2f35ddfb8c633b9 Mon Sep 17 00:00:00 2001 From: VulcanoSoftware Date: Wed, 9 Sep 2026 10:13:47 +0000 Subject: [PATCH 2/2] feat(tools): add download throughput benchmark to locate the bottleneck Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tools/bench_download.py | 148 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 148 insertions(+) create mode 100644 tools/bench_download.py diff --git a/tools/bench_download.py b/tools/bench_download.py new file mode 100644 index 0000000..1de2424 --- /dev/null +++ b/tools/bench_download.py @@ -0,0 +1,148 @@ +"""Measure where DCFS download throughput is actually lost. + +Runs against a live DCFS configuration (same config file the server uses) and +reports throughput for each layer of the download path, so a slow transfer can +be attributed to the Discord CDN, to DCFS's own plumbing, or to the protocol +server on top of it. + +Usage: + + DCFS_CONFIG_FILE=config.yaml python -m tools.bench_download /client/path/file.bin + +Reported measurements: + +* ``dcfs full download`` -- everything DCFS does for an SFTP/WebDAV read: + metadata, decryption, part ordering, parallel + CDN range requests. +* ``cdn single connection`` -- one plain HTTP GET of one 8 MB part. This is the + per-connection speed Discord gives this host. +* ``cdn N connections`` -- the same part fetched as N parallel byte ranges. + If this is not ~N times the single-connection + number, the link (or Discord) is the ceiling and + no amount of client-side parallelism will help. +""" + +import asyncio +import logging +import sys +import time +from typing import Optional + +import aiohttp + +from dcfs.config import get_config +from dcfs.core import Clients, Ops +from dcfs.utils.others import is_big_file + +logger = logging.getLogger(__name__) + +PARALLEL_PROBES = (1, 4, 8, 16) + + +def _mbps(nbytes: int, seconds: float) -> float: + return nbytes / seconds / 1e6 if seconds > 0 else 0.0 + + +async def _time_stream(stream) -> tuple[int, float]: + t0 = time.monotonic() + total = 0 + async for chunk in stream: + total += len(chunk) + return total, time.monotonic() - t0 + + +async def _bench_dcfs(ops: Ops, path: str, limit: int) -> None: + stream = await ops.download(path, 0, limit - 1, "bench", validate=False) + total, dt = await _time_stream(stream) + logger.info( + f"dcfs full download {_mbps(total, dt):8.1f} MB/s " + f"({total / 1e6:.0f} MB in {dt:.1f}s)" + ) + + +async def _bench_cdn(url: str, size: int) -> None: + async with aiohttp.ClientSession() as session: + + async def fetch(begin: int, end: int) -> int: + headers = {"Range": f"bytes={begin}-{end}"} + async with session.get(url, headers=headers) as resp: + resp.raise_for_status() + got = 0 + async for chunk in resp.content.iter_chunked(1024 * 1024): + got += len(chunk) + return got + + for n in PARALLEL_PROBES: + per = size // n + ranges = [(i * per, (i + 1) * per - 1) for i in range(n)] + t0 = time.monotonic() + totals = await asyncio.gather(*[fetch(b, e) for b, e in ranges]) + dt = time.monotonic() - t0 + label = "cdn single connection" if n == 1 else f"cdn {n} connections" + logger.info( + f"{label:23s} {_mbps(sum(totals), dt):8.1f} MB/s " + f"({sum(totals) / 1e6:.0f} MB in {dt:.1f}s)" + ) + + +async def _first_part_url(ops: Ops, path: str) -> Optional[tuple[str, int]]: + """Resolve the CDN URL and size of the file's first Discord part.""" + fd = await ops.desc(path, validate=False) + fv = fd.get_latest_version() + if not fv.message_ids: + return None + repo = ops._client.fc_repo + # Unwrap the encryption decorator if it is installed. + repo = getattr(repo, "_inner", repo) + message_api = repo._message_api # type: ignore[union-attr] + bot = message_api.discord_api.next_bot + channel = await bot._get_channel( # type: ignore[attr-defined] + bot._parse_channel_id(message_api.private_file_channel) # type: ignore[attr-defined] + ) + attachment = await bot._resolve_attachment( # type: ignore[attr-defined] + channel, fv.message_ids[0] + ) + return attachment.url, attachment.size + + +async def main(argv: list[str]) -> int: + if len(argv) < 2: + logger.info(__doc__) + return 2 + + path = argv[1] + limit_mb = int(argv[2]) if len(argv) > 2 else 64 + + # Imported lazily so the module can be read without a Discord login. + from main import create_clients + + config = get_config() + clients: Clients = await create_clients(config) + + client_name, _, sub_path = path.lstrip("/").partition("/") + if client_name not in clients: + logger.info(f"unknown client '{client_name}'; known: {', '.join(clients)}") + return 2 + + ops = Ops(clients[client_name]) + sub_path = "/" + sub_path + + resolved = await _first_part_url(ops, sub_path) + if resolved: + logger.info( + f"part size {resolved[1] / 1e6:.1f} MB, " + f"parallel range split enabled: {is_big_file(resolved[1])}" + ) + else: + logger.info("no parts found") + + await _bench_dcfs(ops, sub_path, limit_mb * 1024 * 1024) + if resolved: + await _bench_cdn(*resolved) + return 0 + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO, format="%(message)s") + logging.getLogger("discord").setLevel(logging.WARNING) + sys.exit(asyncio.run(main(sys.argv)))