diff --git a/superset/utils/screenshot_utils.py b/superset/utils/screenshot_utils.py index 868b8e731d55..82a06fd66c0c 100644 --- a/superset/utils/screenshot_utils.py +++ b/superset/utils/screenshot_utils.py @@ -87,6 +87,29 @@ def resolve_screenshot_task_budget_seconds( return None +# Fallback wall-clock budget, in seconds, for the entire tiled-screenshot +# operation (element lookup plus all per-tile readiness/animation waits +# combined), used when resolve_screenshot_task_budget_seconds() returns None +# (no Celery task context -- e.g. synchronous thumbnail generation -- or no +# usable task limit). The non-tiled readiness path treats None as "keep the +# configured SCREENSHOT_LOAD_WAIT" because it makes exactly one bounded wait; +# the tiled path cannot, because its per-tile waits accumulate: with N tiles, +# an uncapped load_wait allows N * load_wait of total wall-clock time, so the +# operation still needs one fixed total ceiling. Sized against the longest +# Celery hard task_time_limit observed in production for report execution +# (1740s), minus the same 300s cleanup margin the runtime derivation reserves +# for combining tiles, building the PDF, and delivering the notification. +TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS = 1440 # 1740s limit - 300s margin + + +class ScreenshotTaskBudgetExceededError(RuntimeError): + """Raised when no safe task budget remains before screenshot capture.""" + + +class TiledScreenshotBudgetExceededError(ScreenshotTaskBudgetExceededError): + """Raised when the tiled-screenshot time budget runs out mid-capture.""" + + try: from playwright.sync_api import TimeoutError as PlaywrightTimeout except ImportError: @@ -251,7 +274,7 @@ def combine_screenshot_tiles(screenshot_tiles: list[bytes]) -> bytes: return screenshot_tiles[0] -def take_tiled_screenshot( +def take_tiled_screenshot( # noqa: C901 page: "Page", element_name: str, tile_height: int, @@ -274,6 +297,12 @@ def take_tiled_screenshot( Returns: Combined screenshot bytes or None if failed + + Raises: + TiledScreenshotBudgetExceededError: If the total time budget for the + tiled-screenshot operation runs out before every tile has been + verifiably captured. Callers must treat this as a hard failure + rather than fall back to an unchecked/partial screenshot. """ context_suffix = f" [{log_context}]" if log_context else "" # Set right before re-raising the per-tile readiness timeout below, and @@ -286,6 +315,15 @@ def take_tiled_screenshot( # match `except PlaywrightTimeout` and incorrectly propagate instead of # degrading to `None` like every other unexpected error in this function. readiness_timeout = False + # Cap the whole tiled operation against the running Celery task's own + # time limit, using the same runtime derivation as the non-tiled + # readiness wait (#42253/#42427). Unlike that path, a None budget does + # not mean "keep the configured timeout": per-tile waits accumulate, so + # the operation falls back to a fixed total ceiling instead. + wait_budget_seconds = resolve_screenshot_task_budget_seconds(log_context) + if wait_budget_seconds is None: + wait_budget_seconds = float(TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS) + start_time = time.monotonic() try: # Get the target element element = page.locator(f".{element_name}") @@ -320,9 +358,44 @@ def take_tiled_screenshot( num_tiles = max(1, (dashboard_height + tile_height - 1) // tile_height) logger.info("Taking %s screenshot tiles", num_tiles) - screenshot_tiles = [] + screenshot_tiles: list[bytes] = [] + + def _raise_if_budget_exhausted(elapsed: float, remaining_budget: float) -> None: + if remaining_budget > 0: + return + # A customer-side chart-loading issue (a slow/hung dashboard), + # not a Superset system fault, so this is a WARNING rather + # than an ERROR -- consistent with #38130/#38441, which + # deliberately downgraded screenshot timeout logs the same way. + logger.warning( + "Tiled screenshot time budget exhausted on tile %s/%s: " + "%s/%s tiles captured so far, %.1fs elapsed of a %.1fs " + "budget. Aborting instead of capturing remaining tiles " + "unchecked.%s", + i + 1, + num_tiles, + len(screenshot_tiles), + num_tiles, + elapsed, + wait_budget_seconds, + context_suffix, + ) + raise TiledScreenshotBudgetExceededError( + f"Tiled screenshot budget of " + f"{wait_budget_seconds:.1f}s exhausted " + f"after {len(screenshot_tiles)}/{num_tiles} tiles" + ) for i in range(num_tiles): + # Check the time budget before starting this tile's readiness wait. + # If it's already exhausted, we can no longer verify this (or any + # later) tile is actually ready to capture -- fail loudly instead + # of silently snapshotting a spinner or blank chart, or running + # past the Celery task time limit and getting SIGKILLed. + elapsed = time.monotonic() - start_time + remaining_budget = wait_budget_seconds - elapsed + _raise_if_budget_exhausted(elapsed, remaining_budget) + # Calculate scroll position to show this tile's content scroll_y = dashboard_top + (i * tile_height) @@ -332,17 +405,31 @@ def take_tiled_screenshot( ) # Wait for scroll to settle and content to load page.wait_for_timeout(SCROLL_SETTLE_TIMEOUT_MS) + + # Recompute the remaining budget after the scroll-settle sleep -- + # which itself consumes real wall-clock time -- rather than + # reusing the value from before it, so the readiness-check + # timeout below is capped against a fresh number instead of a + # stale one that would let each tile overrun the budget by up + # to one settle interval. + tile_wait_start = time.monotonic() + elapsed = tile_wait_start - start_time + remaining_budget = wait_budget_seconds - elapsed + _raise_if_budget_exhausted(elapsed, remaining_budget) + # Wait for every chart holder visible in the current viewport to reach - # a terminal state (rendered chart or error/empty state). Only check + # a terminal state (rendered chart or error/empty state), capped at + # whatever remains of the total time budget so a slow dashboard + # degrades gracefully instead of exceeding it. Only check # viewport-visible chart holders to avoid blocking on virtualization # placeholders rendered for off-screen charts. A holder that hasn't # mounted anything yet does not satisfy this check -- unlike checking # for the absence of `.loading`, which passes vacuously in that case. - tile_wait_start = time.monotonic() + tile_load_wait = min(load_wait, remaining_budget) try: page.wait_for_function( CHART_HOLDERS_READY_JS, - timeout=load_wait * 1000, + timeout=tile_load_wait * 1000, ) except PlaywrightTimeout: elapsed = time.monotonic() - tile_wait_start @@ -354,14 +441,21 @@ def take_tiled_screenshot( # made the same call for the other screenshot timeout paths. logger.warning( "Timed out after %.2fs waiting for %s chart container(s) to " - "become ready on tile %s/%s (load_wait=%ss)%s; unready chart " - "holders (chart id, state): %s. Aborting tiled screenshot " - "rather than capturing a blank or partially-loaded tile.", + "become ready on tile %s/%s (waited %.1fs of a %ss requested " + "load_wait; %.1fs elapsed of a %.1fs total budget; %s/%s " + "tiles captured so far)%s; unready chart holders (chart id, " + "state): %s. Aborting tiled screenshot rather than capturing " + "a blank or partially-loaded tile.", elapsed, len(unready_chart_holders), i + 1, num_tiles, + tile_load_wait, load_wait, + time.monotonic() - start_time, + wait_budget_seconds, + len(screenshot_tiles), + num_tiles, context_suffix, unready_chart_holders, ) @@ -377,12 +471,36 @@ def take_tiled_screenshot( load_wait, context_suffix, ) + readiness_wait_elapsed = time.monotonic() - tile_wait_start # Wait for chart animations (e.g. ECharts) to finish after spinner clears. # The global animation wait before tiling only covers the first tile; - # subsequent tiles need their own wait after data loads. + # subsequent tiles need their own wait after data loads. Capped at + # whatever remains of the budget; unlike the readiness wait above this + # is cosmetic settling, not a readiness check, so we simply skip it + # (rather than raise) once the budget runs out. + animation_wait_elapsed = 0.0 if animation_wait > 0: - page.wait_for_timeout(animation_wait * 1000) + elapsed = time.monotonic() - start_time + remaining_budget = wait_budget_seconds - elapsed + tile_animation_wait = max(0, min(animation_wait, remaining_budget)) + if tile_animation_wait > 0: + animation_wait_start = time.monotonic() + page.wait_for_timeout(tile_animation_wait * 1000) + animation_wait_elapsed = time.monotonic() - animation_wait_start + + # Per-tile timing breakdown so slow dashboards can be profiled from + # logs alone. DEBUG rather than INFO: this fires once per tile, and + # large dashboards can have dozens of tiles per report run. + logger.debug( + "Tile %s/%s timing: %.2fs waiting for chart readiness, " + "%.2fs waiting for animations.%s", + i + 1, + num_tiles, + readiness_wait_elapsed, + animation_wait_elapsed, + context_suffix, + ) # Calculate what portion of the element we want to capture for this tile tile_start_in_element = i * tile_height @@ -431,6 +549,12 @@ def take_tiled_screenshot( return combined_screenshot + except TiledScreenshotBudgetExceededError: + # Budget exhaustion must fail cleanly, not be swallowed into the + # generic `return None` degradation below -- the raise carries the + # budget diagnostics to the caller, which fails the capture loudly + # (#42273) instead of receiving an anonymous empty result. + raise except Exception as e: if readiness_timeout: # Let the per-tile readiness timeout propagate so the caller diff --git a/superset/utils/webdriver.py b/superset/utils/webdriver.py index 75eeb1d2497b..19715a4c21f8 100644 --- a/superset/utils/webdriver.py +++ b/superset/utils/webdriver.py @@ -47,6 +47,7 @@ CHART_HOLDERS_READY_JS, FIND_CHART_HOLDER_STATES_JS, resolve_screenshot_task_budget_seconds, + ScreenshotTaskBudgetExceededError, take_tiled_screenshot, ) @@ -61,10 +62,6 @@ ) -class ScreenshotTaskBudgetExceededError(RuntimeError): - """Raised when no safe task budget remains before screenshot capture.""" - - if TYPE_CHECKING: from typing import Any diff --git a/tests/unit_tests/utils/test_screenshot_utils.py b/tests/unit_tests/utils/test_screenshot_utils.py index 438290015968..f5569e8a7419 100644 --- a/tests/unit_tests/utils/test_screenshot_utils.py +++ b/tests/unit_tests/utils/test_screenshot_utils.py @@ -25,8 +25,11 @@ combine_screenshot_tiles, resolve_screenshot_task_budget_seconds, SCREENSHOT_TASK_BUDGET_MAX_MARGIN_SECONDS, + ScreenshotTaskBudgetExceededError, SCROLL_SETTLE_TIMEOUT_MS, take_tiled_screenshot, + TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS, + TiledScreenshotBudgetExceededError, ) @@ -453,12 +456,17 @@ def test_per_tile_readiness_timeout_raises_and_skips_capture(self, mock_page): assert warning_args[2] == 1 # count of unready chart containers assert warning_args[3] == 1 # tile index assert warning_args[4] == 3 # total tiles - assert warning_args[5] == 30 # load_wait - assert warning_args[6] == "" # no log_context passed + assert warning_args[5] == 30 # tile_load_wait (uncapped: budget remains) + assert warning_args[6] == 30 # requested load_wait + assert isinstance(warning_args[7], float) # total elapsed vs budget + assert warning_args[8] == 1440 # total budget (fixed fallback) + assert warning_args[9] == 0 # tiles captured so far + assert warning_args[10] == 3 # total tiles + assert warning_args[11] == "" # no log_context passed # Diagnostic payload identifies chart id AND the state it's stuck in # (spinner mounted vs nothing mounted vs waiting-on-database) so a # slow query can be told apart from the virtualization race. - assert warning_args[7] == [{"chartId": "42", "state": "waiting_on_database"}] + assert warning_args[12] == [{"chartId": "42", "state": "waiting_on_database"}] def test_timeout_warning_includes_log_context(self, mock_page): """The log context (e.g. report execution id) is threaded through for @@ -484,7 +492,7 @@ def test_timeout_warning_includes_log_context(self, mock_page): ) warning_args = mock_logger.warning.call_args[0] - assert warning_args[6] == " [execution_id=abc-123]" + assert warning_args[11] == " [execution_id=abc-123]" def test_chart_holder_with_nothing_mounted_blocks_wait(self, mock_page): """Regression test for the vacuous-pass race (PR #39895). @@ -646,3 +654,311 @@ def test_animation_wait_default_is_zero(self): sig = inspect.signature(take_tiled_screenshot) assert sig.parameters["animation_wait"].default == 0 + + +class TestTileWaitBudget: + """The tiled operation's cumulative per-tile waits are capped by one + wall-clock budget derived from the running Celery task's own time limit + (resolve_screenshot_task_budget_seconds), falling back to a fixed total + ceiling outside Celery because per-tile waits accumulate.""" + + @pytest.fixture + def mock_page(self): + """Create a mock Playwright page object for a 3-tile (5000px) dashboard.""" + page = MagicMock() + element = MagicMock() + page.locator.return_value = element + page.evaluate.return_value = { + "height": 5000, + "top": 100, + "left": 50, + "width": 800, + } + page.screenshot.return_value = b"fake_screenshot_data" + return page + + class _FakeClock: + """Stateful monotonic() stand-in the test advances explicitly. + + Robust to how many times the code under test samples the clock per + tile (budget check, per-tile wait timing, animation budget) -- only + explicit advances move time forward. + """ + + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now + + def test_budget_error_is_task_budget_error_subclass(self): + """Callers can catch the whole budget-error family with the base + ScreenshotTaskBudgetExceededError type.""" + assert issubclass( + TiledScreenshotBudgetExceededError, ScreenshotTaskBudgetExceededError + ) + + def test_per_tile_wait_shrinks_as_budget_depletes(self, mock_page, monkeypatch): + """Each tile's readiness-wait timeout is capped at the remaining budget.""" + monkeypatch.setattr( + "superset.utils.screenshot_utils.TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS", # noqa: E501 + 1000, + ) + clock = self._FakeClock() + # Simulate slow tiles: the readiness wait itself consumes wall time, + # so each subsequent tile sees less remaining budget. + wait_durations = iter([950, 40, 5]) + + def slow_wait(*args, **kwargs): + clock.now += next(wait_durations) + + mock_page.wait_for_function.side_effect = slow_wait + + with patch("superset.utils.screenshot_utils.current_task", None): + with patch("superset.utils.screenshot_utils.time.monotonic", new=clock): + with patch("superset.utils.screenshot_utils.combine_screenshot_tiles"): + result = take_tiled_screenshot( + mock_page, "dashboard", tile_height=2000, load_wait=100 + ) + + assert result is not None + timeouts = [ + call[1]["timeout"] for call in mock_page.wait_for_function.call_args_list + ] + # remaining budget at each tile's wait: 1000, 50, 10 seconds + # -> capped timeouts shrink + assert timeouts == [100 * 1000, 50 * 1000, 10 * 1000] + assert timeouts == sorted(timeouts, reverse=True) + + def test_readiness_wait_uses_budget_recomputed_after_scroll_settle( + self, mock_page, monkeypatch + ): + """The readiness-wait timeout must be capped using the budget + recomputed *after* the scroll-settle sleep, not the stale value from + before it -- otherwise each tile could overrun the total budget by up + to one settle interval.""" + monkeypatch.setattr( + "superset.utils.screenshot_utils.TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS", # noqa: E501 + 1000, + ) + # A single-tile dashboard to keep the scenario simple. + mock_page.evaluate.return_value = { + "height": 1000, + "top": 100, + "left": 50, + "width": 800, + } + clock = self._FakeClock() + # The scroll-settle sleep itself consumes 950s of wall-clock time, + # leaving only 50s of the 1000s budget by the time the readiness + # wait is capped. + mock_page.wait_for_timeout.side_effect = lambda *args, **kwargs: setattr( + clock, "now", clock.now + 950 + ) + + with patch("superset.utils.screenshot_utils.current_task", None): + with patch("superset.utils.screenshot_utils.time.monotonic", new=clock): + with patch("superset.utils.screenshot_utils.combine_screenshot_tiles"): + take_tiled_screenshot( + mock_page, "dashboard", tile_height=2000, load_wait=999 + ) + + timeout = mock_page.wait_for_function.call_args_list[0][1]["timeout"] + # Must reflect the post-settle remaining budget (50s), not the + # stale pre-settle value (1000s, which would have let load_wait's + # full 999s through uncapped). + assert timeout == 50 * 1000 + + def test_budget_exhausted_raises_and_stops_capturing(self, mock_page, monkeypatch): + """Exhausting the budget aborts cleanly instead of capturing unchecked.""" + monkeypatch.setattr( + "superset.utils.screenshot_utils.TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS", # noqa: E501 + 1000, + ) + clock = self._FakeClock() + # Tile 0's readiness wait consumes the whole budget; tile 1's budget + # check then sees remaining <= 0 and raises before capturing. + mock_page.wait_for_function.side_effect = lambda *args, **kwargs: setattr( + clock, "now", 1000.0 + ) + + with patch("superset.utils.screenshot_utils.current_task", None): + with patch("superset.utils.screenshot_utils.time.monotonic", new=clock): + with patch( + "superset.utils.screenshot_utils.combine_screenshot_tiles" + ) as mock_combine: + with patch("superset.utils.screenshot_utils.logger") as mock_logger: + with pytest.raises(TiledScreenshotBudgetExceededError): + take_tiled_screenshot( + mock_page, "dashboard", tile_height=2000, load_wait=100 + ) + + # Only the first tile was captured before the budget ran out. + assert mock_page.screenshot.call_count == 1 + # Tiles were never combined -- the function raised before that point. + mock_combine.assert_not_called() + + # Budget exhaustion is a customer chart-loading issue, not a Superset + # system fault, so it must log at WARNING (not ERROR) -- consistent + # with the #38130/#38441 precedent for screenshot timeout logging. + assert mock_logger.error.call_count == 0 + mock_logger.warning.assert_called_once() + warning_args = mock_logger.warning.call_args[0] + assert "budget exhausted" in warning_args[0] + # tile index, tiles total, tiles captured, tiles total, + # elapsed seconds, budget seconds, log-context suffix + assert warning_args[1] == 2 + assert warning_args[2] == 3 + assert warning_args[3] == 1 + assert warning_args[4] == 3 + assert warning_args[5] == 1000 + assert warning_args[6] == 1000 + assert warning_args[7] == "" + + def test_budget_exhausted_warning_includes_log_context( + self, mock_page, monkeypatch + ): + """log_context (e.g. report execution id) is appended to the warning.""" + monkeypatch.setattr( + "superset.utils.screenshot_utils.TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS", # noqa: E501 + 1000, + ) + clock = self._FakeClock() + # Tile 0's readiness wait consumes the whole budget; tile 1's budget + # check then sees remaining <= 0 and raises. + mock_page.wait_for_function.side_effect = lambda *args, **kwargs: setattr( + clock, "now", 1000.0 + ) + with patch("superset.utils.screenshot_utils.current_task", None): + with patch("superset.utils.screenshot_utils.time.monotonic", new=clock): + with patch("superset.utils.screenshot_utils.combine_screenshot_tiles"): + with patch("superset.utils.screenshot_utils.logger") as mock_logger: + with pytest.raises(TiledScreenshotBudgetExceededError): + take_tiled_screenshot( + mock_page, + "dashboard", + tile_height=2000, + load_wait=100, + log_context="execution_id=abc-123", + ) + + warning_args = mock_logger.warning.call_args[0] + assert warning_args[-1] == " [execution_id=abc-123]" + + def test_budget_exhausted_before_first_tile_raises_without_capture( + self, mock_page, monkeypatch + ): + """No budget floor: a budget already exhausted by setup (element + lookup/dimension probing) raises before the first tile is captured, + matching the non-tiled path's raise-before-capture semantics.""" + monkeypatch.setattr( + "superset.utils.screenshot_utils.TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS", # noqa: E501 + 1000, + ) + clock = self._FakeClock() + # The dashboard-dimension evaluate() itself consumes the whole budget. + original_return = {"height": 5000, "top": 100, "left": 50, "width": 800} + + def slow_evaluate(*args, **kwargs): + clock.now = 1000.0 + return original_return + + mock_page.evaluate.side_effect = slow_evaluate + + with patch("superset.utils.screenshot_utils.current_task", None): + with patch("superset.utils.screenshot_utils.time.monotonic", new=clock): + with patch( + "superset.utils.screenshot_utils.combine_screenshot_tiles" + ) as mock_combine: + with pytest.raises(TiledScreenshotBudgetExceededError): + take_tiled_screenshot( + mock_page, "dashboard", tile_height=2000, load_wait=100 + ) + + mock_page.screenshot.assert_not_called() + mock_combine.assert_not_called() + + def test_no_celery_context_uses_fixed_total_fallback(self, mock_page): + """Outside Celery the helper returns None; the tiled path must fall + back to the fixed total ceiling rather than running uncapped, because + per-tile waits accumulate across tiles.""" + clock = self._FakeClock() + with patch("superset.utils.screenshot_utils.current_task", None): + with patch("superset.utils.screenshot_utils.time.monotonic", new=clock): + with patch("superset.utils.screenshot_utils.combine_screenshot_tiles"): + take_tiled_screenshot( + mock_page, + "dashboard", + tile_height=2000, + load_wait=10_000, # deliberately above the fallback + ) + + first_timeout = mock_page.wait_for_function.call_args_list[0][1]["timeout"] + assert first_timeout == TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS * 1000 + + def test_derived_task_budget_caps_tile_wait(self, mock_page): + """Inside Celery, the tiled path caps waits using the same + task-derived budget as the non-tiled path (helper reuse, #42427).""" + task = MagicMock() + task.request.timelimit = (120, None) # (hard, soft): 120s hard limit + + clock = self._FakeClock() + with patch("superset.utils.screenshot_utils.current_task", task): + with patch("superset.utils.screenshot_utils.time.monotonic", new=clock): + with patch("superset.utils.screenshot_utils.combine_screenshot_tiles"): + take_tiled_screenshot( + mock_page, "dashboard", tile_height=2000, load_wait=200 + ) + + # margin = min(300, 120 * 0.2) = 24; budget = 120 - 24 = 96 + first_timeout = mock_page.wait_for_function.call_args_list[0][1]["timeout"] + assert first_timeout == 96 * 1000 + assert first_timeout < 200 * 1000 + + def test_fast_dashboard_matches_default_behavior(self, mock_page): + """Well under budget, waits are not capped and behavior is unchanged.""" + with patch("superset.utils.screenshot_utils.current_task", None): + with patch("superset.utils.screenshot_utils.combine_screenshot_tiles"): + result = take_tiled_screenshot( + mock_page, + "dashboard", + tile_height=2000, + load_wait=30, + animation_wait=5, + ) + + assert result is not None + assert mock_page.screenshot.call_count == 3 + + for call in mock_page.wait_for_function.call_args_list: + assert call[1]["timeout"] == 30 * 1000 + + animation_calls = [ + call + for call in mock_page.wait_for_timeout.call_args_list + if call[0][0] == 5 * 1000 + ] + assert len(animation_calls) == 3 + + def test_per_tile_timing_debug_line_logged(self, mock_page): + """Each tile logs a DEBUG timing breakdown (readiness wait, animation + wait) so slow dashboards can be profiled from logs alone.""" + with patch("superset.utils.screenshot_utils.current_task", None): + with patch("superset.utils.screenshot_utils.logger") as mock_logger: + with patch("superset.utils.screenshot_utils.combine_screenshot_tiles"): + take_tiled_screenshot( + mock_page, + "dashboard", + tile_height=2000, + log_context="cache_key=xyz", + ) + + timing_calls = [ + call for call in mock_logger.debug.call_args_list if "timing" in call[0][0] + ] + assert len(timing_calls) == 3 + for i, call in enumerate(timing_calls): + args = call[0] + assert args[1] == i + 1 # tile index + assert args[2] == 3 # total tiles + assert args[-1] == " [cache_key=xyz]"