diff --git a/.gitignore b/.gitignore index 115b5b3..861a99b 100644 --- a/.gitignore +++ b/.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/ diff --git a/dcfs/app/sftp/handler.py b/dcfs/app/sftp/handler.py index 52da962..5f6b64b 100644 --- a/dcfs/app/sftp/handler.py +++ b/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) diff --git a/dcfs/config.py b/dcfs/config.py index ce55146..e1de21f 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 = 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( diff --git a/dcfs/core/client.py b/dcfs/core/client.py index cafe240..5f12df8 100644 --- a/dcfs/core/client.py +++ b/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) diff --git a/dcfs/utils/others.py b/dcfs/utils/others.py index 7d43dfd..47a626a 100644 --- a/dcfs/utils/others.py +++ b/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 diff --git a/tests/dcfs/config/test_config.py b/tests/dcfs/config/test_config.py index f25671e..2aea00f 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 == 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} diff --git a/tests/dcfs/core/api/message/test_overflow.py b/tests/dcfs/core/api/message/test_overflow.py index d5d840a..034eed9 100644 --- a/tests/dcfs/core/api/message/test_overflow.py +++ b/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): diff --git a/tests/dcfs/utils/test_others.py b/tests/dcfs/utils/test_others.py index 18794ca..eae6e6a 100644 --- a/tests/dcfs/utils/test_others.py +++ b/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): diff --git a/tests/test_asgidav/test_app.py b/tests/test_asgidav/test_app.py index eff98ad..fd83ba0 100644 --- a/tests/test_asgidav/test_app.py +++ b/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) diff --git a/tests/test_asgidav/test_reqres.py b/tests/test_asgidav/test_reqres.py index a60d64d..18b7a57 100644 --- a/tests/test_asgidav/test_reqres.py +++ b/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 diff --git a/tests/test_sftp_factory.py b/tests/test_sftp_factory.py index ec5125a..507fff7 100644 --- a/tests/test_sftp_factory.py +++ b/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 diff --git a/tests/test_sftp_handler.py b/tests/test_sftp_handler.py index 9667416..dbed880 100644 --- a/tests/test_sftp_handler.py +++ b/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()