diff --git a/esphome/framework_helpers.py b/esphome/framework_helpers.py index ee50723a39..d185320f3b 100644 --- a/esphome/framework_helpers.py +++ b/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,15 +910,46 @@ 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 +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: - return min(_part_path(dest).stat().st_size, size) + done = _part_path(dest).stat().st_size except OSError: - return 0 + 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: @@ -1330,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 diff --git a/esphome/platformio/prefetch.py b/esphome/platformio/prefetch.py index 41df598bd7..5cd31bf284 100644 --- a/esphome/platformio/prefetch.py +++ b/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 @@ -39,6 +40,7 @@ from esphome.framework_helpers import ( 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 @@ -69,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() @@ -467,59 +466,42 @@ def _serialized_fetch_job( lock_path: str, body: Any, size: int, - stream_dest: Path, + 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). 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. + """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 - waiting = False - while True: - try: - lock.acquire(timeout=_URI_LOCK_POLL) - break - except Timeout: - 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 - # 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(): tracker(size) # another process finished it while we waited @@ -558,7 +540,6 @@ def _registry_fetch_job( f"{dl_path}.esphome.lock", resume_fetch_job(url, dl_path, sha256=checksum, size=size), size, - dl_path, ) def run(tracker: Any) -> None: diff --git a/esphome/platformio/registry.py b/esphome/platformio/registry.py index 9538a28ff4..c0dfac8874 100644 --- a/esphome/platformio/registry.py +++ b/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", diff --git a/tests/unit_tests/conftest.py b/tests/unit_tests/conftest.py index 9de8f715ef..4835bf7289 100644 --- a/tests/unit_tests/conftest.py +++ b/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 diff --git a/tests/unit_tests/test_framework_helpers.py b/tests/unit_tests/test_framework_helpers.py index b5d535330c..d77f543cb1 100644 --- a/tests/unit_tests/test_framework_helpers.py +++ b/tests/unit_tests/test_framework_helpers.py @@ -2356,8 +2356,7 @@ def test_discard_partial_download_logs_undeletable( 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.""" + """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" @@ -2367,3 +2366,5 @@ def test_downloaded_bytes_reports_what_is_on_disk(tmp_path: Path) -> None: 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 diff --git a/tests/unit_tests/test_platformio_prefetch.py b/tests/unit_tests/test_platformio_prefetch.py index c8e511b7c5..fe19fc957f 100644 --- a/tests/unit_tests/test_platformio_prefetch.py +++ b/tests/unit_tests/test_platformio_prefetch.py @@ -458,13 +458,10 @@ def test_uri_fetch_job_waits_out_a_briefly_held_lock(tmp_path: Path) -> None: 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, and - it reports what the holder has staged so far.""" + """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" - if staged: - (tmp_path / "archive.prefetch.part").write_bytes(staged) + (tmp_path / "archive.prefetch.part").write_bytes(staged) ticks: list[int] = [] with ( patch("esphome.framework_helpers.download_with_resume") as mock_download, @@ -477,67 +474,54 @@ def test_lock_deadline_leaves_download_to_the_holder( 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 + tmp_path: Path, + caplog: pytest.LogCaptureFixture, + held_lock, + job, + part_name: str, + chunks: list[bytes], + expected: list[int], ) -> 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() + """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" - 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") - + 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), ): - pf._registry_fetch_job(manager, "https://x/a.tar.gz", dl_path, "ab" * 32, 4)( - ticks.append - ) + job(dl_path)(ticks.append) mock_download.assert_not_called() - assert ticks == [1, 3, 4] + assert ticks == expected 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: diff --git a/tests/unit_tests/test_platformio_registry.py b/tests/unit_tests/test_platformio_registry.py index 6ba8691c4e..e36d2be21a 100644 --- a/tests/unit_tests/test_platformio_registry.py +++ b/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"