Make save_large_file robust to transient errors

CHANGELOG.md

@@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Make `save_large_file` robust against transient HTTP and host/connection errors by streaming with range-based resume and resetting retries on progress (#348)
 - Carry a script's top-level `const`, `let` and `class` declarations across the wombat block, so other scripts on the page can still see them (#329)
 - Keep a rewritten `<script>`'s original `src` in `__wb_orig_src`, so wombat can return it from `getAttribute("src")` and bundlers that identify chunks by it (Next.js/Turbopack, Vite) still work (openzim/warc2zim#473)
 - Keep URL-encoded `#` characters of the path or querystring when rewriting document URIs, instead of dropping everything after them (#341)

src/zimscraperlib/download.py

@@ -1,5 +1,6 @@
 import pathlib
-import subprocess
+import random
+import time
 from concurrent.futures import Future, ThreadPoolExecutor
 from typing import Any, ClassVar
 from urllib.parse import urlparse
@@ -126,33 +127,105 @@ class BestMp4(YoutubeConfig):
     }
 
 
+HTTP_OK = 200
+RETRYABLE_HTTP_STATUSES = {429, 500, 502, 503, 504}
+CHUNK_SIZE = 1024 * 1024
+CONNECT_TIMEOUT_SECONDS = 30
+READ_TIMEOUT_SECONDS = 60
+BACKOFF_BASE_SECONDS = 2
+BACKOFF_MAX_SECONDS = 60
+
+
+class DownloadError(Exception):
+    """Raised when a download fails and cannot be retried any further."""
+
+
 def save_large_file(
-    url: str, fpath: pathlib.Path, retries: int = 5, user_agent: str | None = None
+    url: str,
+    fpath: pathlib.Path,
+    retries: int = 5,
+    user_agent: str | None = None,
 ) -> None:
-    """download a binary file from its URL, using wget
+    """Download a binary file from its URL, streaming the response with resume.
 
-    Arguments -
-        url:
+    Retries on connection errors, timeouts, incomplete reads, and HTTP
+    429/500/502/503/504, resuming from the last byte written each time via a
+    `Range` request. `retries` is the number of consecutive attempts that may
+    fail *without making any progress* before giving up; an attempt that
+    writes at least one byte resets this counter.
     """
-    command = [
-        "/usr/bin/env",
-        "wget",
-        "-t",
-        f"{retries}",
-        "--retry-connrefused",
-        "--random-wait",
-        "--progress=dot:giga",
-        "-O",
-        str(fpath),
-        "-c",
-        url,
-    ]
-    if user_agent:
-        command += ["-U", user_agent]
-    subprocess.run(
-        command,
-        check=True,
-    )
+    headers = {"User-Agent": user_agent} if user_agent else {}
+    consecutive_failures = 0
+
+    with requests.Session() as session:
+        while True:
+            resume_from = fpath.stat().st_size if fpath.exists() else 0
+            request_headers = dict(headers)
+            if resume_from:
+                request_headers["Range"] = f"bytes={resume_from}-"
+
+            bytes_written = 0
+            error: Exception | None = None
+            try:
+                with session.get(
+                    url,
+                    headers=request_headers,
+                    stream=True,
+                    timeout=(CONNECT_TIMEOUT_SECONDS, READ_TIMEOUT_SECONDS),
+                ) as response:
+                    if resume_from and response.status_code == HTTP_OK:
+                        logger.warning(
+                            f"{url} does not support resuming downloads, "
+                            "restarting from scratch"
+                        )
+                        fpath.unlink(missing_ok=True)
+                        resume_from = 0
+
+                    response.raise_for_status()
+
+                    mode = "ab" if resume_from else "wb"
+                    with fpath.open(mode) as fh:
+                        for chunk in response.iter_content(chunk_size=CHUNK_SIZE):
+                            if not chunk:
+                                continue
+                            fh.write(chunk)
+                            bytes_written += len(chunk)
+                return  # response body was fully consumed without error
+            except (
+                requests.ConnectionError,
+                requests.Timeout,
+                requests.exceptions.ChunkedEncodingError,
+            ) as exc:
+                error = exc
+            except requests.HTTPError as exc:
+                status = exc.response.status_code if exc.response is not None else None
+                if status not in RETRYABLE_HTTP_STATUSES:
+                    raise DownloadError(
+                        f"Download of {url} failed with non-retryable error"
+                    ) from exc
+                error = exc
+
+            if bytes_written:
+                consecutive_failures = 0
+            else:
+                consecutive_failures += 1
+
+            if consecutive_failures > retries:
+                raise DownloadError(
+                    f"Download of {url} failed after {retries} consecutive "
+                    "attempts without progress"
+                ) from error
+
+            wait = min(
+                BACKOFF_MAX_SECONDS,
+                BACKOFF_BASE_SECONDS * 2 ** (consecutive_failures - 1),
+            )
+            wait += random.uniform(0, wait * 0.1)  # noqa: S311
+            logger.warning(
+                f"Download of {url} interrupted ({error}), retrying in "
+                f"{wait:.1f}s from byte {resume_from}"
+            )
+            time.sleep(wait)
 
 
 def get_retry_adapter(

tests/download/test_download.py

@@ -2,18 +2,20 @@
 import io
 import pathlib
 import re
+from collections.abc import Iterator
 from typing import ClassVar
-from unittest.mock import Mock
+from unittest.mock import MagicMock, Mock
 
 import pytest
 import requests
 import requests.structures
-from yt_dlp.utils import DownloadError
+from yt_dlp.utils import DownloadError as YtDlpDownloadError
 
 from zimscraperlib.constants import DEFAULT_USER_AGENT, DEFAULT_WEB_REQUESTS_TIMEOUT
 from zimscraperlib.download import (
     BestMp4,
     BestWebm,
+    DownloadError,
     YoutubeDownloader,
     has_user_agent_header,
     is_wikimedia_upload_url,
@@ -311,6 +313,176 @@ def test_large_download_https_custom_ua(tmp_path: pathlib.Path, valid_https_url:
     assert_downloaded_file(valid_https_url, dest_file)
 
 
+class FakeResponse:
+    """Minimal stand-in for requests.Response, usable as a context manager."""
+
+    def __init__(
+        self,
+        status_code: int,
+        chunks: list[bytes] | None = None,
+        raise_while_streaming: Exception | None = None,
+    ) -> None:
+        self.status_code = status_code
+        self._chunks = chunks or []
+        self._raise_while_streaming = raise_while_streaming
+
+    def __enter__(self) -> FakeResponse:
+        return self
+
+    def __exit__(self, *_args: object) -> None:
+        return None
+
+    def raise_for_status(self) -> None:
+        if self.status_code >= 400:
+            response = requests.Response()
+            response.status_code = self.status_code
+            raise requests.HTTPError(response=response)
+
+    def iter_content(self, chunk_size: int) -> Iterator[bytes]:  # noqa: ARG002
+        yield from self._chunks
+        if self._raise_while_streaming:
+            raise self._raise_while_streaming
+
+
+def test_save_large_file_single_attempt(
+    tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """A clean download writes all chunks in one attempt, no Range header."""
+    get_mock = MagicMock(return_value=FakeResponse(200, [b"hello ", b"world"]))
+    monkeypatch.setattr(requests.Session, "get", get_mock)
+
+    fpath = tmp_path / "out.zim"
+    save_large_file("https://example.com/file.zim", fpath, user_agent="some-agent")
+
+    assert fpath.read_bytes() == b"hello world"
+    _args, kwargs = get_mock.call_args
+    assert kwargs["headers"] == {"User-Agent": "some-agent"}
+    assert "Range" not in kwargs["headers"]
+
+
+def test_save_large_file_without_user_agent(
+    tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    get_mock = MagicMock(return_value=FakeResponse(200, [b"data"]))
+    monkeypatch.setattr(requests.Session, "get", get_mock)
+
+    fpath = tmp_path / "out.zim"
+    save_large_file("https://example.com/file.zim", fpath)
+
+    _args, kwargs = get_mock.call_args
+    assert "User-Agent" not in kwargs["headers"]
+
+
+def test_save_large_file_resumes_after_connection_error(
+    tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """A dropped connection resumes with a Range header for the missing bytes."""
+    responses = [
+        FakeResponse(
+            200,
+            [b"hello "],
+            raise_while_streaming=requests.ConnectionError("dropped"),
+        ),
+        FakeResponse(206, [b"world"]),
+    ]
+    get_mock = MagicMock(side_effect=responses)
+    monkeypatch.setattr(requests.Session, "get", get_mock)
+    monkeypatch.setattr("zimscraperlib.download.time.sleep", MagicMock())
+
+    fpath = tmp_path / "out.zim"
+    save_large_file("https://example.com/file.zim", fpath)
+
+    assert fpath.read_bytes() == b"hello world"
+    assert get_mock.call_count == 2
+    second_call_headers = get_mock.call_args_list[1].kwargs["headers"]
+    assert second_call_headers["Range"] == "bytes=6-"
+
+
+def test_save_large_file_retries_on_retryable_http_status(
+    tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    responses = [FakeResponse(503), FakeResponse(200, [b"ok"])]
+    get_mock = MagicMock(side_effect=responses)
+    monkeypatch.setattr(requests.Session, "get", get_mock)
+    monkeypatch.setattr("zimscraperlib.download.time.sleep", MagicMock())
+
+    fpath = tmp_path / "out.zim"
+    save_large_file("https://example.com/file.zim", fpath)
+
+    assert fpath.read_bytes() == b"ok"
+    assert get_mock.call_count == 2
+
+
+def test_save_large_file_restarts_when_server_ignores_range(
+    tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """If the server replies 200 (not 206) to a Range request, start over."""
+    fpath = tmp_path / "out.zim"
+    fpath.write_bytes(b"stale-partial-data")
+
+    get_mock = MagicMock(return_value=FakeResponse(200, [b"fresh"]))
+    monkeypatch.setattr(requests.Session, "get", get_mock)
+
+    save_large_file("https://example.com/file.zim", fpath)
+
+    assert fpath.read_bytes() == b"fresh"
+
+
+def test_save_large_file_does_not_retry_non_retryable_http_status(
+    tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    get_mock = MagicMock(return_value=FakeResponse(404))
+    monkeypatch.setattr(requests.Session, "get", get_mock)
+
+    fpath = tmp_path / "out.zim"
+    with pytest.raises(DownloadError):
+        save_large_file("https://example.com/file.zim", fpath)
+
+    assert get_mock.call_count == 1
+
+
+def test_save_large_file_gives_up_after_consecutive_failures_without_progress(
+    tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    get_mock = MagicMock(
+        side_effect=requests.ConnectionError("always fails before any byte")
+    )
+    monkeypatch.setattr(requests.Session, "get", get_mock)
+    monkeypatch.setattr("zimscraperlib.download.time.sleep", MagicMock())
+
+    fpath = tmp_path / "out.zim"
+    with pytest.raises(DownloadError):
+        save_large_file("https://example.com/file.zim", fpath, retries=2)
+
+    assert get_mock.call_count == 3  # initial attempt + 2 retries
+
+
+def test_save_large_file_keeps_retrying_as_long_as_progress_is_made(
+    tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Every attempt but the last writes a byte before failing; progress keeps
+    resetting the failure counter, so more than `retries` attempts in a row are
+    allowed as long as each one makes progress."""
+    failing_responses = [
+        FakeResponse(
+            200 if i == 0 else 206,
+            [b"x"],
+            raise_while_streaming=requests.ConnectionError("dropped"),
+        )
+        for i in range(4)
+    ]
+    responses = [*failing_responses, FakeResponse(206, [b"x"])]
+    get_mock = MagicMock(side_effect=responses)
+    monkeypatch.setattr(requests.Session, "get", get_mock)
+    monkeypatch.setattr("zimscraperlib.download.time.sleep", MagicMock())
+
+    fpath = tmp_path / "out.zim"
+    save_large_file("https://example.com/file.zim", fpath, retries=1)
+
+    assert get_mock.call_count == 5
+    assert fpath.read_bytes() == b"xxxxx"
+
+
 @pytest.mark.slow
 @pytest.mark.parametrize(
     "url,video_id",
@@ -353,7 +525,7 @@ def test_youtube_download_nowait(tmp_path: pathlib.Path):
 @pytest.mark.slow
 def test_youtube_download_error():
     yt_downloader = YoutubeDownloader(threads=1)
-    with pytest.raises(DownloadError, match=re.escape("is not a valid URL")):
+    with pytest.raises(YtDlpDownloadError, match=re.escape("is not a valid URL")):
         yt_downloader.download("11", BestMp4.get_options())
     yt_downloader.shutdown()