From 85a33d5d7ff5331ee7d97be97bdfcaf6782366dd Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Thu, 27 Aug 2026 18:17:29 -0500 Subject: [PATCH] Announce batched refetches and cap preinstall workers --- esphome/platformio/prefetch.py | 3 +- esphome/platformio/registry.py | 24 +++++++--- tests/unit_tests/test_platformio_prefetch.py | 49 ++++++++++++++++++++ tests/unit_tests/test_platformio_registry.py | 38 +++++++++++++++ 4 files changed, 106 insertions(+), 8 deletions(-) diff --git a/esphome/platformio/prefetch.py b/esphome/platformio/prefetch.py index ef8c27c9aa..8575ca246f 100644 --- a/esphome/platformio/prefetch.py +++ b/esphome/platformio/prefetch.py @@ -32,6 +32,7 @@ import time from typing import Any, NamedTuple from esphome.framework_helpers import ( + BATCH_EXTRACT_WORKERS, content_length, discard_partial_download, failure_reason, @@ -681,7 +682,7 @@ def _preinstall( would hang, not fail). Waves skip dependencies; the installed manifests feed the next wave. Any failure falls back to pio run. """ - workers = min(get_usable_cpu_count(), len(entries)) + workers = min(get_usable_cpu_count(), len(entries), BATCH_EXTRACT_WORKERS) # One manager per worker (_install mutates instance state); built # serially because construction rewires the shared manager logger managers: SimpleQueue = SimpleQueue() diff --git a/esphome/platformio/registry.py b/esphome/platformio/registry.py index 11f7c49c44..af2817c531 100644 --- a/esphome/platformio/registry.py +++ b/esphome/platformio/registry.py @@ -323,13 +323,23 @@ def install_package( ) else: url, sha256, size = registry_download(name, version) - # Batched: no private bar, no bytes (the shared bar must never - # run backwards), but the zero tick keeps cancellation observable - download_progress = ( - None - if extract_progress is None - else lambda _done: extract_progress(0.0) - ) + if extract_progress is None: + download_progress = None + else: + announced = False + + # Batched: no private bar, no bytes (the shared bar must + # never run backwards); the zero tick keeps cancellation + # observable and a real refetch is announced once + def _batched_progress(done: int) -> None: + nonlocal announced + if not announced and size and done < size: + _LOGGER.info("Re-downloading %s %s ...", name, version) + announced = True + extract_progress(0.0) + + download_progress = _batched_progress + download_with_resume( url, archive, sha256=sha256, size=size, progress=download_progress ) diff --git a/tests/unit_tests/test_platformio_prefetch.py b/tests/unit_tests/test_platformio_prefetch.py index 91fb78c6af..ac74417c5d 100644 --- a/tests/unit_tests/test_platformio_prefetch.py +++ b/tests/unit_tests/test_platformio_prefetch.py @@ -1729,3 +1729,52 @@ def test_platformio_private_api_contract() -> None: derived = PackageSpec("https://x/y/archive/master.zip") assert derived.name and not derived.has_custom_name() assert PackageSpec("Foo=https://x/y/archive/master.zip").has_custom_name() + + +def test_preinstall_caps_workers( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + """A high core count is capped; the workers share one disk.""" + + class _Mgr: + package_dir = str(tmp_path) + compatibility = None + + def __init__(self, package_dir, **kwargs) -> None: + pass + + def lock(self) -> None: + pass + + def unlock(self) -> None: + pass + + def memcache_reset(self) -> None: + pass + + def get_tmp_dir(self) -> str: + return str(tmp_path) + + def get_download_dir(self) -> str: + return str(tmp_path) + + def get_package(self, spec): + return None + + def get_pkg_dependencies(self, pkg): + return None + + def _install(self, spec, skip_dependencies, compatibility=None) -> None: + pass + + with ( + caplog.at_level(logging.INFO), + patch.object(pf, "get_usable_cpu_count", return_value=64), + ): + pf._preinstall( + _Mgr(str(tmp_path)), + [(f"p{i}@1", _FakeSpec(name=f"p{i}")) for i in range(11)], + ) + assert "Installing 11 PlatformIO package(s) with 10 extraction worker(s)" in ( + caplog.text + ) diff --git a/tests/unit_tests/test_platformio_registry.py b/tests/unit_tests/test_platformio_registry.py index 7a42dceb19..5957d87411 100644 --- a/tests/unit_tests/test_platformio_registry.py +++ b/tests/unit_tests/test_platformio_registry.py @@ -927,3 +927,41 @@ def test_install_packages_caps_workers(tmp_path: Path) -> None: ): registry.install_packages(specs, dl) assert batch.call_args.kwargs["max_workers"] == 10 + + +def test_install_package_batched_refetch_announced_once( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + """A batched archive that fails verification and refetches is announced; + a verify no-op (full size credited immediately) stays silent.""" + dest = tmp_path / "pkg" + (tmp_path / "dl").mkdir() + (tmp_path / "dl" / "pkg-1.0.0").write_bytes(b"x") + with ( + caplog.at_level(logging.INFO), + patch.object(registry, "download_with_resume") as mock_download, + patch.object(registry, "archive_extract_all") as mock_extract, + patch.object( + registry, + "registry_download", + return_value=("http://x/pkg.tar.gz", "abc123", 42), + ), + ): + mock_extract.side_effect = lambda *_a, **_kw: (dest / "payload").mkdir( + parents=True + ) + registry.install_package( + "pkg", + "1.0.0", + dest, + [], + tmp_path / "dl", + expect=("payload",), + extract_progress=lambda _frac: None, + ) + progress = mock_download.call_args[1]["progress"] + progress(42) + assert "Re-downloading pkg 1.0.0" not in caplog.text + progress(10) + progress(20) + assert caplog.text.count("Re-downloading pkg 1.0.0") == 1