[espidf] Download tool archives in parallel during the install prefetch

esphome/espidf/framework.py

@@ -1,6 +1,7 @@
 """ESP-IDF framework tools for ESPHome."""
 
 from collections.abc import Callable
+from concurrent.futures import ThreadPoolExecutor
 from ctypes.util import find_library
 import json
 import logging
@@ -15,6 +16,7 @@ import platformdirs
 
 from esphome.core import CORE, Version
 from esphome.framework_helpers import (
+    BatchDownloadProgress,
     PathType,
     archive_extract_all,
     create_venv,
@@ -690,6 +692,12 @@ def _patch_tools_json_demote_unused_tools(framework_path: Path) -> None:
     )
 
 
+# Tool archives are large (tens to hundreds of MB) and served by GitHub /
+# dl.espressif.com; a few streams at once saturate most links without
+# hammering the host. Smaller than external_files' 8: those are tiny files.
+_PREFETCH_WORKERS = 4
+
+
 def _prefetch_idf_tool_archives(
     framework_path: Path,
     targets_str: str,
@@ -702,10 +710,10 @@ def _prefetch_idf_tool_archives(
     which makes large archives effectively impossible to fetch on unstable
     connections (#17703). This asks the framework's idf_tools (via
     ``get_tool_downloads.py``) which archives the coming install needs, then
-    downloads each into ``<IDF_TOOLS_PATH>/dist`` with
-    ``download_with_resume``. The installer then finds the verified archives
-    already in place ("file ... is already downloaded") and never touches the
-    network.
+    downloads them into ``<IDF_TOOLS_PATH>/dist`` with
+    ``download_with_resume``, a few at a time under one combined progress
+    bar. The installer then finds the verified archives already in place
+    ("file ... is already downloaded") and never touches the network.
 
     Strictly best-effort: any failure here just logs and returns, leaving
     ``idf_tools.py install`` to download whatever is missing exactly as
@@ -732,21 +740,34 @@ def _prefetch_idf_tool_archives(
             for entry in json.loads(stdout)
             if not (dist_path / entry["dest"]).is_file()
         ]
-        for index, entry in enumerate(entries, start=1):
-            _LOGGER.info(
-                "Downloading %s (%d/%d) ...", entry["name"], index, len(entries)
-            )
+        if not entries:
+            return
+        _LOGGER.info(
+            "Downloading %d ESP-IDF tool archive(s): %s",
+            len(entries),
+            ", ".join(entry["name"] for entry in entries),
+        )
+        progress = BatchDownloadProgress(
+            "Downloading ESP-IDF tools",
+            sum(entry["size"] or 0 for entry in entries),
+        )
+
+        def _download(entry: dict) -> None:
             try:
                 download_with_resume(
                     entry["url"],
                     dist_path / entry["dest"],
                     sha256=entry["sha256"],
                     size=entry["size"],
+                    progress=progress.tracker(),
                 )
             except Exception as e:  # noqa: BLE001  # pylint: disable=broad-exception-caught
                 # Keep prefetching the remaining archives; the installer
                 # will retry this one itself (without resume).
                 _LOGGER.warning("Could not prefetch %s: %s", entry["name"], e)
+
+        with ThreadPoolExecutor(max_workers=min(_PREFETCH_WORKERS, len(entries))) as ex:
+            list(ex.map(_download, entries))
     except Exception as e:  # noqa: BLE001  # pylint: disable=broad-exception-caught
         # The installer downloads anything missing itself; never let the
         # prefetch become a new way for the install to fail.

esphome/framework_helpers.py

@@ -1,6 +1,6 @@
 """Generic toolchain installation helpers shared across framework implementations."""
 
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
 from contextlib import ExitStack
 import hashlib
 import io
@@ -10,6 +10,7 @@ import os
 from pathlib import Path
 import subprocess
 import sys
+import threading
 import time
 from typing import IO, TYPE_CHECKING
 
@@ -697,7 +698,11 @@ def _response_validator(resp: "requests.Response") -> str | None:
 
 
 def _stream_response_to_file(
-    resp: "requests.Response", f: IO[bytes], offset: int, size: int | None = None
+    resp: "requests.Response",
+    f: IO[bytes],
+    offset: int,
+    size: int | None = None,
+    progress: Callable[[int], None] | None = None,
 ) -> None:
     """Stream an open ``_open_ranged`` response body into ``f`` at ``offset``.
 
@@ -705,21 +710,60 @@ def _stream_response_to_file(
     (effective offset 0) discards the stale bytes. ``offset`` also seeds the
     progress bar so a resumed download shows overall progress. ``size`` is
     the known full file size; when None it is derived from the response's
-    content-length, and without either there is no progress bar.
+    content-length, and without either there is no progress bar. With
+    ``progress`` set, no bar is drawn here; the callback gets the absolute
+    byte count, seeded with ``offset`` and then after each chunk.
     """
     f.seek(offset)
     f.truncate(offset)
     total_size = size or offset + _content_length(resp)
     downloaded = offset
-    progress = ProgressBar("Downloading") if total_size > 0 else None
+    own_bar: ProgressBar | None = None
+    if progress is None:
+        own_bar = ProgressBar("Downloading") if total_size > 0 else None
+        progress = (
+            (lambda done: own_bar.update(done / total_size))
+            if own_bar
+            else (lambda _: None)
+        )
+    progress(downloaded)
     for chunk in resp.iter_content(chunk_size=256 * 1024):
         if chunk:
             f.write(chunk)
             downloaded += len(chunk)
-        if progress is not None:
-            progress.update(downloaded / total_size)
-    if progress is not None:
-        progress.update(1)
+        progress(downloaded)
+    if own_bar is not None:
+        own_bar.update(1)
+
+
+class BatchDownloadProgress:
+    """One progress bar across several concurrent ``download_with_resume`` calls.
+
+    Each ``tracker()`` is a ``progress`` callback for one download; it reports
+    that file's absolute byte count and the bar shows the sum over ``total``.
+    The lock also serialises the bar's stderr writes, so worker threads never
+    interleave frames. With an unknown ``total`` (0) nothing is drawn.
+    """
+
+    def __init__(self, header: str, total: int) -> None:
+        self._bar = ProgressBar(header) if total > 0 else None
+        self._total = total
+        self._sum = 0
+        self._lock = threading.Lock()
+
+    def tracker(self) -> Callable[[int], None]:
+        last = 0
+
+        def update(done: int) -> None:
+            nonlocal last
+            if self._bar is None:
+                return
+            with self._lock:
+                self._sum += done - last
+                last = done
+                self._bar.update(min(self._sum / self._total, 1))
+
+        return update
 
 
 def download_with_resume(
@@ -732,6 +776,7 @@ def download_with_resume(
     attempts: int = 5,
     timeout: int = 30,
     retry_connect_errors: bool = True,
+    progress: Callable[[int], None] | None = None,
 ) -> None:
     """Download ``url`` to ``dest``, resuming partial downloads.
 
@@ -754,6 +799,12 @@ def download_with_resume(
     of consuming attempts — for callers with their own fallback, like
     ``download_from_mirrors``.
 
+    ``progress``, when given, replaces the built-in progress bar: it is called
+    with the absolute number of bytes of ``dest`` obtained so far (including
+    a resumed prefix, and the final size once the file is verified), so a
+    caller running several downloads at once can draw one combined bar (see
+    ``BatchDownloadProgress``).
+
     Raises EsphomeError when all attempts are exhausted.
     """
     # Imported lazily: requests is a heavy import (~85ms) and is only needed
@@ -777,6 +828,8 @@ def download_with_resume(
     if dest.is_file() and (sha256 is not None or size is not None):
         try:
             _verify_file(dest, sha256, size)
+            if progress is not None:
+                progress(dest.stat().st_size)
             return
         except EsphomeError:
             dest.unlink()
@@ -822,7 +875,7 @@ def download_with_resume(
                             # Recorded so a later run can prove an If-Range
                             # resume of this part file safe.
                             _write_download_meta(meta, url, validator, expected_total)
-                        _stream_response_to_file(resp, f, offset, size)
+                        _stream_response_to_file(resp, f, offset, size, progress)
             # else: a previous run already wrote every byte (or more) but
             # was killed before the rename below. Skip the network entirely
             # — a Range request past EOF would draw HTTP 416 — and let
@@ -831,6 +884,10 @@ def download_with_resume(
 
             expected_size = size if size is not None else expected_total
             _verify_file(part, sha256, expected_size or None)
+            if progress is not None:
+                # Also credits a part file an earlier run completed without
+                # streaming anything this time.
+                progress(part.stat().st_size)
             if not expected_size and sha256 is None:
                 # No sha, no size, and the server sent no usable
                 # content-length: nothing can prove the download complete

tests/unit_tests/test_espidf_framework.py

@@ -2,6 +2,7 @@
 
 # pylint: disable=protected-access
 
+from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
 import importlib.util
 import io
@@ -899,12 +900,68 @@ def test_prefetch_downloads_each_archive_with_resume(tmp_path: Path) -> None:
         _prefetch_idf_tool_archives(tmp_path, "esp32", ["required"], None)
 
     dist = get_idf_tools_path() / "dist"
-    assert download.call_count == 2
-    assert download.call_args_list[0][0] == (
-        "https://example.com/cmake.tar.gz",
-        dist / "cmake-3.30.2.tar.gz",
+    # Archives download concurrently, so the call order is not fixed.
+    calls = {call[0]: call[1] for call in download.call_args_list}
+    assert set(calls) == {
+        ("https://example.com/cmake.tar.gz", dist / "cmake-3.30.2.tar.gz"),
+        ("https://example.com/ninja.zip", dist / "ninja.zip"),
+    }
+    kwargs = calls[("https://example.com/cmake.tar.gz", dist / "cmake-3.30.2.tar.gz")]
+    assert kwargs["sha256"] == "ab" * 32
+    assert kwargs["size"] == 123
+    # every archive reports into the one combined progress bar
+    assert all(
+        kw["progress"].__qualname__.startswith("BatchDownloadProgress.tracker")
+        for kw in calls.values()
     )
-    assert download.call_args_list[0][1] == {"sha256": "ab" * 32, "size": 123}
+
+
+def test_prefetch_downloads_archives_concurrently(tmp_path: Path) -> None:
+    """More than one archive fans out over a bounded thread pool."""
+    entries = [
+        {
+            "name": f"tool{i}@1",
+            "url": f"https://example.com/tool{i}.tar.gz",
+            "size": 10,
+            "sha256": "ab" * 32,
+            "dest": f"tool{i}.tar.gz",
+        }
+        for i in range(6)
+    ]
+    with (
+        patch(
+            "esphome.espidf.framework.run_command",
+            return_value=(True, json.dumps(entries), ""),
+        ),
+        patch("esphome.espidf.framework.download_with_resume") as download,
+        patch("esphome.espidf.framework.get_system_python_path", return_value="python"),
+        patch(
+            "esphome.espidf.framework.ThreadPoolExecutor", wraps=ThreadPoolExecutor
+        ) as pool,
+    ):
+        _prefetch_idf_tool_archives(tmp_path, "esp32", ["required"], None)
+
+    pool.assert_called_once_with(max_workers=4)
+    assert download.call_count == 6
+
+
+def test_prefetch_single_archive_uses_one_worker(tmp_path: Path) -> None:
+    entries = json.loads(_PREFETCH_JSON)[:1]
+    with (
+        patch(
+            "esphome.espidf.framework.run_command",
+            return_value=(True, json.dumps(entries), ""),
+        ),
+        patch("esphome.espidf.framework.download_with_resume") as download,
+        patch("esphome.espidf.framework.get_system_python_path", return_value="python"),
+        patch(
+            "esphome.espidf.framework.ThreadPoolExecutor", wraps=ThreadPoolExecutor
+        ) as pool,
+    ):
+        _prefetch_idf_tool_archives(tmp_path, "esp32", ["required"], None)
+
+    pool.assert_called_once_with(max_workers=1)
+    assert download.call_count == 1
 
 
 def test_prefetch_skips_already_downloaded_archives(tmp_path: Path) -> None:
@@ -964,6 +1021,11 @@ def test_prefetch_one_failed_archive_does_not_stop_the_rest(
 ) -> None:
     """A single archive failing its download must not abort the prefetch of
     the remaining archives."""
+
+    def _fail_cmake_download(url: str, *args, **kwargs) -> None:
+        if "cmake" in url:
+            raise OSError("network down")
+
     with (
         patch(
             "esphome.espidf.framework.run_command",
@@ -971,7 +1033,7 @@ def test_prefetch_one_failed_archive_does_not_stop_the_rest(
         ),
         patch(
             "esphome.espidf.framework.download_with_resume",
-            side_effect=[OSError("network down"), None],
+            side_effect=_fail_cmake_download,
         ) as download,
         patch("esphome.espidf.framework.get_system_python_path", return_value="python"),
     ):

tests/unit_tests/test_framework_helpers.py

@@ -21,6 +21,7 @@ import requests as req
 from esphome import framework_helpers
 from esphome.core import EsphomeError
 from esphome.framework_helpers import (
+    BatchDownloadProgress,
     _7z_extract_all,
     _detect_archive_root,
     _is_transient_download_error,
@@ -1112,6 +1113,77 @@ class TestDownloadWithResume:
         assert mock_get.call_args[1]["headers"] == {}
         assert dest.read_bytes() == b"data"
 
+    def test_progress_callback_reports_absolute_bytes(self, tmp_path: Path) -> None:
+        """With a callback no bar is drawn; the callback sees the running
+        byte count of this file, then its final verified size."""
+        dest = tmp_path / "tool.tar.gz"
+        resp = _mock_response(b"")
+        resp.headers = {"content-length": "7"}
+        resp.iter_content.return_value = [b"1234", b"567"]
+        seen: list[int] = []
+        with (
+            patch("requests.get", return_value=resp),
+            patch("esphome.framework_helpers.ProgressBar") as bar,
+        ):
+            download_with_resume(
+                "https://example.com/t", dest, size=7, progress=seen.append
+            )
+        assert seen == [0, 4, 7, 7]
+        bar.assert_not_called()
+
+    def test_progress_callback_seeds_with_resume_offset(self, tmp_path: Path) -> None:
+        dest = tmp_path / "tool.tar.gz"
+        (tmp_path / "tool.tar.gz.part").write_bytes(b"12345")
+        good = hashlib.sha256(b"12345678").hexdigest()
+        seen: list[int] = []
+        with patch("requests.get", return_value=_resumed_response(b"678")):
+            download_with_resume(
+                "https://example.com/t", dest, sha256=good, size=8, progress=seen.append
+            )
+        assert seen[0] == 5
+        assert seen[-1] == 8
+
+    def test_progress_callback_credits_already_complete_download(
+        self, tmp_path: Path
+    ) -> None:
+        """A verified dest from an earlier run still counts toward the batch."""
+        dest = tmp_path / "tool.tar.gz"
+        dest.write_bytes(b"12345678")
+        seen: list[int] = []
+        with patch("requests.get") as mock_get:
+            download_with_resume(
+                "https://example.com/t", dest, size=8, progress=seen.append
+            )
+        mock_get.assert_not_called()
+        assert seen == [8]
+
+
+class TestBatchDownloadProgress:
+    def test_sums_trackers_into_one_bar(self) -> None:
+        with patch("esphome.framework_helpers.ProgressBar") as bar_cls:
+            progress = BatchDownloadProgress("Downloading", 100)
+            a = progress.tracker()
+            b = progress.tracker()
+            a(10)
+            b(20)
+            a(30)
+            a(0)  # a restart from zero takes that file's bytes back out
+        bar_cls.assert_called_once_with("Downloading")
+        updates = [c[0][0] for c in bar_cls.return_value.update.call_args_list]
+        assert updates == [0.1, 0.3, 0.5, 0.2]
+
+    def test_clamps_at_one(self) -> None:
+        """Sizes are advisory; an over-delivering server never pushes past 100%."""
+        with patch("esphome.framework_helpers.ProgressBar") as bar_cls:
+            progress = BatchDownloadProgress("Downloading", 10)
+            progress.tracker()(25)
+        assert bar_cls.return_value.update.call_args[0][0] == 1
+
+    def test_unknown_total_draws_nothing(self) -> None:
+        with patch("esphome.framework_helpers.ProgressBar") as bar_cls:
+            BatchDownloadProgress("Downloading", 0).tracker()(5)
+        bar_cls.assert_not_called()
+
 
 class TestDownloadFromMirrors:
     def test_success_returns_url_and_writes_content(self, tmp_path: Path) -> None: