fix(reports): time-budget tiled screenshot to fail cleanly instead of hitting Celery kill

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

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