From b6908326594df570f9cd8efb116e4fc20e536cff Mon Sep 17 00:00:00 2001 From: splitsec2 <35583321+splitsec2@users.noreply.github.com> Date: Fri, 25 Sep 2026 16:06:15 -0600 Subject: [PATCH] fix(download): stream a completed book instead of buffering it in RAM - lowering memory needs significantly (#1378) Credit where it is due: this defect was found and measured by **@DrNgo** in [DrNgo/shelfmark-fork@ae2185c](https://github.com/DrNgo/shelfmark-fork/commit/ae2185c8a83d537e55f1dc03f745ab844b5cdcdb). He recorded a 493 MB audiobook peaking near 963 MB and OOM-killing a 1 GiB container, with the proxy access log showing a single `GET /api/localdownload` returning 502 at the exact second of the kill. The analysis is his; I am sending it because it is still open here. Serving a completed book from the live queue calls `get_book_data`, which reads the whole file into bytes so the route can wrap it in a `BytesIO` for `send_file`. That is two copies of the book, with one resident for the length of the client transfer, to hand over a file that is already sitting on disk. `get_book_path` returns the path the task already holds and `send_file` streams it. The history fallback a few lines up in the same route has always worked this way, so this makes the two paths consistent rather than introducing anything new. `get_book_data` stays for callers that genuinely want the bytes, now documented as the expensive option. ## The fetch side has the same problem, and this PR does not fix it I should note: `download_url` in `shelfmark/download/http.py` still builds the inbound file in a `BytesIO` and returns it, so a large download is fully resident while it fetches. @DrNgo's commit fixes that too, with a `tempfile.SpooledTemporaryFile(max_size=...)` so small payloads stay in memory exactly as they do now and large ones spill to disk. I left it out deliberately. It changes the return type of `download_url` from `BytesIO` to a file object, which touches several callers, and it lands in a file that has been reworked around bypass handling, waiting rooms and resume since his branch point. That deserves its own PR rather than riding along with a two-function change. I may send it later; if @DrNgo sends it first, his should win, and if you would rather have both together say so and I will hold this one. ## Verification - `tests/download/test_orchestrator_retry.py`: `get_book_path` returns the path without opening the file (the test fails the run if it does), and reports a missing file rather than handing back a dead path. - The existing `/api/localdownload` tests still pass unchanged, including the history fallback and the ownership checks. - Full suite (3222), ruff, ruff format, basedpyright, vulture green. --- shelfmark/download/orchestrator.py | 32 ++++++++++++++- shelfmark/main.py | 11 +++--- tests/download/test_orchestrator_retry.py | 47 +++++++++++++++++++++++ 3 files changed, 83 insertions(+), 7 deletions(-) diff --git a/shelfmark/download/orchestrator.py b/shelfmark/download/orchestrator.py index dc1e55cd..841cdcaa 100644 --- a/shelfmark/download/orchestrator.py +++ b/shelfmark/download/orchestrator.py @@ -367,8 +367,38 @@ def queue_status(user_id: int | None = None) -> dict[str, dict[str, Any]]: } +def get_book_path(task_id: str) -> tuple[str | None, DownloadTask | None]: + """Path of a task's downloaded file, without reading it into memory. + + Serving a completed book used to go through :func:`get_book_data`, which read the + whole file so the route could wrap it in a BytesIO: two copies resident for a + transfer that the web server can stream straight off disk. On a large audiobook + that is enough to reach a container's memory limit. + """ + task = None + try: + task = book_queue.get_task(task_id) + if not task: + return None, None + + path = task.download_path + if not path or not Path(path).is_file(): + return None, task + except OSError as e: + logger.error_trace(f"Error resolving book path: {e}") + if task: + task.download_path = None + return None, task + else: + return path, task + + def get_book_data(task_id: str) -> tuple[bytes | None, DownloadTask | None]: - """Get downloaded file data for a specific task.""" + """Get downloaded file data for a specific task. + + Prefer :func:`get_book_path` when the bytes are only going to be written straight + back out; this reads the entire file into memory. + """ task = None try: task = book_queue.get_task(task_id) diff --git a/shelfmark/main.py b/shelfmark/main.py index 4c6b640a..e1ad0960 100644 --- a/shelfmark/main.py +++ b/shelfmark/main.py @@ -1,7 +1,6 @@ """Flask app - routes, WebSocket handlers, and middleware.""" import binascii -import io import logging import os import re @@ -1698,8 +1697,10 @@ def api_local_download() -> Response | tuple[Response, int]: return jsonify({"error": "No book ID provided"}), 400 try: - file_data, book_info = backend.get_book_data(book_id) - if file_data is None: + # The path, not the bytes: send_file streams it off disk, where reading it + # first cost two resident copies of the whole book. + file_path, book_info = backend.get_book_path(book_id) + if file_path is None: # Fallback for dismissed/history entries where queue task may no longer exist. if download_history_service is not None: is_admin, db_user_id, can_access_status = _resolve_status_scope() @@ -1732,9 +1733,7 @@ def api_local_download() -> Response | tuple[Response, int]: return jsonify({"error": "File not found"}), 404 file_name = book_info.get_filename() if book_info is not None else Path(book_id).name - # Prepare the file for sending to the client - data = io.BytesIO(file_data) - return send_file(data, download_name=file_name, as_attachment=True) + return send_file(file_path, download_name=file_name, as_attachment=True) except _OPERATIONAL_ERRORS as e: logger.error_trace(f"Local download error: {e}") diff --git a/tests/download/test_orchestrator_retry.py b/tests/download/test_orchestrator_retry.py index fc822f98..d938c3bb 100644 --- a/tests/download/test_orchestrator_retry.py +++ b/tests/download/test_orchestrator_retry.py @@ -285,3 +285,50 @@ def test_get_book_data_clears_download_path_when_file_read_fails(monkeypatch, tm assert file_data is None assert returned_task is task assert task.download_path is None + + +def test_get_book_path_returns_the_path_without_reading_the_file(monkeypatch, tmp_path): + import shelfmark.download.orchestrator as orchestrator + + book = tmp_path / "book.epub" + book.write_bytes(b"x" * 4096) + task = DownloadTask( + task_id="task-book-path-1", + source="direct_download", + title="Streamed Book", + download_path=str(book), + ) + + mock_queue = MagicMock() + mock_queue.get_task.return_value = task + monkeypatch.setattr(orchestrator, "book_queue", mock_queue) + + def _no_reads(*_args, **_kwargs): + raise AssertionError("get_book_path must not open the file") + + monkeypatch.setattr(Path, "open", _no_reads) + + path, returned_task = orchestrator.get_book_path(task.task_id) + + assert path == str(book) + assert returned_task is task + + +def test_get_book_path_reports_a_missing_file_rather_than_a_dead_path(monkeypatch, tmp_path): + import shelfmark.download.orchestrator as orchestrator + + task = DownloadTask( + task_id="task-book-path-2", + source="direct_download", + title="Gone", + download_path=str(tmp_path / "gone.epub"), + ) + + mock_queue = MagicMock() + mock_queue.get_task.return_value = task + monkeypatch.setattr(orchestrator, "book_queue", mock_queue) + + path, returned_task = orchestrator.get_book_path(task.task_id) + + assert path is None + assert returned_task is task