From 062c542c3692d9a25f69ca7e40a145e35476b3b8 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Sun, 23 Aug 2026 14:42:00 -0500 Subject: [PATCH] Give worker log records their own line, harden the cancellation sentinel A logging filter installed for the batch's duration ends a partial bar row (under the bar's lock) before any record is emitted, so a mirror retry warning from a worker is no longer overwritten by the next CR frame. The Ctrl-C sentinel derives from BaseException so a broad except Exception in the download layers cannot convert an abort into a retry, an abandoned job is reported like a failure instead of reading as a completed download, and the mirror-sweep backoff sleeps through the cancellation tick so a Ctrl-C during a retry window aborts at the boundary instead of after it. The docstring states the parked-socket worst case. --- esphome/framework_helpers.py | 77 +++++++++++++++++++--- tests/unit_tests/test_framework_helpers.py | 67 ++++++++++++++++++- 2 files changed, 133 insertions(+), 11 deletions(-) diff --git a/esphome/framework_helpers.py b/esphome/framework_helpers.py index c84b6ce503..29bd9313b4 100644 --- a/esphome/framework_helpers.py +++ b/esphome/framework_helpers.py @@ -1,8 +1,8 @@ """Generic toolchain installation helpers shared across framework implementations.""" -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Iterator from concurrent.futures import ThreadPoolExecutor -from contextlib import ExitStack +from contextlib import ExitStack, contextmanager import hashlib import io import json @@ -752,7 +752,8 @@ def run_batch_downloads( Each ``fetch(tracker)`` reports absolute byte counts; the bar total is the sum of the sizes. Failures are returned after the bar is done so warnings never land on its row. Ctrl-C drops queued jobs and aborts - in-flight ones at their next tick; resumable destinations + in-flight ones at their next progress tick or backoff boundary (a + parked socket read defers that by its timeout); resumable destinations (``download_with_resume``) keep their fetched ``.part`` bytes. ``jobs`` must be non-empty. """ @@ -771,8 +772,11 @@ def run_batch_downloads( try: fetch(checked) - except _BatchDownloadCancelled: + except _BatchDownloadCancelled as err: + # Reported like a failure: an abandoned job must never read as + # a completed download if a caller sees the list after Ctrl-C tracker(0) + return (name, err) except Exception as err: # noqa: BLE001 # pylint: disable=broad-exception-caught tracker(0) return (name, err) @@ -780,8 +784,9 @@ def run_batch_downloads( ex = ThreadPoolExecutor(max_workers=max_workers) try: - futures = [ex.submit(_run, name, fetch) for name, _, fetch in jobs] - return [failure for f in futures if (failure := f.result()) is not None] + with progress.logging_guard(): + futures = [ex.submit(_run, name, fetch) for name, _, fetch in jobs] + return [failure for f in futures if (failure := f.result()) is not None] except BaseException: # Without this the non-daemon workers download to completion before # the interpreter can exit, making Ctrl-C ineffective for minutes @@ -792,8 +797,12 @@ def run_batch_downloads( progress.done() -class _BatchDownloadCancelled(Exception): - """Raised inside a download job to abandon it after Ctrl-C.""" +class _BatchDownloadCancelled(BaseException): + """Raised inside a download job to abandon it after Ctrl-C. + + BaseException, like KeyboardInterrupt: a broad ``except Exception`` in + the download layers must not convert an abort into a retry. + """ class _BatchDownloadProgress: @@ -828,6 +837,56 @@ class _BatchDownloadProgress: if self._bar is not None: self._bar.done() + @contextmanager + def logging_guard(self) -> Iterator[None]: + r"""End a partial bar row before any log record while active. + + Worker warnings (mirror retries) share stderr with the bar's \r + frames; without this the record lands mid-row and the next frame + overwrites it. A handler-level filter runs just before emit, so + only a tiny window remains for a concurrent frame. + """ + the_bar = self._bar + if the_bar is None: + yield + return + lock = self._lock + + class _EndRow(logging.Filter): + def filter(self, record: logging.LogRecord) -> bool: + with lock: + if ( + the_bar.last_progress is not None + and the_bar.last_progress != 100 + ): + the_bar.done() + # Force the next tick to redraw the frame + the_bar.last_progress = None + return True + + end_row = _EndRow() + handlers = logging.getLogger().handlers + for handler in handlers: + handler.addFilter(end_row) + try: + yield + finally: + for handler in handlers: + handler.removeFilter(end_row) + + +def _cancellable_sleep( + delay: float, progress: Callable[[int], None] | None, done: int +) -> None: + """Backoff sleep that still observes a batch cancellation tick.""" + if progress is None: + time.sleep(delay) + return + end = time.monotonic() + delay + while (remaining := end - time.monotonic()) > 0: + progress(done) # raises when the batch was cancelled + time.sleep(min(0.5, remaining)) + def download_with_resume( url: str, @@ -1270,7 +1329,7 @@ def download_from_mirrors( sweep + 1, _MIRROR_SWEEP_ATTEMPTS, ) - time.sleep(delay) + _cancellable_sleep(delay, progress, 0) # 4. Report every attempted URL if all mirrors failed. failures spans # all sweeps (deduplicated by URL and reason), so neither an early diff --git a/tests/unit_tests/test_framework_helpers.py b/tests/unit_tests/test_framework_helpers.py index 772257ce14..28e8f2f106 100644 --- a/tests/unit_tests/test_framework_helpers.py +++ b/tests/unit_tests/test_framework_helpers.py @@ -1126,13 +1126,13 @@ class TestDownloadWithResume: seen: list[int] = [] with ( patch("requests.get", return_value=resp), - patch("esphome.framework_helpers.ProgressBar") as bar, + patch("esphome.framework_helpers.ProgressBar") as bar_cls, ): download_with_resume( "https://example.com/t", dest, size=7, progress=seen.append ) assert seen == [0, 4, 7, 7] - bar.assert_not_called() + bar_cls.assert_not_called() def test_progress_callback_seeds_with_resume_offset(self, tmp_path: Path) -> None: dest = tmp_path / "tool.tar.gz" @@ -1190,6 +1190,69 @@ def test_run_batch_downloads_ctrl_c_aborts_in_flight_jobs() -> None: assert len(ticks) < 500 +def test_cancellation_escapes_broad_except_in_fetch() -> None: + """A fetch that wraps its work in except Exception cannot swallow the + Ctrl-C sentinel (it is a BaseException).""" + from esphome.framework_helpers import _BatchDownloadCancelled + + started = threading.Event() + swallowed = [] + + def interrupter(tracker) -> None: + started.wait(5) + raise KeyboardInterrupt + + def greedy_fetch(tracker) -> None: + started.set() + try: + for i in range(500): + tracker(i) + time.sleep(0.01) + except Exception as err: # noqa: BLE001 # pylint: disable=broad-exception-caught + swallowed.append(err) + + t0 = time.monotonic() + with pytest.raises(KeyboardInterrupt): + run_batch_downloads( + "Downloading", + [("boom", 0, interrupter), ("greedy", 0, greedy_fetch)], + max_workers=2, + ) + assert time.monotonic() - t0 < 3 + assert not swallowed + assert issubclass(_BatchDownloadCancelled, BaseException) + assert not issubclass(_BatchDownloadCancelled, Exception) + + +def test_logging_guard_ends_the_bar_row_before_a_record() -> None: + r"""A worker warning gets its own line instead of the bar's \r row.""" + stream = io.StringIO() + stream.isatty = lambda: True # type: ignore[method-assign] + with patch("esphome.helpers.sys.stderr", stream): + progress = _BatchDownloadProgress("Downloading", 10) + progress.tracker()(5) + with progress.logging_guard(): + logging.getLogger("esphome.test").warning("mirror retry") + # The partial 50% frame ended its line before the record was emitted + assert stream.getvalue().endswith("50% \n") + # And the next tick redraws the frame on a fresh row + progress.tracker()(2) + assert stream.getvalue().endswith("70% ") + + +def test_cancellable_sleep_aborts_at_the_tick() -> None: + """A backoff sleep observes the cancellation raise promptly.""" + from esphome.framework_helpers import _BatchDownloadCancelled, _cancellable_sleep + + def cancelled_tick(done: int) -> None: + raise _BatchDownloadCancelled + + t0 = time.monotonic() + with pytest.raises(_BatchDownloadCancelled): + _cancellable_sleep(30, cancelled_tick, 0) + assert time.monotonic() - t0 < 1 + + class Test_BatchDownloadProgress: def test_sums_trackers_into_one_bar(self) -> None: with patch("esphome.framework_helpers.ProgressBar") as bar_cls: