Derive the registry archive path in one place

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,61 @@ 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`` while streaming), capped at ``size``."""
+    done = 0
+    for candidate in (_part_path(dest), dest):
+        try:
+            done = candidate.stat().st_size
+            break
+        except FileNotFoundError:
+            continue
+    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
+
+# Waiting on another process's download; past this the caller leaves the
+# file to its holder (the later sequential install waits on the same lock)
+DOWNLOAD_LOCK_TIMEOUT = 60
+
+
+class DownloadLockUnavailable(OSError):
+    """The lock file cannot be used at all (a lock-less filesystem)."""
+
+
+def wait_for_download_lock(
+    lock: "FileLock",
+    tracker: Callable[[int], None],
+    on_disk: Callable[[], int],
+    name: str,
+) -> None:
+    """Acquire ``lock``, reporting ``on_disk()`` to ``tracker`` each poll so the
+    bar follows the holder's download. Raises filelock's ``Timeout`` once
+    ``DOWNLOAD_LOCK_TIMEOUT`` seconds pass."""
+    from filelock import Timeout
+
+    deadline = time.monotonic() + DOWNLOAD_LOCK_TIMEOUT
+    waiting = False
+    while True:
+        try:
+            lock.acquire(timeout=_DOWNLOAD_LOCK_POLL)
+            return
+        except Timeout:
+            pass
+        except OSError as err:
+            # Distinct from an OSError out of on_disk(), which must not
+            # read as "locks unsupported"
+            raise DownloadLockUnavailable(*err.args) from err
+        if not waiting:
+            waiting = True
+            _LOGGER.info("Waiting for another process downloading %s", name)
+        tracker(on_disk())  # raises when the batch is cancelled
+        if time.monotonic() >= deadline:
+            raise Timeout(lock.lock_file)
+
+
 def discard_partial_download(dest: Path) -> None:
     """Remove ``dest`` and the resume sidecars of an abandoned download."""
     part = _part_path(dest)
@@ -1319,10 +1375,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

@@ -33,11 +33,14 @@ import time
 from typing import Any, NamedTuple
 
 from esphome.framework_helpers import (
+    DownloadLockUnavailable,
     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
@@ -61,16 +64,10 @@ _RESOLVE_WORKERS = 8
 # A hung child must not block the build; downloads resume on the next run
 _PREFETCH_TIMEOUT = 20 * 60
 
-# Waiting on another process's URL download; past this, leave it to pio
-_DOWNLOAD_LOCK_TIMEOUT = 60
-
 # Child exit for a handled, already-warned failure; 1 would collide with
 # 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 +459,54 @@ 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.
     """
 
+    def on_disk() -> int:
+        # A URL job's holder streams beside the staging path until it
+        # promotes; after that only dl_path is left
+        done = downloaded_bytes(dl_path, size)
+        if not done and stream_dest is not None:
+            done = downloaded_bytes(stream_dest, size)
+        return done
+
     def run(tracker: Any) -> None:
         from filelock import FileLock, Timeout
 
         # 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:
+            wait_for_download_lock(lock, tracker, on_disk, dl_path.name)
+        except 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 DownloadLockUnavailable 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 +540,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 +572,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
 
@@ -164,11 +166,17 @@ class _PendingArchive(NamedTuple):
     name: str
     version: str
     dest: Path
+    archive: Path
     url: str
     sha256: str
     size: int
 
 
+def _archive_path(downloads_dir: Path, name: str, version: str) -> Path:
+    """The one archive path the prefetch and the sequential install share."""
+    return downloads_dir / f"{name}-{version}"
+
+
 def _already_installed(dest: Path) -> bool:
     """Whether ``dest`` holds a completed install (extraction marker)."""
     return (dest / ".esphome_extracted").is_file()
@@ -187,18 +195,18 @@ def prefetch_packages(
     lock as ``install_package``: the archive's ``.part`` file is shared, and
     two concurrent writers would truncate each other's bytes.
     """
-    from filelock import FileLock
+    from filelock import FileLock, Timeout
 
     pending: list[_PendingArchive] = []
-    seen: set[str] = set()
+    seen: set[Path] = set()
     for name, version, dest, mirrors in packages:
         if mirrors or (dest / ".esphome_extracted").is_file():
             continue
-        archive_name = f"{name}-{version}"
-        if archive_name in seen:
+        archive = _archive_path(downloads_dir, name, version)
+        if archive in seen:
             # A duplicate entry would race itself between two workers
             continue
-        seen.add(archive_name)
+        seen.add(archive)
         try:
             url, sha256, size = registry_download(name, version)
         except EsphomeError as err:
@@ -207,10 +215,9 @@ def prefetch_packages(
             continue
         if not size:
             continue
-        archive = downloads_dir / archive_name
         if archive.is_file() and archive.stat().st_size == size:
             continue
-        pending.append(_PendingArchive(name, version, dest, url, sha256, size))
+        pending.append(_PendingArchive(name, version, dest, archive, url, sha256, size))
     if len(pending) < 2:
         return
     downloads_dir.mkdir(parents=True, exist_ok=True)
@@ -222,20 +229,36 @@ 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,
-                )
+
+        def on_disk() -> int:
+            if done := downloaded_bytes(entry.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)
+        try:
+            wait_for_download_lock(lock, tracker, on_disk, entry.name)
+        except Timeout:
+            # install_package waits on this same lock and verifies the
+            # holder's copy
+            _LOGGER.debug("Leaving %s to its current downloader", entry.name)
+            return
+        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,
+                entry.archive,
+                sha256=entry.sha256,
+                size=entry.size,
+                progress=tracker,
+            )
+        finally:
+            lock.release()
 
     failures = run_batch_downloads(
         "Downloading packages",
@@ -288,7 +311,7 @@ def install_package(
         rmdir(dest, msg=f"Clean up incomplete {name} install")
         # Persistent location so an interrupted download resumes across runs.
         downloads_dir.mkdir(parents=True, exist_ok=True)
-        archive = downloads_dir / f"{name}-{version}"
+        archive = _archive_path(downloads_dir, name, version)
         _LOGGER.info("Downloading %s %s ...", name, version)
         if mirrors:
             _LOGGER.warning(

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,40 @@ 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`` (or runs it, for a callable)
+    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 | Callable[[], None]],
+        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
+            if callable(chunk):
+                chunk()
+            else:
+                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,20 @@ 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:
+    """Part file first, then the landed file, both capped at size; else 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
+    part.unlink()
+    dest.write_bytes(b"abc")
+    assert framework_helpers.downloaded_bytes(dest, 4) == 3
+    assert framework_helpers.downloaded_bytes(dest) == 3
+    dest.write_bytes(b"abcdef")
+    assert framework_helpers.downloaded_bytes(dest, 4) == 4

tests/unit_tests/test_platformio_prefetch.py

@@ -454,23 +454,96 @@ 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,
         patch("filelock.FileLock.acquire", side_effect=Timeout("held")),
-        patch.object(pf, "_DOWNLOAD_LOCK_TIMEOUT", 0),
+        patch("esphome.framework_helpers.DOWNLOAD_LOCK_TIMEOUT", 0),
     ):
         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_uri_lock_wait_prefers_the_landed_archive(tmp_path: Path, held_lock) -> None:
+    """Between the holder's promotion rename and its release the staging
+    part is gone; the landed cache file is credited instead of 0."""
+    dl_path = tmp_path / "archive"
+    ticks: list[int] = []
+    acquire = held_lock(
+        tmp_path / "archive.prefetch.part",
+        [b"ab", lambda: dl_path.write_bytes(b"abcd")],
+        lambda: None,
+    )
+    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, 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."""
@@ -479,7 +552,7 @@ def test_registry_lock_deadline_skips_registration(tmp_path: Path) -> None:
     with (
         patch("esphome.framework_helpers.download_with_resume") as mock_download,
         patch("filelock.FileLock.acquire", side_effect=Timeout("held")),
-        patch.object(pf, "_DOWNLOAD_LOCK_TIMEOUT", 0),
+        patch("esphome.framework_helpers.DOWNLOAD_LOCK_TIMEOUT", 0),
     ):
         pf._registry_fetch_job(manager, "https://x/a.tar.gz", dl_path, "ab" * 32, 4)(
             lambda done: None

tests/unit_tests/test_platformio_registry.py

@@ -8,6 +8,7 @@ import os
 from pathlib import Path
 from unittest.mock import MagicMock, patch
 
+from filelock import Timeout
 import pytest
 
 from esphome.core import EsphomeError
@@ -540,16 +541,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 +557,69 @@ 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] = []
+    part = tmp_path / "dl" / "a-1.0.part"
+
+    def installed_and_pruned() -> None:
+        # install_package touches the marker, then unlinks the archive
+        (dest / ".esphome_extracted").touch()
+        part.unlink()
+
+    acquire = held_lock(
+        part,
+        [lambda: None, b"abc", installed_and_pruned],
+        (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 == [0, 3, 10, 10]
+    mock_download.assert_called_once()
+
+
+def test_prefetch_packages_leaves_a_long_held_lock_to_its_holder(
+    tmp_path: Path,
+) -> None:
+    """Past the deadline the worker skips; install_package waits on the same
+    lock later and verifies whatever the holder produced."""
+    with (
+        patch("filelock.FileLock.acquire", side_effect=Timeout("held")),
+        patch("esphome.framework_helpers.DOWNLOAD_LOCK_TIMEOUT", 0),
+        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", tmp_path / "a", []), ("b", "2.0", tmp_path / "b", [])],
+            tmp_path / "dl",
+        )
+    mock_download.assert_not_called()
+
+
 def test_already_installed_probe(tmp_path: Path) -> None:
     """Both arms of the marker probe the prefetch worker keys on."""
     dest = tmp_path / "pkg"