Share the lock wait with the ESP-IDF registry prefetch and trim the docstrings

esphome/framework_helpers.py

@@ -23,6 +23,7 @@ from esphome.net_retry import (
 )
 
 if TYPE_CHECKING:
+    from filelock import FileLock
     import requests
 
 PathType = str | os.PathLike
@@ -909,6 +910,48 @@ def _part_path(dest: Path) -> Path:
     return dest.with_name(dest.name + ".part")
 
 
+def downloaded_bytes(dest: Path, size: int | None = None) -> int:
+    """Bytes of ``dest`` on disk: its ``.part`` so far (capped at ``size``),
+    else the landed file, else 0."""
+    try:
+        done = _part_path(dest).stat().st_size
+    except OSError:
+        if not dest.is_file():
+            return 0
+        done = dest.stat().st_size if size is None else size
+    return done if size is None else min(done, size)
+
+
+# Short lock-acquire slices so a waiting worker still observes Ctrl-C
+_DOWNLOAD_LOCK_POLL = 1
+
+
+def wait_for_download_lock(
+    lock: "FileLock",
+    tracker: Callable[[int], None],
+    on_disk: Callable[[], int],
+    name: str,
+    timeout: float | None = None,
+) -> bool:
+    """Acquire ``lock``, reporting ``on_disk()`` to ``tracker`` each poll so the
+    bar follows the holder's download. False once ``timeout`` seconds pass."""
+    from filelock import Timeout
+
+    deadline = None if timeout is None else time.monotonic() + timeout
+    waiting = False
+    while True:
+        try:
+            lock.acquire(timeout=_DOWNLOAD_LOCK_POLL)
+            return True
+        except Timeout:
+            if not waiting:
+                waiting = True
+                _LOGGER.info("Waiting for another process downloading %s", name)
+            tracker(on_disk())  # raises when the batch is cancelled
+            if deadline is not None and time.monotonic() >= deadline:
+                return False
+
+
 def discard_partial_download(dest: Path) -> None:
     """Remove ``dest`` and the resume sidecars of an abandoned download."""
     part = _part_path(dest)
@@ -1319,10 +1362,7 @@ def download_from_mirrors(
             )
             # Tick with the bytes already on disk so a combined bar holds
             # steady during the backoff instead of rewinding to zero
-            done = 0
-            if progress is not None:
-                part = _part_path(path_target)
-                done = part.stat().st_size if part.is_file() else 0
+            done = downloaded_bytes(path_target) if progress is not None else 0
             _cancellable_sleep(delay, progress, done)
 
     # 3. Report every attempted URL if all mirrors failed. failures spans

esphome/platformio/prefetch.py

@@ -19,6 +19,7 @@ from __future__ import annotations
 from collections.abc import Iterator
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager, suppress
+from functools import partial
 import hashlib
 import json
 import logging
@@ -35,9 +36,11 @@ 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,
+    wait_for_download_lock,
     warn_prefetch_failures,
 )
 from esphome.helpers import get_bool_env, get_usable_cpu_count, rmtree
@@ -68,9 +71,6 @@ _DOWNLOAD_LOCK_TIMEOUT = 60
 # the interpreter's own import-failure exit
 _EXIT_HANDLED = 3
 
-# Short lock-acquire slices so a waiting worker still observes Ctrl-C
-_URI_LOCK_POLL = 1
-
 # Resolution errored (vs a clean skip); suppresses the warm sentinel
 _RESOLVE_FAILED = object()
 
@@ -462,51 +462,50 @@ 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 | None = None,
+    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.
+    """Wrap ``body`` so the shared destination is single-writer (interleaved
+    writers truncate each other's ``.part``, see registry.py). A blown deadline
+    is a clean skip. On a lock-less filesystem a sha256-verified body runs
+    unlocked with one warning; a checksum-less one (``unlocked_ok=False``) fails.
     """
+    # The holder's part file sits beside dl_path, or beside the staging
+    # path a URL job promotes from
+    on_disk = partial(downloaded_bytes, stream_dest or dl_path, size)
 
     def run(tracker: Any) -> None:
-        from filelock import FileLock, Timeout
+        from filelock import FileLock
 
         # fallback_to_soft would leave a stale marker on lock-less
         # filesystems that blocks every later build (see git.py)
         lock = FileLock(lock_path, fallback_to_soft=False)
-        deadline = time.monotonic() + _DOWNLOAD_LOCK_TIMEOUT
-        while True:
-            try:
-                lock.acquire(timeout=_URI_LOCK_POLL)
-                break
-            except Timeout:
-                tracker(0)  # raises when the batch is cancelled
-                if time.monotonic() >= deadline:
-                    # Another process is fetching this same file; its copy
-                    # is what the build needs (a large framework archive
-                    # can hold the lock far longer than this deadline)
-                    _LOGGER.debug("Leaving %s to its current downloader", dl_path.name)
-                    return
-            except OSError as err:
-                if not unlocked_ok:
-                    # A body with no checksum to catch interleaved corruption
-                    raise
-                lock = None
-                _LOGGER.warning(
-                    "Could not lock %s (%s); downloading unlocked",
-                    dl_path.name,
-                    err,
-                )
-                break
+        try:
+            if not wait_for_download_lock(
+                lock, tracker, on_disk, dl_path.name, _DOWNLOAD_LOCK_TIMEOUT
+            ):
+                # The holder's copy is what the build needs (a large
+                # framework archive can outlast this deadline)
+                _LOGGER.debug("Leaving %s to its current downloader", dl_path.name)
+                return
+        except OSError as err:
+            if not unlocked_ok:
+                # A body with no checksum to catch interleaved corruption
+                raise
+            lock = None
+            _LOGGER.warning(
+                "Could not lock %s (%s); downloading unlocked",
+                dl_path.name,
+                err,
+            )
         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 +539,7 @@ def _registry_fetch_job(
         dl_path,
         f"{dl_path}.esphome.lock",
         resume_fetch_job(url, dl_path, sha256=checksum, size=size),
+        size,
     )
 
     def run(tracker: Any) -> None:
@@ -571,9 +571,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

esphome/platformio/registry.py

@@ -17,8 +17,10 @@ from esphome.framework_helpers import (
     archive_extract_all,
     download_from_mirrors,
     download_with_resume,
+    downloaded_bytes,
     rmdir,
     run_batch_downloads,
+    wait_for_download_lock,
 )
 from esphome.net_retry import fetch_with_retry, http_request
 
@@ -222,20 +224,31 @@ def prefetch_packages(
 
     def _fetch(entry: _PendingArchive, tracker: Callable[[int], None]) -> None:
         entry.dest.parent.mkdir(parents=True, exist_ok=True)
-        with FileLock(f"{entry.dest}.lock", fallback_to_soft=False):
-            # Marker re-check: a concurrent build may have installed (and
-            # deleted the archive of) this package while we waited;
-            # re-downloading would orphan a fresh copy in downloads_dir
-            # no branch: the thread tracer misses the skip edge; both
-            # arms of _already_installed are pinned directly
-            if not _already_installed(entry.dest):  # pragma: no branch
-                download_with_resume(
-                    entry.url,
-                    downloads_dir / f"{entry.name}-{entry.version}",
-                    sha256=entry.sha256,
-                    size=entry.size,
-                    progress=tracker,
-                )
+        archive = downloads_dir / f"{entry.name}-{entry.version}"
+
+        def on_disk() -> int:
+            if done := downloaded_bytes(archive, entry.size):
+                return done
+            # The holder deletes the archive once it has installed it
+            return entry.size if _already_installed(entry.dest) else 0
+
+        lock = FileLock(f"{entry.dest}.lock", fallback_to_soft=False)
+        wait_for_download_lock(lock, tracker, on_disk, entry.name)
+        try:
+            if _already_installed(entry.dest):
+                # A concurrent build installed it while we waited; a
+                # re-download would orphan a fresh copy in downloads_dir
+                tracker(entry.size)
+                return
+            download_with_resume(
+                entry.url,
+                archive,
+                sha256=entry.sha256,
+                size=entry.size,
+                progress=tracker,
+            )
+        finally:
+            lock.release()
 
     failures = run_batch_downloads(
         "Downloading packages",

tests/unit_tests/conftest.py

@@ -9,7 +9,7 @@ not be part of a unit test suite.
 
 """
 
-from collections.abc import Generator
+from collections.abc import Callable, Generator
 import os
 from pathlib import Path
 import sys
@@ -137,3 +137,34 @@ def mock_get_component() -> Generator[Mock, None, None]:
     """Mock get_component for config module."""
     with patch("esphome.config.get_component") as mock:
         yield mock
+
+
+@pytest.fixture
+def held_lock() -> Callable[..., Callable[..., None]]:
+    """Factory for a ``FileLock.acquire`` fake held by another downloader.
+
+    Each poll writes the next chunk to ``part`` and raises ``Timeout``; when
+    the chunks run out the part is removed, ``land()`` runs, and the acquire
+    succeeds (also for any later job, so ``land`` must be idempotent).
+    """
+    from filelock import Timeout
+
+    def make(
+        part: Path, chunks: list[bytes], land: Callable[[], None]
+    ) -> Callable[..., None]:
+        polls = iter(chunks)
+
+        def acquire(*args, **kwargs) -> None:
+            try:
+                chunk = next(polls)
+            except StopIteration:
+                part.unlink(missing_ok=True)
+                land()
+                return
+            part.parent.mkdir(parents=True, exist_ok=True)
+            part.write_bytes(chunk)
+            raise Timeout("held")
+
+        return acquire
+
+    return make

tests/unit_tests/test_framework_helpers.py

@@ -2353,3 +2353,18 @@ 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:
+    """Landed file: size; part file: its bytes, capped at size; nothing: 0."""
+    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
+    part.unlink()
+    assert framework_helpers.downloaded_bytes(dest) == 4

tests/unit_tests/test_platformio_prefetch.py

@@ -454,11 +454,14 @@ 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:
-    """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."""
+@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 is another process's download; skip
+    cleanly, polling the tracker with what the holder has staged so far."""
     dl_path = tmp_path / "archive"
+    (tmp_path / "archive.prefetch.part").write_bytes(staged)
     ticks: list[int] = []
     with (
         patch("esphome.framework_helpers.download_with_resume") as mock_download,
@@ -467,10 +470,60 @@ 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()
 
 
+@pytest.mark.parametrize(
+    ("job", "part_name", "chunks", "expected"),
+    [
+        (
+            lambda dl_path: pf._registry_fetch_job(
+                MagicMock(), "https://x/a.tar.gz", dl_path, "ab" * 32, 4
+            ),
+            "archive.part",
+            [b"a", b"abc"],
+            [1, 3, 4],
+        ),
+        (
+            lambda dl_path: pf._uri_fetch_job(
+                MagicMock(), "https://x/a.zip", dl_path, 4
+            ),
+            "archive.prefetch.part",
+            [b"ab"],
+            [2, 4],
+        ),
+    ],
+    ids=["registry", "uri"],
+)
+def test_lock_wait_reports_the_holders_progress(
+    tmp_path: Path,
+    caplog: pytest.LogCaptureFixture,
+    held_lock,
+    job,
+    part_name: str,
+    chunks: list[bytes],
+    expected: list[int],
+) -> None:
+    """A waiting job reports the holder's part file (the staging one for a
+    URL job), then the full size once the holder lands the archive."""
+    dl_path = tmp_path / "archive"
+    ticks: list[int] = []
+    acquire = held_lock(
+        tmp_path / part_name, chunks, lambda: dl_path.write_bytes(b"abcd")
+    )
+    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),
+    ):
+        job(dl_path)(ticks.append)
+    mock_download.assert_not_called()
+    assert ticks == expected
+    assert caplog.text.count("Waiting for another process downloading archive") == 1
+
+
 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."""

tests/unit_tests/test_platformio_registry.py

@@ -540,16 +540,13 @@ def test_prefetch_packages_skips_freshly_installed_dest(tmp_path: Path) -> None:
     dest = tmp_path / "a"
     dest.mkdir()
 
-    from contextlib import contextmanager
-
-    @contextmanager
-    def marker_appears_under_lock(path, **kwargs):
+    def marker_appears_under_lock(*args, **kwargs):
         # Simulates the concurrent build finishing while we waited
         (dest / ".esphome_extracted").touch()
-        yield
 
     with (
-        patch("filelock.FileLock", side_effect=marker_appears_under_lock),
+        patch("filelock.FileLock.acquire", side_effect=marker_appears_under_lock),
+        patch("filelock.FileLock.release"),
         patch.object(registry, "download_with_resume") as mock_download,
         patch.object(
             registry, "registry_download", side_effect=_resolve_for({"a": 10})
@@ -559,6 +556,42 @@ def test_prefetch_packages_skips_freshly_installed_dest(tmp_path: Path) -> None:
     mock_download.assert_not_called()
 
 
+def test_prefetch_packages_waits_with_the_holders_progress(
+    tmp_path: Path, held_lock
+) -> None:
+    """A worker parked on another build's lock reports that build's part
+    file, then the full size once the marker appears."""
+    dest = tmp_path / "a"
+    dest.mkdir()
+    ticks: list[int] = []
+    acquire = held_lock(
+        tmp_path / "dl" / "a-1.0.part",
+        [b"abc"],
+        (dest / ".esphome_extracted").touch,
+    )
+
+    def fake_batch(header, jobs):
+        for _name, _size, fetch in jobs:
+            fetch(ticks.append)
+        return []
+
+    with (
+        patch("filelock.FileLock.acquire", side_effect=acquire),
+        patch("filelock.FileLock.release"),
+        patch.object(registry, "run_batch_downloads", side_effect=fake_batch),
+        patch.object(registry, "download_with_resume") as mock_download,
+        patch.object(
+            registry, "registry_download", side_effect=_resolve_for({"a": 10, "b": 5})
+        ),
+    ):
+        registry.prefetch_packages(
+            [("a", "1.0", dest, []), ("b", "2.0", tmp_path / "b", [])],
+            tmp_path / "dl",
+        )
+    assert ticks == [3, 10]
+    mock_download.assert_called_once()
+
+
 def test_already_installed_probe(tmp_path: Path) -> None:
     """Both arms of the marker probe the prefetch worker keys on."""
     dest = tmp_path / "pkg"