diff --git a/dcfs/app/sftp/__init__.py b/dcfs/app/sftp/__init__.py index 1fb3b20..a25aa05 100644 --- a/dcfs/app/sftp/__init__.py +++ b/dcfs/app/sftp/__init__.py @@ -58,7 +58,7 @@ 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, ) 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..60a8077 100644 --- a/dcfs/app/sftp/handler.py +++ b/dcfs/app/sftp/handler.py @@ -322,7 +322,7 @@ async def fstat(self) -> asyncssh.SFTPAttrs: class DCFSSFTPBufferedFile(DCFSSFTPFileBase): - MAX_FORWARD_SKIP = 2 * 1024 * 1024 # 2 MB forward skip + MAX_FORWARD_SKIP = 8 * 1024 * 1024 # 8 MB forward skip MAX_BACKWARD_RETAIN = 4 * 1024 * 1024 # 4 MB backward retain def __init__(self, ops: Ops, path: str, mode: str, client_name: str): @@ -391,7 +391,7 @@ async def _start_prefetch(self, offset: int) -> None: os.path.basename(self.path), validate=False, ) - self._prefetch_queue = asyncio.Queue(maxsize=64) + self._prefetch_queue = asyncio.Queue(maxsize=128) self._prefetch_eof = False self._prefetch_task = asyncio.create_task( self._run_prefetch(self._read_stream, self._prefetch_queue) @@ -401,7 +401,17 @@ async def read(self, offset: int, size: int) -> bytes: if "r" not in self.mode: raise asyncssh.SFTPPermissionDenied("File not open for reading") + # Fast lock-free path for sequential/pipelined buffer hits + rel_offset = offset - self._buf_offset + if 0 <= rel_offset and rel_offset + size <= len(self._read_buf): + return bytes(self._read_buf[rel_offset : rel_offset + size]) + async with self._read_lock: + # Re-check inside lock after acquiring + rel_offset = offset - self._buf_offset + if 0 <= rel_offset and rel_offset + size <= len(self._read_buf): + return bytes(self._read_buf[rel_offset : rel_offset + size]) + buf_end = self._buf_offset + len(self._read_buf) can_reuse_stream = ( @@ -450,12 +460,12 @@ async def read(self, offset: int, size: int) -> bytes: self._highest_offset = max(self._highest_offset, offset + len(data)) - # Prune buffer behind prune_target to keep memory bounded + # Prune buffer behind prune_target to keep memory bounded without reallocation 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:] + del self._read_buf[:discard] self._buf_offset += discard return data diff --git a/dcfs/config.py b/dcfs/config.py index e1de21f..ee2006a 100644 --- a/dcfs/config.py +++ b/dcfs/config.py @@ -15,7 +15,7 @@ @dataclass class DownloadConfig: chunk_size_kb: int - download_max_concurrent_parts: int = 6 + download_max_concurrent_parts: int = 12 upload_max_retries: int = 10 upload_retry_interval: int = 5 upload_base_retry_delay: float = 2.0 @@ -25,7 +25,7 @@ def from_dict(cls, data: dict) -> Self: return cls( chunk_size_kb=data["chunk_size_kb"], download_max_concurrent_parts=int( - data.get("download_max_concurrent_parts", 6) + data.get("download_max_concurrent_parts", 12) ), upload_max_retries=int(data.get("upload_max_retries", 10)), upload_retry_interval=int( diff --git a/dcfs/core/api/message/__init__.py b/dcfs/core/api/message/__init__.py index da59b5f..620744c 100644 --- a/dcfs/core/api/message/__init__.py +++ b/dcfs/core/api/message/__init__.py @@ -1,6 +1,6 @@ import asyncio import logging -from typing import AsyncIterator, Iterable, Iterator, List +from typing import Iterable, Iterator, List, cast from pyrate_limiter import Duration, InMemoryBucket, Limiter, Rate @@ -178,44 +178,41 @@ def _size(begin: int, end: int) -> int: 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 = 16 sub_ranges = list(self.split_download_tasks(begin, end, n)) - - resps = await asyncio.gather(*[ - self.discord_api.next_bot.download_file( - DownloadFileReq( - chat=self.private_file_channel, - message_id=message_id, - chunk_size=get_config().dcfs.download.chunk_size_kb, - begin=b, - end=e, - ) - ) - for b, e in sub_ranges - ]) + chunk_size_kb = get_config().dcfs.download.chunk_size_kb queues: list[asyncio.Queue[object]] = [ - asyncio.Queue(maxsize=32) for _ in range(n) + asyncio.Queue(maxsize=32) for _ in sub_ranges ] async def _producer( - chunks_iter: Iterator[bytes] | AsyncIterator[bytes], - q: asyncio.Queue[object], + b: int, e: int, q: asyncio.Queue[object] ) -> None: try: + resp = await self.discord_api.next_bot.download_file( + DownloadFileReq( + chat=self.private_file_channel, + message_id=message_id, + chunk_size=chunk_size_kb, + begin=b, + end=e, + ) + ) + chunks_iter = resp.chunks if hasattr(chunks_iter, "__anext__"): async for chunk in chunks_iter: # type: ignore[union-attr] await q.put(chunk) else: - for chunk in chunks_iter: # type: ignore[union-attr] + for chunk in cast(Iterator[bytes], chunks_iter): await q.put(chunk) await q.put(None) except Exception as ex: await q.put(ex) producer_tasks = [ - asyncio.create_task(_producer(resp.chunks, q)) - for resp, q in zip(resps, queues) + asyncio.create_task(_producer(b, e, q)) + for (b, e), q in zip(sub_ranges, queues) ] async def _parallel_chunks(): diff --git a/dcfs/core/repository/impl/file_content/__init__.py b/dcfs/core/repository/impl/file_content/__init__.py index 88a3286..f926923 100644 --- a/dcfs/core/repository/impl/file_content/__init__.py +++ b/dcfs/core/repository/impl/file_content/__init__.py @@ -233,7 +233,7 @@ async def _stream_parts( ) queue: asyncio.Queue["tuple[int, Optional[bytes]]"] = asyncio.Queue( - maxsize=64 + maxsize=128 ) # Shared dict: producer stores its exception HERE before the diff --git a/dcfs/crypto/repository.py b/dcfs/crypto/repository.py index 430d648..10f2e58 100644 --- a/dcfs/crypto/repository.py +++ b/dcfs/crypto/repository.py @@ -252,11 +252,15 @@ async def content_length(self, fv: "DCFSFileVersion") -> int: return 0 # ``_detect`` already caches per file, so the second call from a HEAD # request or a Content-Range computation is free. - detected = await self._detect(fv, "") - if detected is None: + try: + detected = await self._detect(fv, "") + if detected is None: + return fv.size + header, _ = detected + return _plaintext_size_from_ciphertext(fv.size, header.chunk_size) + except InvalidHeaderError as ex: + logger.warning(f"Could not verify encryption header for version {fv.id}: {ex}") return fv.size - header, _ = detected - return _plaintext_size_from_ciphertext(fv.size, header.chunk_size) # -- internals --------------------------------------------------------- @@ -428,7 +432,8 @@ def _trim_and_count(plaintext: bytes) -> bytes: # Yield to the event loop between decrypting chunks so the # Discord gateway heartbeat and other async tasks can make # progress during large downloads. - await asyncio.sleep(0) + if chunks_done % 16 == 0: + await asyncio.sleep(0) # Stream exhausted. Anything left in ``buf`` is the final requested # chunk -- its on-wire size is ``len(buf)``, which may equal ``stride`` diff --git a/dcfs/discord/impl/discord_bot.py b/dcfs/discord/impl/discord_bot.py index e3bdb3b..321b4c4 100644 --- a/dcfs/discord/impl/discord_bot.py +++ b/dcfs/discord/impl/discord_bot.py @@ -32,7 +32,7 @@ logger = logging.getLogger(__name__) -CHUNK_SIZE = 1024 * 1024 # 1 MB chunks for downloads +CHUNK_SIZE = 1024 * 1024 # 1 MB chunks for efficient streaming class DiscordBotAPI(IDiscordClient): @@ -41,10 +41,19 @@ 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._url_cache: dict[int, tuple[str, int]] = {} + self._inflight_fetches: dict[int, asyncio.Future[tuple[str, int]]] = {} + self._url_cache_lock = asyncio.Lock() + self._fetch_message_semaphore = asyncio.Semaphore(3) 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, + ttl_dns_cache=300, + enable_cleanup_closed=True, + ) + self._http_session = aiohttp.ClientSession(connector=connector) return self._http_session async def _get_channel(self, channel_id: int) -> Any: @@ -77,7 +86,14 @@ async def get_messages(self, req: GetMessagesReq) -> GetMessagesResp: async def _fetch(m_id: int) -> Optional[MessageResp]: try: - msg = await channel.fetch_message(m_id) + async with self._fetch_message_semaphore: + msg = await channel.fetch_message(m_id) + if msg.attachments: + att = msg.attachments[0] + async with self._url_cache_lock: + if len(self._url_cache) > 2048: + self._url_cache.clear() + self._url_cache[m_id] = (att.url, att.size) return self._to_message_dto(msg) except discord.NotFound: return None @@ -153,21 +169,55 @@ 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] + + async with self._url_cache_lock: + cached = self._url_cache.get(req.message_id) + if cached is not None: + url, attach_size = cached + is_owner = False + fut = None + elif req.message_id in self._inflight_fetches: + is_owner = False + fut = self._inflight_fetches[req.message_id] + else: + is_owner = True + fut = asyncio.get_running_loop().create_future() + self._inflight_fetches[req.message_id] = fut + + if cached is None and fut is not None: + if is_owner: + try: + channel = await self._get_channel(channel_id) + try: + async with self._fetch_message_semaphore: + 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] + res = (attachment.url, attachment.size) + async with self._url_cache_lock: + if len(self._url_cache) > 2048: + self._url_cache.clear() + self._url_cache[req.message_id] = res + self._inflight_fetches.pop(req.message_id, None) + fut.set_result(res) + except Exception as ex: + async with self._url_cache_lock: + self._inflight_fetches.pop(req.message_id, None) + fut.set_exception(ex) + raise ex + else: + res = await fut + + url, attach_size = res session = await self._ensure_http_session() # Build optional Range header so the CDN only streams the requested # byte range (critical for download_file_parallel sub-requests). should_range = req.begin > 0 or req.end != -1 - url = attachment.url headers = {} if should_range: range_end = "" if req.end == -1 else str(req.end) @@ -175,7 +225,7 @@ async def download_file(self, req: DownloadFileReq) -> DownloadFileResp: logger.info( "CDN download: msg=%d range=%d-%d should_range=%s attach_size=%d", - req.message_id, req.begin, req.end, should_range, attachment.size, + req.message_id, req.begin, req.end, should_range, attach_size, ) # Timeout: connect within 15s, download within 120s. Without a @@ -244,7 +294,7 @@ async def _chunk_generator(): finally: response.close() - return DownloadFileResp(chunks=_chunk_generator(), size=attachment.size) + return DownloadFileResp(chunks=_chunk_generator(), size=attach_size) async def search_messages(self, req: SearchMessageReq) -> GetMessagesRespNoNone: channel_id = self._parse_channel_id(req.chat) diff --git a/tests/dcfs/config/test_config.py b/tests/dcfs/config/test_config.py index 2aea00f..1c15dea 100644 --- a/tests/dcfs/config/test_config.py +++ b/tests/dcfs/config/test_config.py @@ -21,7 +21,7 @@ def test_from_dict(self): config = DownloadConfig.from_dict(data) assert config.chunk_size_kb == 1024 - assert config.download_max_concurrent_parts == 6 # default + assert config.download_max_concurrent_parts == 12 # default def test_from_dict_custom_concurrent(self): data = {"chunk_size_kb": 1024, "download_max_concurrent_parts": 5} diff --git a/tests/dcfs/core/api/message/test_parallel.py b/tests/dcfs/core/api/message/test_parallel.py index d4aa504..09c1843 100644 --- a/tests/dcfs/core/api/message/test_parallel.py +++ b/tests/dcfs/core/api/message/test_parallel.py @@ -24,9 +24,21 @@ async def test_download_file_parallel(): DownloadFileResp(chunks=mock_chunks([b"part2_chunk1"]), size=10), DownloadFileResp(chunks=mock_chunks([b"part3_chunk1"]), size=10), DownloadFileResp(chunks=mock_chunks([b"part4_chunk1"]), size=10), + DownloadFileResp(chunks=mock_chunks([b"part5_chunk1"]), size=10), + DownloadFileResp(chunks=mock_chunks([b"part6_chunk1"]), size=10), + DownloadFileResp(chunks=mock_chunks([b"part7_chunk1"]), size=10), + DownloadFileResp(chunks=mock_chunks([b"part8_chunk1"]), size=10), + DownloadFileResp(chunks=mock_chunks([b"part9_chunk1"]), size=10), + DownloadFileResp(chunks=mock_chunks([b"part10_chunk1"]), size=10), + DownloadFileResp(chunks=mock_chunks([b"part11_chunk1"]), size=10), + DownloadFileResp(chunks=mock_chunks([b"part12_chunk1"]), size=10), + DownloadFileResp(chunks=mock_chunks([b"part13_chunk1"]), size=10), + DownloadFileResp(chunks=mock_chunks([b"part14_chunk1"]), size=10), + DownloadFileResp(chunks=mock_chunks([b"part15_chunk1"]), size=10), + DownloadFileResp(chunks=mock_chunks([b"part16_chunk1"]), size=10), ] - resp = await message_api.download_file_parallel(message_id=999, begin=0, end=39) + resp = await message_api.download_file_parallel(message_id=999, begin=0, end=159) chunks = [] async for chunk in resp.chunks: @@ -38,5 +50,17 @@ async def test_download_file_parallel(): b"part2_chunk1", b"part3_chunk1", b"part4_chunk1", + b"part5_chunk1", + b"part6_chunk1", + b"part7_chunk1", + b"part8_chunk1", + b"part9_chunk1", + b"part10_chunk1", + b"part11_chunk1", + b"part12_chunk1", + b"part13_chunk1", + b"part14_chunk1", + b"part15_chunk1", + b"part16_chunk1", ] - assert bot.download_file.call_count == 4 + assert bot.download_file.call_count == 16