From d10a9ee4338f7c333a37527ff0f27e0eb0e8ef35 Mon Sep 17 00:00:00 2001
From: pierrejeambrun <pierrejbrun@gmail.com>
Date: Wed, 15 Jul 2026 13:41:09 +0200
Subject: [PATCH] Re-render only the changed column when Grid summaries stream
 in
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The Dag Grid streams TI summaries as one NDJSON line per run and stores them in
a Map; every line replaced the Map and re-rendered Grid. Because useVirtualizer
(@tanstack/react-virtual) is on the React Compiler's known-incompatible list, the
compiler skips optimizing all of Grid — so flatNodes and the click handlers got
fresh references on every render and every column re-rendered on every line, not
just the run whose summary arrived. On a ~20-run x 30-task grid that was 117
spurious column re-renders per load.

Memoize the values handed down to the columns by hand — the escape hatch the
compiler's own bailout message points to for incompatible-library APIs. Measured
on the same grid: column re-renders per load dropped 171 -> 54 (all legitimate),
grid load render work ~47% lower.

Also drop a redundant immediate stream restart: an unconditional refresh tick
fired the instant the interval effect mounted, aborting and reopening the
just-opened mount stream (an AbortError on every grid mount with active runs).

related: #69531
---
 .../airflow/ui/src/layouts/Details/Grid/Grid.tsx  | 15 ++++++++++-----
 .../airflow/ui/src/queries/useGridTISummaries.ts  |  6 ++++--
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/airflow-core/src/airflow/ui/src/layouts/Details/Grid/Grid.tsx b/airflow-core/src/airflow/ui/src/layouts/Details/Grid/Grid.tsx
index eed38a10ad48f..a1cb9d36e991c 100644
--- a/airflow-core/src/airflow/ui/src/layouts/Details/Grid/Grid.tsx
+++ b/airflow-core/src/airflow/ui/src/layouts/Details/Grid/Grid.tsx
@@ -20,7 +20,7 @@ import { Box, Flex } from "@chakra-ui/react";
 import { useVirtualizer } from "@tanstack/react-virtual";
 import dayjs from "dayjs";
 import dayjsDuration from "dayjs/plugin/duration";
-import { useRef } from "react";
+import { useCallback, useMemo, useRef } from "react";
 import type { RefObject } from "react";
 import { useParams, useSearchParams } from "react-router-dom";
 
@@ -141,7 +141,12 @@ export const Grid = ({
     showVersionIndicatorMode,
   });
 
-  const { flatNodes } = flattenNodes(dagStructure, openGroupIds);
+  // React Compiler skips optimizing this whole component: `useVirtualizer` (@tanstack/react-virtual) is
+  // on the compiler's known-incompatible list — its return value exposes functions that can't be
+  // memoized safely — so it declines to memoize anything in Grid. Without the manual memoization here
+  // and on the click handlers below, `flatNodes` and the handlers get fresh references every render, so
+  // each TI-summaries stream line re-renders every column instead of only the run whose summary changed.
+  const { flatNodes } = useMemo(() => flattenNodes(dagStructure, openGroupIds), [dagStructure, openGroupIds]);
 
   const taskNameColumnWidthPx = showGantt ? estimateTaskNameColumnWidthPx(flatNodes) : undefined;
 
@@ -166,9 +171,9 @@ export const Grid = ({
     tasks: flatNodes,
   });
 
-  const handleRowClick = () => setMode(NavigationModes.TASK);
-  const handleCellClick = () => setMode(NavigationModes.TI);
-  const handleColumnClick = () => setMode(NavigationModes.RUN);
+  const handleRowClick = useCallback(() => setMode(NavigationModes.TASK), [setMode]);
+  const handleCellClick = useCallback(() => setMode(NavigationModes.TI), [setMode]);
+  const handleColumnClick = useCallback(() => setMode(NavigationModes.RUN), [setMode]);
 
   const rowVirtualizer = useVirtualizer({
     count: flatNodes.length,
diff --git a/airflow-core/src/airflow/ui/src/queries/useGridTISummaries.ts b/airflow-core/src/airflow/ui/src/queries/useGridTISummaries.ts
index 1ec67127b3f02..d90b2f0078ece 100644
--- a/airflow-core/src/airflow/ui/src/queries/useGridTISummaries.ts
+++ b/airflow-core/src/airflow/ui/src/queries/useGridTISummaries.ts
@@ -144,8 +144,10 @@ export const useGridTiSummariesStream = ({
       return undefined;
     }
 
-    // Kick off an immediate refresh so the stream doesn't have to wait for the first interval to elapse.
-    setRefreshTick((tick) => tick + 1);
+    // The stream already fetches on mount and whenever runIdsKey changes, so there is no first-interval
+    // wait to avoid. Bumping refreshTick here would abort that just-opened mount stream and immediately
+    // reopen it — a redundant connection plus an AbortError on every grid mount — so let the interval be
+    // the only re-stream trigger.
     const timer = setInterval(() => {
       setRefreshTick((tick) => tick + 1);
     }, baseRefetchInterval);