From 712ee229e58bb621b2f4b90197ebf37ff2d8c602 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Sat, 22 Aug 2026 16:58:24 -0500 Subject: [PATCH] Prefetch pending registry packages in parallel under one combined bar --- esphome/platformio/registry.py | 64 ++++++++ tests/unit_tests/test_platformio_registry.py | 152 +++++++++++++++++++ 2 files changed, 216 insertions(+) diff --git a/esphome/platformio/registry.py b/esphome/platformio/registry.py index 99ccbe2334..acf5b69801 100644 --- a/esphome/platformio/registry.py +++ b/esphome/platformio/registry.py @@ -4,6 +4,7 @@ platformio package (identical bits, esphome's own download machinery).""" from __future__ import annotations from collections.abc import Collection +from concurrent.futures import ThreadPoolExecutor import io import json import logging @@ -13,6 +14,7 @@ import platform from esphome.core import EsphomeError from esphome.framework_helpers import ( + BatchDownloadProgress, archive_extract_all, download_from_mirrors, download_with_resume, @@ -140,6 +142,68 @@ def _check_layout(name: str, dest: Path, expect: Collection[str]) -> None: ) +def prefetch_packages( + packages: list[tuple[str, str, Path, list[str]]], downloads_dir: Path +) -> None: + """Download pending package archives in parallel under one combined bar. + + ``packages`` holds ``(name, version, dest, mirrors)`` per package. Purely + an optimization: ``install_package`` verifies every archive and + re-downloads anything this pass left unfinished. Mirror overrides and + registry entries without a size stay on the sequential path so its + per-file bars remain trustworthy. + """ + pending: list[tuple[str, str, str, str, int]] = [] + for name, version, dest, mirrors in packages: + if mirrors or (dest / ".esphome_extracted").is_file(): + continue + try: + url, sha256, size = registry_download(name, version) + except EsphomeError as err: + # The sequential install reports the real failure with context + _LOGGER.debug("Prefetch resolve for %s failed: %s", name, err) + continue + if not size: + continue + archive = downloads_dir / f"{name}-{version}" + if archive.is_file() and archive.stat().st_size == size: + continue + pending.append((name, version, url, sha256, size)) + if len(pending) < 2: + return + downloads_dir.mkdir(parents=True, exist_ok=True) + _LOGGER.info( + "Downloading %d package archive(s): %s", + len(pending), + ", ".join(name for name, _, _, _, _ in pending), + ) + progress = BatchDownloadProgress( + "Downloading packages", sum(size for *_, size in pending) + ) + + def _fetch(entry: tuple[str, str, str, str, int]) -> None: + name, version, url, sha256, size = entry + try: + download_with_resume( + url, + downloads_dir / f"{name}-{version}", + sha256=sha256, + size=size, + progress=progress.tracker(), + ) + except Exception as err: # noqa: BLE001 # pylint: disable=broad-exception-caught + # install_package retries this one itself, with a visible bar + _LOGGER.debug("Prefetch of %s failed: %s", name, err) + + ex = ThreadPoolExecutor(max_workers=len(pending)) + try: + for future in [ex.submit(_fetch, entry) for entry in pending]: + future.result() + finally: + ex.shutdown(wait=True, cancel_futures=True) + progress.done() + + def install_package( name: str, version: str, diff --git a/tests/unit_tests/test_platformio_registry.py b/tests/unit_tests/test_platformio_registry.py index 4f83048aab..cf995ad2f4 100644 --- a/tests/unit_tests/test_platformio_registry.py +++ b/tests/unit_tests/test_platformio_registry.py @@ -452,3 +452,155 @@ def test_registry_download_non_list_system_is_named() -> None: pytest.raises(EsphomeError, match="Unexpected package registry response"), ): registry.registry_download("pkg", "1.0.0") + + +def _resolve_for(sizes: dict[str, int | None]): + def resolve(name: str, version: str): + size = sizes[name] + if size == -1: + raise EsphomeError("registry down") + return (f"http://x/{name}.tar.gz", "abc123", size) + + return resolve + + +def test_prefetch_packages_downloads_pending_in_parallel(tmp_path: Path) -> None: + """Two uninstalled packages download together under one combined bar, + with the registry's sha256 and size and a batch progress tracker.""" + with ( + patch.object(registry, "download_with_resume") as mock_download, + patch.object( + registry, "registry_download", side_effect=_resolve_for({"a": 10, "b": 20}) + ), + ): + registry.prefetch_packages( + [ + ("a", "1.0", tmp_path / "a", []), + ("b", "2.0", tmp_path / "b", []), + ], + tmp_path / "dl", + ) + assert mock_download.call_count == 2 + for call, (name, version, size) in zip( + mock_download.call_args_list, [("a", "1.0", 10), ("b", "2.0", 20)], strict=True + ): + assert call[0][0] == f"http://x/{name}.tar.gz" + assert call[0][1] == tmp_path / "dl" / f"{name}-{version}" + assert call[1]["sha256"] == "abc123" + assert call[1]["size"] == size + assert callable(call[1]["progress"]) + + +def test_prefetch_packages_single_pending_skips(tmp_path: Path) -> None: + """One pending package has nothing to parallelize; the sequential + install keeps its own bar.""" + marker_dest = tmp_path / "a" + marker_dest.mkdir() + (marker_dest / ".esphome_extracted").touch() + with ( + patch.object(registry, "download_with_resume") as mock_download, + patch.object( + registry, "registry_download", side_effect=_resolve_for({"b": 20}) + ), + ): + registry.prefetch_packages( + [ + ("a", "1.0", marker_dest, []), + ("b", "2.0", tmp_path / "b", []), + ], + tmp_path / "dl", + ) + mock_download.assert_not_called() + + +def test_prefetch_packages_mirror_and_sizeless_stay_sequential( + tmp_path: Path, +) -> None: + """Mirror overrides and size-less registry entries are left to the + sequential path so its per-file bars stay trustworthy.""" + with ( + patch.object(registry, "download_with_resume") as mock_download, + patch.object( + registry, + "registry_download", + side_effect=_resolve_for({"b": None, "c": 30}), + ), + ): + registry.prefetch_packages( + [ + ("a", "1.0", tmp_path / "a", ["http://mirror/{VERSION}"]), + ("b", "2.0", tmp_path / "b", []), + ("c", "3.0", tmp_path / "c", []), + ], + tmp_path / "dl", + ) + mock_download.assert_not_called() + + +def test_prefetch_packages_resolve_failure_defers_to_install( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + """A registry failure only skips the prefetch; install_package reports + the real error with context.""" + caplog.set_level("DEBUG") + with ( + patch.object(registry, "download_with_resume") as mock_download, + patch.object( + registry, "registry_download", side_effect=_resolve_for({"a": -1, "b": 20}) + ), + ): + registry.prefetch_packages( + [ + ("a", "1.0", tmp_path / "a", []), + ("b", "2.0", tmp_path / "b", []), + ], + tmp_path / "dl", + ) + mock_download.assert_not_called() + assert "Prefetch resolve for a failed" in caplog.text + + +def test_prefetch_packages_complete_archive_skipped(tmp_path: Path) -> None: + """An archive already fully downloaded is not re-fetched.""" + dl = tmp_path / "dl" + dl.mkdir() + (dl / "a-1.0").write_bytes(b"x" * 10) + with ( + patch.object(registry, "download_with_resume") as mock_download, + patch.object( + registry, "registry_download", side_effect=_resolve_for({"a": 10, "b": 20}) + ), + ): + registry.prefetch_packages( + [ + ("a", "1.0", tmp_path / "a", []), + ("b", "2.0", tmp_path / "b", []), + ], + dl, + ) + mock_download.assert_not_called() + + +def test_prefetch_packages_download_failure_is_debug( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + """A failed prefetch download is logged and left for install_package.""" + caplog.set_level("DEBUG") + with ( + patch.object( + registry, "download_with_resume", side_effect=OSError("boom") + ) as mock_download, + patch.object( + registry, "registry_download", side_effect=_resolve_for({"a": 10, "b": 20}) + ), + ): + registry.prefetch_packages( + [ + ("a", "1.0", tmp_path / "a", []), + ("b", "2.0", tmp_path / "b", []), + ], + tmp_path / "dl", + ) + assert mock_download.call_count == 2 + assert "Prefetch of a failed" in caplog.text + assert "Prefetch of b failed" in caplog.text