diff --git a/esphome/framework_helpers.py b/esphome/framework_helpers.py index 82bc0d3727..fc2a18a6ec 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,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 diff --git a/esphome/platformio/prefetch.py b/esphome/platformio/prefetch.py index 5097239065..17a06cb9c1 100644 --- a/esphome/platformio/prefetch.py +++ b/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 diff --git a/esphome/platformio/registry.py b/esphome/platformio/registry.py index 9538a28ff4..75df82da0e 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 @@ -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( diff --git a/tests/unit_tests/conftest.py b/tests/unit_tests/conftest.py index 9de8f715ef..ad9c0bb11f 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,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 diff --git a/tests/unit_tests/test_framework_helpers.py b/tests/unit_tests/test_framework_helpers.py index fcc5572f51..22b34c9df5 100644 --- a/tests/unit_tests/test_framework_helpers.py +++ b/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 diff --git a/tests/unit_tests/test_platformio_prefetch.py b/tests/unit_tests/test_platformio_prefetch.py index fb79885736..77490fd861 100644 --- a/tests/unit_tests/test_platformio_prefetch.py +++ b/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 diff --git a/tests/unit_tests/test_platformio_registry.py b/tests/unit_tests/test_platformio_registry.py index 6ba8691c4e..9d5f6c4ce5 100644 --- a/tests/unit_tests/test_platformio_registry.py +++ b/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"