Optimize SFTP download performance with background prefetching and range parallelism

.gitignore

@@ -6,6 +6,9 @@ CLAUDE.md
 dist/
 coverage.xml
 .coverage
+.mypy_cache/
+.pytest_cache/
+.venv/
 dcfs-gh-pages/.next/
 dcfs-gh-pages/out/
 dcfs-gh-pages/node_modules/

dcfs/app/sftp/handler.py

@@ -332,12 +332,50 @@ def __init__(self, ops: Ops, path: str, mode: str, client_name: str):
 
         # Read streaming state
         self._read_stream: Optional[AsyncIterator[bytes]] = None
-        self._read_iter: Optional[AsyncIterator[bytes]] = None
         self._buf_offset = 0
         self._read_buf = bytearray()
         self._read_lock = asyncio.Lock()
         self._cached_attrs: Optional[asyncssh.SFTPAttrs] = None
 
+        # Prefetch state
+        self._prefetch_queue: Optional[asyncio.Queue[Optional[Any]]] = None
+        self._prefetch_task: Optional[asyncio.Task[None]] = None
+        self._prefetch_eof = False
+
+    async def _run_prefetch(
+        self,
+        stream: AsyncIterator[bytes],
+        queue: asyncio.Queue[Optional[Any]],
+    ) -> None:
+        try:
+            async for chunk in stream:
+                if chunk:
+                    await queue.put(chunk)
+            await queue.put(None)
+        except asyncio.CancelledError:
+            raise
+        except Exception as ex:
+            await queue.put(ex)
+
+    async def _stop_prefetch(self) -> None:
+        if self._prefetch_task is not None and not self._prefetch_task.done():
+            self._prefetch_task.cancel()
+            try:
+                await self._prefetch_task
+            except (asyncio.CancelledError, Exception) as ex:
+                logger.debug(f"Prefetch task stopped: {ex}")
+            self._prefetch_task = None
+
+        if self._read_stream is not None:
+            try:
+                await cast(AsyncGenerator[bytes, None], self._read_stream).aclose()
+            except Exception as ex:
+                logger.debug(f"Error closing read stream: {ex}")
+            self._read_stream = None
+
+        self._prefetch_queue = None
+        self._prefetch_eof = False
+
     async def read(self, offset: int, size: int) -> bytes:
         if "r" not in self.mode:
             raise asyncssh.SFTPPermissionDenied("File not open for reading")
@@ -352,10 +390,7 @@ async def read(self, offset: int, size: int) -> bytes:
             )
 
             if not in_buffer:
-                if self._read_stream is not None:
-                    await cast(AsyncGenerator[bytes, None], self._read_stream).aclose()
-                    self._read_stream = None
-                    self._read_iter = None
+                await self._stop_prefetch()
 
                 self._read_buf = bytearray()
                 self._buf_offset = offset
@@ -366,20 +401,28 @@ async def read(self, offset: int, size: int) -> bytes:
                     os.path.basename(self.path),
                     validate=False,
                 )
-                self._read_iter = self._read_stream.__aiter__()
+                self._prefetch_queue = asyncio.Queue(maxsize=64)
+                self._prefetch_eof = False
+                self._prefetch_task = asyncio.create_task(
+                    self._run_prefetch(self._read_stream, self._prefetch_queue)
+                )
             else:
                 discard = offset - self._buf_offset
                 if discard > 0:
                     self._read_buf = self._read_buf[discard:]
                     self._buf_offset = offset
 
-            it = cast(AsyncIterator[bytes], self._read_iter)
-            while len(self._read_buf) < size:
-                try:
-                    chunk = await anext(it)
-                    self._read_buf.extend(chunk)
-                except StopAsyncIteration:
+            while len(self._read_buf) < size 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)
 
             data = self._read_buf[:size]
             self._read_buf = self._read_buf[size:]
@@ -410,10 +453,7 @@ async def close(self) -> None:
 
         self.closed = True
 
-        if self._read_stream is not None:
-            await cast(AsyncGenerator[bytes, None], self._read_stream).aclose()
-            self._read_stream = None
-            self._read_iter = None
+        await self._stop_prefetch()
 
         if "w" in self.mode and self.buffer:
             data = bytes(self.buffer)

dcfs/config.py

@@ -15,7 +15,7 @@
 @dataclass
 class DownloadConfig:
     chunk_size_kb: int
-    download_max_concurrent_parts: int = 3
+    download_max_concurrent_parts: int = 6
     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", 3)
+                data.get("download_max_concurrent_parts", 6)
             ),
             upload_max_retries=int(data.get("upload_max_retries", 10)),
             upload_retry_interval=int(

dcfs/core/client.py

@@ -36,7 +36,7 @@ async def create(
         metadata_cfg: MetadataConfig,
         discord_api: DiscordApi,
         encryption_cfg: Optional[EncryptionConfig] = None,
-        download_max_concurrent_parts: int = 3,
+        download_max_concurrent_parts: int = 6,
     ) -> "Client":
         channel = await discord_api.next_bot.resolve_channel_id(channel_id)
         message_api = MessageApi(discord_api, channel)

dcfs/utils/others.py

@@ -8,4 +8,4 @@ def exclude_none(iterable: Iterable[Optional[T]]) -> Iterable[T]:
 
 
 def is_big_file(size: int) -> bool:
-    return size > 10 * 1024 * 1024  # 10 MB
+    return size >= 2 * 1024 * 1024  # 2 MB

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 == 3  # default
+        assert config.download_max_concurrent_parts == 6  # default
 
     def test_from_dict_custom_concurrent(self):
         data = {"chunk_size_kb": 1024, "download_max_concurrent_parts": 5}

tests/dcfs/core/api/message/test_overflow.py

@@ -1,7 +1,10 @@
-import pytest
 from unittest.mock import AsyncMock, MagicMock
-from dcfs.core.api.message import MessageApi, OVERFLOW_SENTINEL, OVERFLOW_FILENAME
-from dcfs.reqres import SendMessageResp, MessageResp, Document
+
+import pytest
+
+from dcfs.core.api.message import OVERFLOW_FILENAME, OVERFLOW_SENTINEL, MessageApi
+from dcfs.reqres import Document, MessageResp, SendMessageResp
+
 
 @pytest.mark.asyncio
 async def test_send_text_overflow(mocker):

tests/dcfs/utils/test_others.py

@@ -22,11 +22,11 @@ def test_exclude_none_with_empty_list(self):
 
     def test_is_big_file_small_file(self):
         assert is_big_file(1024) is False
-        assert is_big_file(10 * 1024 * 1024 - 1) is False
+        assert is_big_file(2 * 1024 * 1024 - 1) is False
 
     def test_is_big_file_big_file(self):
-        assert is_big_file(10 * 1024 * 1024) is False
-        assert is_big_file(10 * 1024 * 1024 + 1) is True
+        assert is_big_file(2 * 1024 * 1024) is True
+        assert is_big_file(10 * 1024 * 1024) is True
         assert is_big_file(50 * 1024 * 1024) is True
 
     def test_is_big_file_edge_case(self):

tests/test_asgidav/test_app.py

@@ -1,4 +1,5 @@
 import pytest
+
 from asgidav.app import extract_path_from_destination, split_path
 
 
@@ -48,7 +49,9 @@ class TestAppEndpoints:
     @pytest.mark.asyncio
     async def test_proppatch_endpoint_found(self, mocker):
         from fastapi.testclient import TestClient
+
         from asgidav.app import create_app
+
         from .common import MockResource
 
         mock_get_member = mocker.AsyncMock(return_value=MockResource("/test.txt"))
@@ -62,6 +65,7 @@ async def test_proppatch_endpoint_found(self, mocker):
     @pytest.mark.asyncio
     async def test_proppatch_endpoint_not_found(self, mocker):
         from fastapi.testclient import TestClient
+
         from asgidav.app import create_app
 
         mock_get_member = mocker.AsyncMock(return_value=None)

tests/test_asgidav/test_reqres.py

@@ -1,7 +1,13 @@
 import pytest
 from fastapi import Request
 
-from asgidav.reqres import PropfindRequest, _propfind_response, _propstat, propfind, proppatch
+from asgidav.reqres import (
+    PropfindRequest,
+    _propfind_response,
+    _propstat,
+    propfind,
+    proppatch,
+)
 
 from .common import MockFolder, MockResource
 

tests/test_sftp_factory.py

@@ -1,8 +1,11 @@
-import pytest
 from unittest.mock import MagicMock
+
+import pytest
+
 from dcfs.app.sftp import create_sftp_server
 from dcfs.app.sftp.handler import DCFSSFTPHandler
 
+
 @pytest.mark.asyncio
 async def test_sftp_factory_logic(mocker):
     # Mock dependencies

tests/test_sftp_handler.py

@@ -1,7 +1,7 @@
-import pytest
 from unittest.mock import AsyncMock, MagicMock
 
 import asyncssh
+import pytest
 
 from dcfs.app.sftp.handler import DCFSSFTPBufferedFile
 
@@ -88,3 +88,21 @@ async def test_sftp_buffered_file_eof_and_mode_checks():
         await file_handle.write(0, b"data")
 
     await file_handle.close()
+
+
+@pytest.mark.asyncio
+async def test_sftp_buffered_file_prefetch_error():
+    async def mock_error_gen():
+        yield b"chunk1"
+        raise ValueError("Download failed")
+
+    mock_ops = MagicMock()
+    mock_ops.download = AsyncMock(return_value=mock_error_gen())
+
+    file_handle = DCFSSFTPBufferedFile(mock_ops, "/test.txt", "r", "client1")
+
+    # Reading raises the exception propagated from the prefetch worker
+    with pytest.raises(ValueError, match="Download failed"):
+        await file_handle.read(0, 100)
+
+    await file_handle.close()