From 347aead445218eccc8886b404a421cdc1b03befd Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Tue, 7 Jul 2026 04:57:35 +0000 Subject: [PATCH] fix: ensure SFTP uploads are streamed directly to Discord This change ensures that SFTP uploads are streamed directly to Discord instead of being fully buffered in memory when the file size is unknown. Key changes: - Supported unknown file sizes (sentinel -1) in `DCMsgFileContentRepository.save` and `FileMessageFromStream.read`. - Updated SFTP handler to always use `DCFSSFTPStreamingFile` for writes, enabling part-by-part streaming to Discord. - Ensured `FileUploader` correctly partitions streams with unknown sizes. - This applies natural backpressure via the async queue, matching the user's actual internet upload speed. Co-authored-by: VulcanoSoftware <113239901+VulcanoSoftware@users.noreply.github.com> --- dcfs/app/sftp/handler.py | 18 ++-- .../repository/impl/file_content/__init__.py | 92 ++++++++++++------- .../impl/file_content/file_uploader.py | 6 +- dcfs/reqres.py | 7 +- 4 files changed, 83 insertions(+), 40 deletions(-) diff --git a/dcfs/app/sftp/handler.py b/dcfs/app/sftp/handler.py index b12a2b1..43bdfac 100644 --- a/dcfs/app/sftp/handler.py +++ b/dcfs/app/sftp/handler.py @@ -188,7 +188,7 @@ async def open(self, path: bytes, flags: int, attrs: asyncssh.SFTPAttrs) -> 'DCF if expected_size is not None and expected_size < 0: expected_size = None - if mode == "w" and expected_size is not None: + if mode == "w": return DCFSSFTPStreamingFile(ops, sub_path, expected_size, client_name) return DCFSSFTPBufferedFile(ops, sub_path, mode, client_name) @@ -428,10 +428,16 @@ async def fstat(self) -> asyncssh.SFTPAttrs: class DCFSSFTPStreamingFile(DCFSSFTPFileBase): - def __init__(self, ops: Ops, path: str, expected_size: int, client_name: str): + def __init__( + self, + ops: Ops, + path: str, + expected_size: Optional[int], + client_name: str, + ): self.ops = ops self.path = path - self.expected_size = expected_size + self.expected_size = expected_size if expected_size is not None else -1 self.client_name = client_name self.closed = False @@ -518,7 +524,7 @@ async def write(self, offset: int, data: bytes) -> int: return 0 end = offset + len(data) - if end > self.expected_size: + if self.expected_size >= 0 and end > self.expected_size: raise asyncssh.SFTPFailure("Write exceeds expected file size") existing = self._pending.get(offset) @@ -542,7 +548,7 @@ async def close(self) -> None: self.closed = True - if self._next_offset != self.expected_size: + if self.expected_size >= 0 and self._next_offset != self.expected_size: raise asyncssh.SFTPFailure( f"Upload incomplete: received {self._next_offset} of {self.expected_size} bytes" ) @@ -555,7 +561,7 @@ async def fstat(self) -> asyncssh.SFTPAttrs: now = int(time.time()) return asyncssh.SFTPAttrs( permissions=stat.S_IFREG | 0o644, - size=self.expected_size, + size=max(0, self.expected_size), mtime=now, atime=now, ) diff --git a/dcfs/core/repository/impl/file_content/__init__.py b/dcfs/core/repository/impl/file_content/__init__.py index 4d26550..88a3286 100644 --- a/dcfs/core/repository/impl/file_content/__init__.py +++ b/dcfs/core/repository/impl/file_content/__init__.py @@ -85,42 +85,70 @@ async def save(self, file_msg: UploadableFileMessage) -> List[SentFileMessage]: tasks = [] max_part_size = discord_max_file_size_bytes() - for i, part_size in enumerate(self._partition(total_size, max_part_size)): - # Acquire semaphore BEFORE buffering to keep memory usage under control. - # Only N parts will be buffered and uploading concurrently. - await self._upload_semaphore.acquire() - - # Update file_msg for the current part - file_msg.name = f"[part{i+1}]{original_name}" - file_msg.size = part_size - - # Create an uploader for this part. Using next_bot to distribute load - # across multiple bots if available. - uploader = FileUploader(self._message_api.discord_api.next_bot, file_msg) - - # Read and buffer the part content sequentially from the source stream. - await uploader.upload() - - # Dispatch the upload as a background task. - tasks.append( - asyncio.create_task( - self._upload_part_task( - uploader, - self._message_api.private_file_channel, - i + 1, - part_size, - self._upload_semaphore, + + if total_size >= 0: + # Case 1: Total size is known upfront. + for i, part_size in enumerate(self._partition(total_size, max_part_size)): + await self._upload_semaphore.acquire() + + file_msg.name = f"[part{i+1}]{original_name}" + file_msg.size = part_size + + uploader = FileUploader(self._message_api.discord_api.next_bot, file_msg) + await uploader.upload() + + tasks.append( + asyncio.create_task( + self._upload_part_task( + uploader, + self._message_api.private_file_channel, + i + 1, + part_size, + self._upload_semaphore, + ) ) ) - ) + await asyncio.sleep(0) + file_msg.next_part(part_size) + else: + # Case 2: Total size is unknown (e.g. streaming from SFTP). + # We read and upload parts until the stream is exhausted. + i = 0 + while True: + await self._upload_semaphore.acquire() + + file_msg.name = f"[part{i+1}]{original_name}" + # Set size to -1 for FileUploader to allow reading up to max_part_size + # without being capped by FileMessageFromStream.read(). + file_msg.size = -1 + + uploader = FileUploader(self._message_api.discord_api.next_bot, file_msg) + part_size = await uploader.upload() + + if part_size == 0 and i > 0: + # EOF reached after some parts were uploaded. + self._upload_semaphore.release() + break + + tasks.append( + asyncio.create_task( + self._upload_part_task( + uploader, + self._message_api.private_file_channel, + i + 1, + part_size, + self._upload_semaphore, + ) + ) + ) + await asyncio.sleep(0) - # Yield to the event loop between parts so the Discord - # gateway heartbeat and other async tasks can make progress - # during large uploads. - await asyncio.sleep(0) + if part_size < max_part_size: + # Last part reached. + break - # Advance the source message to the next part - file_msg.next_part(part_size) + file_msg.next_part(part_size) + i += 1 # Wait for all uploads to complete. return list(await asyncio.gather(*tasks)) diff --git a/dcfs/core/repository/impl/file_content/file_uploader.py b/dcfs/core/repository/impl/file_content/file_uploader.py index 905190e..39eccde 100644 --- a/dcfs/core/repository/impl/file_content/file_uploader.py +++ b/dcfs/core/repository/impl/file_content/file_uploader.py @@ -66,7 +66,11 @@ async def upload(self) -> int: await self._file_msg.open() self._buffer = io.BytesIO() while True: - chunk = await self._file_msg.read(1024 * 1024) + remaining = limit - self._buffer.tell() + if remaining <= 0: + break + + chunk = await self._file_msg.read(min(1024 * 1024, remaining)) if not chunk: break self._buffer.write(chunk) diff --git a/dcfs/reqres.py b/dcfs/reqres.py index c605492..0a08193 100644 --- a/dcfs/reqres.py +++ b/dcfs/reqres.py @@ -275,7 +275,12 @@ def new( ) async def read(self, length: int) -> bytes: - size_to_return = min(length, self.get_size() - self._read_size) + total_size = self.get_size() + if total_size >= 0: + size_to_return = min(length, total_size - self._read_size) + else: + size_to_return = length + if size_to_return <= 0: return b""