Prefer the landed archive while waiting, cap credited bytes, and give the registry wait a deadline

esphome/framework_helpers.py

@@ -911,45 +911,58 @@ def _part_path(dest: Path) -> Path:
 
 
 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
+    """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,
-    timeout: float | None = None,
-) -> bool:
+) -> None:
     """Acquire ``lock``, reporting ``on_disk()`` to ``tracker`` each poll so the
-    bar follows the holder's download. False once ``timeout`` seconds pass."""
+    bar follows the holder's download. Raises filelock's ``Timeout`` once
+    ``DOWNLOAD_LOCK_TIMEOUT`` seconds pass."""
     from filelock import Timeout
 
-    deadline = None if timeout is None else time.monotonic() + timeout
+    deadline = time.monotonic() + DOWNLOAD_LOCK_TIMEOUT
     waiting = False
     while True:
         try:
             lock.acquire(timeout=_DOWNLOAD_LOCK_POLL)
-            return True
+            return
         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
+            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:

esphome/platformio/prefetch.py

@@ -19,7 +19,6 @@ 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
@@ -34,6 +33,7 @@ import time
 from typing import Any, NamedTuple
 
 from esphome.framework_helpers import (
+    DownloadLockUnavailable,
     content_length,
     discard_partial_download,
     downloaded_bytes,
@@ -64,9 +64,6 @@ _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
@@ -474,25 +471,29 @@ def _serialized_fetch_job(
     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 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
+        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)
         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:
+            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

esphome/platformio/registry.py

@@ -189,7 +189,7 @@ 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()
@@ -233,7 +233,13 @@ def prefetch_packages(
             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:
+            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

tests/unit_tests/conftest.py

@@ -143,14 +143,17 @@ def mock_get_component() -> Generator[Mock, None, None]:
 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).
+    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], land: Callable[[], None]
+        part: Path,
+        chunks: list[bytes | Callable[[], None]],
+        land: Callable[[], None],
     ) -> Callable[..., None]:
         polls = iter(chunks)
 
@@ -161,8 +164,11 @@ def held_lock() -> Callable[..., Callable[..., None]]:
                 part.unlink(missing_ok=True)
                 land()
                 return
-            part.parent.mkdir(parents=True, exist_ok=True)
-            part.write_bytes(chunk)
+            if callable(chunk):
+                chunk()
+            else:
+                part.parent.mkdir(parents=True, exist_ok=True)
+                part.write_bytes(chunk)
             raise Timeout("held")
 
         return acquire

tests/unit_tests/test_framework_helpers.py

@@ -2356,7 +2356,7 @@ def test_discard_partial_download_logs_undeletable(
 
 
 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."""
+    """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"
@@ -2364,7 +2364,9 @@ def test_downloaded_bytes_reports_what_is_on_disk(tmp_path: Path) -> None:
     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
+    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

@@ -466,7 +466,7 @@ def test_lock_deadline_leaves_download_to_the_holder(
     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()
@@ -524,6 +524,26 @@ def test_lock_wait_reports_the_holders_progress(
     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."""
@@ -532,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
@@ -564,9 +565,16 @@ def test_prefetch_packages_waits_with_the_holders_progress(
     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(
-        tmp_path / "dl" / "a-1.0.part",
-        [b"abc"],
+        part,
+        [lambda: None, b"abc", installed_and_pruned],
         (dest / ".esphome_extracted").touch,
     )
 
@@ -588,10 +596,30 @@ def test_prefetch_packages_waits_with_the_holders_progress(
             [("a", "1.0", dest, []), ("b", "2.0", tmp_path / "b", [])],
             tmp_path / "dl",
         )
-    assert ticks == [3, 10]
+    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"