From 7ba29eeb3beda73b763dcfb857f668a1a125192f Mon Sep 17 00:00:00 2001
From: pawank925 <229111348+pawank925@users.noreply.github.com>
Date: Wed, 14 Oct 2026 11:10:44 +0530
Subject: [PATCH] Make save_large_file resumable with progress-based retries

save_large_file() shelled out to a single wget run with -t {retries}. wget's
-t is a budget for the whole download, not for consecutive failures: a
connection that drops every few hundred MB exhausts it and aborts with no
further retry, even though every attempt resumes correctly. That is what made
long downloads fail in production (see openzim/zimwright#3).

Stream the response with requests instead, deriving the resume point from the
bytes already on disk and sending a Range request. The retry counter now only
grows on attempts that make no progress at all: any attempt that writes at
least one byte resets it, so a flaky-but-advancing download keeps resuming
until it completes. Transient HTTP statuses (matching the existing requests
retry adapter) and connection/timeout/chunked-encoding errors are retried with
exponential backoff; anything else raises the new DownloadError. If the server
ignores the Range request, the file is started over.

This is the same approach already used by zimwright, which can now rely on the
library implementation.

Closes #348.
---
 CHANGELOG.md                    |   1 +
 src/zimscraperlib/download.py   | 103 +++++++++++++------
 tests/download/test_download.py | 175 +++++++++++++++++++++++++++++++-
 3 files changed, 247 insertions(+), 32 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1de7b951..018c058e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Rewrite `save_large_file` to stream and resume the download itself instead of running `wget` once, whose `-t` counter is a budget for the whole download: a flaky connection now keeps resuming and the retry budget only counts attempts that make no progress, so long downloads are no longer abandoned mid-way (#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)
diff --git a/src/zimscraperlib/download.py b/src/zimscraperlib/download.py
index 63c1ed06..7971f321 100644
--- a/src/zimscraperlib/download.py
+++ b/src/zimscraperlib/download.py
@@ -1,5 +1,5 @@
 import pathlib
-import subprocess
+import time
 from concurrent.futures import Future, ThreadPoolExecutor
 from typing import Any, ClassVar
 
@@ -121,33 +121,85 @@ class BestMp4(YoutubeConfig):
     }
 
 
+_RETRYABLE_HTTP_STATUS_CODES = (413, 429, 500, 502, 503, 504)
+_HTTP_PARTIAL_CONTENT = 206
+_DOWNLOAD_BLOCK_SIZE = 1024 * 1024
+_DOWNLOAD_BACKOFF_FACTOR = 1.0
+
+
+class DownloadError(Exception):
+    """Raised when a file could not be completely downloaded."""
+
+
 def save_large_file(
     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 its content with resume.
+
+    Unlike a single wget run, whose -t counter is a budget for the whole download,
+    the retry counter here only grows on attempts that make no progress at all:
+    any attempt that writes at least one byte resets it, so a download that keeps
+    resuming is not abandoned early.
 
     Arguments -
         url:
+        fpath:
+        retries: number of consecutive attempts without progress tolerated
+        user_agent:
     """
-    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: dict[str, str] = {"User-Agent": user_agent} if user_agent else {}
+    attempts_without_progress = 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}-"
+            made_progress = False
+            try:
+                with session.get(
+                    url,
+                    stream=True,
+                    headers=request_headers,
+                    timeout=DEFAULT_WEB_REQUESTS_TIMEOUT,
+                ) as response:
+                    if resume_from and response.status_code != _HTTP_PARTIAL_CONTENT:
+                        resume_from = 0
+                    response.raise_for_status()
+                    with open(fpath, "ab" if resume_from else "wb") as fh:
+                        for chunk in response.iter_content(_DOWNLOAD_BLOCK_SIZE):
+                            if not chunk:
+                                continue
+                            fh.write(chunk)
+                            made_progress = True
+            except requests.HTTPError as exc:
+                status = exc.response.status_code if exc.response is not None else None
+                if status not in _RETRYABLE_HTTP_STATUS_CODES:
+                    raise DownloadError(
+                        f"Failed to download {url}: HTTP status {status}"
+                    ) from exc
+            except (
+                requests.ConnectionError,
+                requests.Timeout,
+                requests.exceptions.ChunkedEncodingError,
+            ):
+                logger.warning(f"Download of {url} interrupted")
+            else:
+                return
+
+            if made_progress:
+                attempts_without_progress = 0
+            else:
+                attempts_without_progress += 1
+            if attempts_without_progress > retries:
+                raise DownloadError(
+                    f"Failed to download {url}: {retries} attempts without progress"
+                )
+            wait = _DOWNLOAD_BACKOFF_FACTOR * 2 ** (attempts_without_progress - 1)
+            logger.warning(
+                f"Retrying download of {url} in {wait:.0f}s from byte {resume_from}"
+            )
+            time.sleep(wait)
 
 
 def get_retry_adapter(
@@ -161,14 +213,7 @@ def get_retry_adapter(
         status=2,  # failure HTTP status (only those bellow)
         redirect=False,  # don't fail on redirections
         backoff_factor=30,  # sleep factor between retries
-        status_forcelist=[
-            413,
-            429,
-            500,
-            502,
-            503,
-            504,
-        ],  # force retry on the following codes
+        status_forcelist=list(_RETRYABLE_HTTP_STATUS_CODES),
     )
 
     return requests.adapters.HTTPAdapter(max_retries=retries)
diff --git a/tests/download/test_download.py b/tests/download/test_download.py
index 3a1bf8e7..d940531a 100644
--- a/tests/download/test_download.py
+++ b/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_WEB_REQUESTS_TIMEOUT
 from zimscraperlib.download import (
     BestMp4,
     BestWebm,
+    DownloadError,
     YoutubeDownloader,
     save_large_file,
     stream_file,
@@ -217,6 +219,173 @@ 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]:
+        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, without a 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
+    assert get_mock.call_args_list[1].kwargs["headers"]["Range"] == "bytes=6-"
+
+
+def test_save_large_file_retries_on_retryable_http_status(
+    tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    get_mock = MagicMock(side_effect=[FakeResponse(503), FakeResponse(200, [b"ok"])])
+    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
+
+
+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 index == 0 else 206,
+            [b"x"],
+            raise_while_streaming=requests.ConnectionError("dropped"),
+        )
+        for index in range(4)
+    ]
+    get_mock = MagicMock(side_effect=[*failing_responses, FakeResponse(206, [b"x"])])
+    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=2)
+
+    assert fpath.read_bytes() == b"xxxxx"
+    assert get_mock.call_count == 5
+
+
 @pytest.mark.slow
 @pytest.mark.parametrize(
     "url,video_id",
@@ -259,7 +428,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()