From a875016b1a6c43bf38d124c3da99c6be17c02c5d Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Sat, 5 Sep 2026 14:05:37 +0200 Subject: [PATCH] Show the other downloader's progress while a prefetch job waits on its lock --- esphome/framework_helpers.py | 11 +++ esphome/platformio/prefetch.py | 41 ++++++++--- tests/unit_tests/test_framework_helpers.py | 14 ++++ tests/unit_tests/test_platformio_prefetch.py | 75 +++++++++++++++++++- 4 files changed, 127 insertions(+), 14 deletions(-) diff --git a/esphome/framework_helpers.py b/esphome/framework_helpers.py index 82bc0d3727..ee50723a39 100644 --- a/esphome/framework_helpers.py +++ b/esphome/framework_helpers.py @@ -909,6 +909,17 @@ def _part_path(dest: Path) -> Path: return dest.with_name(dest.name + ".part") +def downloaded_bytes(dest: Path, size: int) -> int: + """Bytes of ``dest`` on disk: ``size`` once it landed, else what its + ``.part`` holds so far (capped at ``size``), else 0.""" + if dest.is_file(): + return size + try: + return min(_part_path(dest).stat().st_size, size) + except OSError: + return 0 + + def discard_partial_download(dest: Path) -> None: """Remove ``dest`` and the resume sidecars of an abandoned download.""" part = _part_path(dest) diff --git a/esphome/platformio/prefetch.py b/esphome/platformio/prefetch.py index 5097239065..41df598bd7 100644 --- a/esphome/platformio/prefetch.py +++ b/esphome/platformio/prefetch.py @@ -35,6 +35,7 @@ from typing import Any, NamedTuple from esphome.framework_helpers import ( content_length, discard_partial_download, + downloaded_bytes, failure_reason, resume_fetch_job, run_batch_downloads, @@ -462,16 +463,22 @@ def _uri_jobs( def _serialized_fetch_job( - dl_path: Path, lock_path: str, body: Any, unlocked_ok: bool = True + dl_path: Path, + lock_path: str, + body: Any, + size: int, + stream_dest: Path, + unlocked_ok: bool = True, ) -> Any: """Wrap ``body`` so the shared destination is single-writer. Interleaved writers truncate each other's ``.part`` bytes (see - registry.py). The bounded poll observes Ctrl-C via the tracker; a - blown deadline is a clean skip (the holder's copy is what the build - needs). On a lock-less filesystem a sha256-verified body runs - unlocked with one warning; a checksum-less one - (``unlocked_ok=False``) is a counted failure instead. + registry.py). While the lock is held elsewhere the poll reports the + holder's bytes (streamed into ``stream_dest``) so the bar moves, and + observes Ctrl-C via the tracker; a blown deadline is a clean skip + (the holder's copy is what the build needs). On a lock-less + filesystem a sha256-verified body runs unlocked with one warning; a + checksum-less one (``unlocked_ok=False``) is a counted failure instead. """ def run(tracker: Any) -> None: @@ -481,12 +488,21 @@ def _serialized_fetch_job( # filesystems that blocks every later build (see git.py) lock = FileLock(lock_path, fallback_to_soft=False) deadline = time.monotonic() + _DOWNLOAD_LOCK_TIMEOUT + waiting = False while True: try: lock.acquire(timeout=_URI_LOCK_POLL) break except Timeout: - tracker(0) # raises when the batch is cancelled + if not waiting: + waiting = True + _LOGGER.info( + "Waiting for another process downloading %s", dl_path.name + ) + # Raises when the batch is cancelled + tracker( + size if dl_path.is_file() else downloaded_bytes(stream_dest, size) + ) if time.monotonic() >= deadline: # Another process is fetching this same file; its copy # is what the build needs (a large framework archive @@ -506,7 +522,8 @@ def _serialized_fetch_job( break try: if dl_path.is_file(): - return # another process finished it while we waited + tracker(size) # another process finished it while we waited + return body(tracker) finally: if lock is not None: @@ -540,6 +557,8 @@ def _registry_fetch_job( dl_path, f"{dl_path}.esphome.lock", resume_fetch_job(url, dl_path, sha256=checksum, size=size), + size, + dl_path, ) def run(tracker: Any) -> None: @@ -571,9 +590,9 @@ def _uri_fetch_job(manager: Any, url: str, dl_path: Path, size: int) -> Any: tmp.replace(dl_path) def run(tracker: Any) -> None: - _serialized_fetch_job(dl_path, f"{tmp}.lock", promote, unlocked_ok=False)( - tracker - ) + _serialized_fetch_job( + dl_path, f"{tmp}.lock", promote, size, tmp, unlocked_ok=False + )(tracker) if dl_path.is_file(): # Won or lost, the race is over; staging files left behind # are dead weight PlatformIO's cache never prunes diff --git a/tests/unit_tests/test_framework_helpers.py b/tests/unit_tests/test_framework_helpers.py index fcc5572f51..b5d535330c 100644 --- a/tests/unit_tests/test_framework_helpers.py +++ b/tests/unit_tests/test_framework_helpers.py @@ -2353,3 +2353,17 @@ def test_discard_partial_download_logs_undeletable( ): framework_helpers.discard_partial_download(dest) assert "Could not remove" in caplog.text + + +def test_downloaded_bytes_reports_what_is_on_disk(tmp_path: Path) -> None: + """A landed file counts in full, a part file counts what it holds + (never more than the size), and nothing on disk counts zero.""" + dest = tmp_path / "archive" + assert framework_helpers.downloaded_bytes(dest, 4) == 0 + part = tmp_path / "archive.part" + part.write_bytes(b"ab") + assert framework_helpers.downloaded_bytes(dest, 4) == 2 + part.write_bytes(b"abcdef") + assert framework_helpers.downloaded_bytes(dest, 4) == 4 + dest.write_bytes(b"abcd") + assert framework_helpers.downloaded_bytes(dest, 4) == 4 diff --git a/tests/unit_tests/test_platformio_prefetch.py b/tests/unit_tests/test_platformio_prefetch.py index 379ef52ebd..c8e511b7c5 100644 --- a/tests/unit_tests/test_platformio_prefetch.py +++ b/tests/unit_tests/test_platformio_prefetch.py @@ -454,11 +454,17 @@ def test_uri_fetch_job_waits_out_a_briefly_held_lock(tmp_path: Path) -> None: assert dl_path.read_bytes() == b"data" -def test_lock_deadline_leaves_download_to_the_holder(tmp_path: Path) -> None: +@pytest.mark.parametrize("staged", [b"", b"ab"]) +def test_lock_deadline_leaves_download_to_the_holder( + tmp_path: Path, staged: bytes +) -> None: """A lock held past the deadline means another process is fetching the same file; skipping cleanly beats a misleading failure warning. The - tracker is still polled so a parked worker observes cancellation.""" + tracker is still polled so a parked worker observes cancellation, and + it reports what the holder has staged so far.""" dl_path = tmp_path / "archive" + if staged: + (tmp_path / "archive.prefetch.part").write_bytes(staged) ticks: list[int] = [] with ( patch("esphome.framework_helpers.download_with_resume") as mock_download, @@ -467,10 +473,73 @@ def test_lock_deadline_leaves_download_to_the_holder(tmp_path: Path) -> None: ): pf._uri_fetch_job(MagicMock(), "https://x/a.zip", dl_path, 4)(ticks.append) mock_download.assert_not_called() - assert ticks == [0] + assert ticks == [len(staged)] assert not dl_path.exists() +def test_lock_wait_reports_the_holders_progress( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + """While another process holds the lock the bar follows its part file, + and the job credits the full size once that process lands the archive + (the 0% bar of a build racing the dashboard for the same toolchain).""" + manager = MagicMock() + dl_path = tmp_path / "archive" + part = tmp_path / "archive.part" + ticks: list[int] = [] + polls = iter([b"a", b"abc"]) + + def acquire(*args, **kwargs): + try: + part.write_bytes(next(polls)) + except StopIteration: + part.unlink() + dl_path.write_bytes(b"abcd") + return + raise Timeout("held") + + with ( + patch("esphome.framework_helpers.download_with_resume") as mock_download, + patch("filelock.FileLock.acquire", side_effect=acquire), + patch("filelock.FileLock.release"), + caplog.at_level(logging.INFO), + ): + pf._registry_fetch_job(manager, "https://x/a.tar.gz", dl_path, "ab" * 32, 4)( + ticks.append + ) + mock_download.assert_not_called() + assert ticks == [1, 3, 4] + assert caplog.text.count("Waiting for another process downloading archive") == 1 + manager.set_download_utime.assert_called_once_with(str(dl_path)) + + +def test_uri_lock_wait_reads_the_staging_part(tmp_path: Path) -> None: + """A URL job's holder streams into the staging path, so the wait + reports that part file, not one beside the cache path.""" + dl_path = tmp_path / "archive" + staging_part = tmp_path / "archive.prefetch.part" + ticks: list[int] = [] + polls = iter([b"ab"]) + + def acquire(*args, **kwargs): + try: + staging_part.write_bytes(next(polls)) + except StopIteration: + staging_part.unlink() + dl_path.write_bytes(b"abcd") + return + raise Timeout("held") + + with ( + patch("esphome.framework_helpers.download_with_resume") as mock_download, + patch("filelock.FileLock.acquire", side_effect=acquire), + patch("filelock.FileLock.release"), + ): + pf._uri_fetch_job(MagicMock(), "https://x/a.zip", dl_path, 4)(ticks.append) + mock_download.assert_not_called() + assert ticks == [2, 4] + + def test_registry_lock_deadline_skips_registration(tmp_path: Path) -> None: """A registry job that lost the download race to another process must not stamp a nonexistent archive into pio's usage.db."""