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""