From 712ee229e58bb621b2f4b90197ebf37ff2d8c602 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Sat, 22 Aug 2026 16:58:24 -0500 Subject: [PATCH 1/2] 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 From 0c5586c7dcb4c8668545b6240fda4f3aeeb978e6 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Sat, 22 Aug 2026 16:59:07 -0500 Subject: [PATCH 2/2] Download the Arduino core and toolchain archives in parallel --- esphome/arduino8266/framework.py | 22 +++++++++++++++++-- .../unit_tests/test_arduino8266_framework.py | 19 ++++++++++++++++ 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/esphome/arduino8266/framework.py b/esphome/arduino8266/framework.py index f680815c31..125238a30a 100644 --- a/esphome/arduino8266/framework.py +++ b/esphome/arduino8266/framework.py @@ -22,7 +22,7 @@ from esphome.build_helpers.ninja import find_ninja from esphome.build_helpers.tools_cache import tools_cache_path from esphome.core import EsphomeError, Version from esphome.framework_helpers import str_to_lst_of_str -from esphome.platformio.registry import install_package +from esphome.platformio.registry import install_package, prefetch_packages FRAMEWORK_PACKAGE = "framework-arduinoespressif8266" TOOLCHAIN_PACKAGE = "toolchain-xtensa" @@ -101,6 +101,25 @@ def check_and_install(framework_version: Version) -> InstalledPaths: package_version = framework_package_version(framework_version) framework_path = get_framework_path(package_version) downloads_dir = get_arduino8266_tools_path() / "downloads" + toolchain_path = get_toolchain_path() + # Fetch both archives at once; the installs below verify and extract + prefetch_packages( + [ + ( + FRAMEWORK_PACKAGE, + package_version, + framework_path, + ESPHOME_ARDUINO8266_FRAMEWORK_MIRRORS, + ), + ( + TOOLCHAIN_PACKAGE, + TOOLCHAIN_VERSION, + toolchain_path, + ESPHOME_ARDUINO8266_TOOLCHAIN_MIRRORS, + ), + ], + downloads_dir, + ) install_package( FRAMEWORK_PACKAGE, package_version, @@ -109,7 +128,6 @@ def check_and_install(framework_version: Version) -> InstalledPaths: downloads_dir, expect=("cores/esp8266", "tools/sdk", "libraries"), ) - toolchain_path = get_toolchain_path() install_package( TOOLCHAIN_PACKAGE, TOOLCHAIN_VERSION, diff --git a/tests/unit_tests/test_arduino8266_framework.py b/tests/unit_tests/test_arduino8266_framework.py index 2837be1905..932958d667 100644 --- a/tests/unit_tests/test_arduino8266_framework.py +++ b/tests/unit_tests/test_arduino8266_framework.py @@ -63,6 +63,7 @@ def test_check_and_install_returns_paths(tmp_path: Path) -> None: with ( patch.dict(os.environ, {"ESPHOME_ARDUINO8266_PREFIX": str(tmp_path)}), patch.object(framework, "install_package") as mock_install, + patch.object(framework, "prefetch_packages") as mock_prefetch, patch.object(framework, "find_ninja", return_value=tmp_path / "ninja"), ): paths = framework.check_and_install(cv.Version(3, 1, 2)) @@ -89,6 +90,24 @@ def test_check_and_install_returns_paths(tmp_path: Path) -> None: tmp_path / "downloads", ) assert tc_call.kwargs["expect"] == ("bin", "xtensa-lx106-elf") + # The prefetch sees the same package specs as the installs + assert mock_prefetch.call_args.args == ( + [ + ( + framework.FRAMEWORK_PACKAGE, + "3.30102.0", + tmp_path / "frameworks" / "3.30102.0", + framework.ESPHOME_ARDUINO8266_FRAMEWORK_MIRRORS, + ), + ( + framework.TOOLCHAIN_PACKAGE, + framework.TOOLCHAIN_VERSION, + tmp_path / "toolchains" / framework.TOOLCHAIN_VERSION, + framework.ESPHOME_ARDUINO8266_TOOLCHAIN_MIRRORS, + ), + ], + tmp_path / "downloads", + ) def test_get_build_env_prepends_toolchain_bin(tmp_path: Path) -> None: