ui: Flamegraph rendering improvements (#6117)

* Flamegraph rendering improvements

* [pre-commit.ci lite] apply automatic fixes

* Linter fixes

---------

Co-authored-by: pre-commit-ci-lite[bot] <117423508+pre-commit-ci-lite[bot]@users.noreply.github.com>

Author: manojVivek

ui/packages/shared/profile/src/ProfileFlameGraph/FlameGraphArrow/FlameGraphNodes.tsx

@@ -11,7 +11,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import React, {useMemo} from 'react';
+import React, {useCallback, useMemo} from 'react';
 
 import {Table} from 'apache-arrow';
 import cx from 'classnames';
@@ -101,36 +101,52 @@ export const FlameNode = React.memo(
     effectiveDepth,
     tooltipId = 'default',
   }: FlameNodeProps): React.JSX.Element {
-    // get the columns to read from
-    const mappingColumn = table.getChild(FIELD_MAPPING_FILE);
-    const functionNameColumn = table.getChild(FIELD_FUNCTION_NAME);
-    const cumulativeColumn = table.getChild(FIELD_CUMULATIVE);
-    const depthColumn = table.getChild(FIELD_DEPTH);
-    const diffColumn = table.getChild(FIELD_DIFF);
-    const filenameColumn = table.getChild(FIELD_FUNCTION_FILE_NAME);
-    const valueOffsetColumn = table.getChild(FIELD_VALUE_OFFSET);
-    const tsColumn = table.getChild(FIELD_TIMESTAMP);
+    // Memoize column references - only changes when table changes
+    const columns = useMemo(
+      () => ({
+        mapping: table.getChild(FIELD_MAPPING_FILE),
+        functionName: table.getChild(FIELD_FUNCTION_NAME),
+        cumulative: table.getChild(FIELD_CUMULATIVE),
+        depth: table.getChild(FIELD_DEPTH),
+        diff: table.getChild(FIELD_DIFF),
+        filename: table.getChild(FIELD_FUNCTION_FILE_NAME),
+        valueOffset: table.getChild(FIELD_VALUE_OFFSET),
+        ts: table.getChild(FIELD_TIMESTAMP),
+      }),
+      [table]
+    );
 
     // get the actual values from the columns
     const binaries = useAppSelector(selectBinaries);
 
-    const mappingFile: string | null = arrowToString(mappingColumn?.get(row));
-    const functionName: string | null = arrowToString(functionNameColumn?.get(row));
-    const cumulative = cumulativeColumn?.get(row) != null ? BigInt(cumulativeColumn?.get(row)) : 0n;
-    const diff: bigint | null = diffColumn?.get(row) != null ? BigInt(diffColumn?.get(row)) : null;
-    const filename: string | null = arrowToString(filenameColumn?.get(row));
-    const depth: number = depthColumn?.get(row) ?? 0;
-
-    const valueOffset: bigint =
-      valueOffsetColumn?.get(row) !== null && valueOffsetColumn?.get(row) !== undefined
-        ? BigInt(valueOffsetColumn?.get(row))
-        : 0n;
+    // Memoize row data extraction - only changes when table or row changes
+    const rowData = useMemo(() => {
+      const mappingFile: string | null = arrowToString(columns.mapping?.get(row));
+      const functionName: string | null = arrowToString(columns.functionName?.get(row));
+      const cumulative =
+        columns.cumulative?.get(row) != null ? BigInt(columns.cumulative?.get(row)) : 0n;
+      const diff: bigint | null =
+        columns.diff?.get(row) != null ? BigInt(columns.diff?.get(row)) : null;
+      const filename: string | null = arrowToString(columns.filename?.get(row));
+      const depth: number = columns.depth?.get(row) ?? 0;
+      const valueOffset: bigint =
+        columns.valueOffset?.get(row) !== null && columns.valueOffset?.get(row) !== undefined
+          ? BigInt(columns.valueOffset?.get(row))
+          : 0n;
+
+      return {mappingFile, functionName, cumulative, diff, filename, depth, valueOffset};
+    }, [columns, row]);
+
+    const {mappingFile, functionName, cumulative, diff, filename, depth, valueOffset} = rowData;
 
     const colorAttribute =
       colorBy === 'filename' ? filename : colorBy === 'binary' ? mappingFile : null;
 
-    const hoveringName =
-      hoveringRow !== undefined ? arrowToString(functionNameColumn?.get(hoveringRow)) : '';
+    // Memoize hovering name lookup
+    const hoveringName = useMemo(() => {
+      return hoveringRow !== undefined ? arrowToString(columns.functionName?.get(hoveringRow)) : '';
+    }, [columns.functionName, hoveringRow]);
+
     const shouldBeHighlighted =
       functionName != null && hoveringName != null && functionName === hoveringName;
 
@@ -147,18 +163,59 @@ export const FlameNode = React.memo(
       return row === 0 ? 'root' : nodeLabel(table, row, binaries.length > 1);
     }, [table, row, binaries]);
 
+    // Memoize selection data - only changes when selectedRow changes
+    const selectionData = useMemo(() => {
+      const selectionOffset =
+        columns.valueOffset?.get(selectedRow) !== null &&
+        columns.valueOffset?.get(selectedRow) !== undefined
+          ? BigInt(columns.valueOffset?.get(selectedRow))
+          : 0n;
+      const selectionCumulative =
+        columns.cumulative?.get(selectedRow) !== null
+          ? BigInt(columns.cumulative?.get(selectedRow))
+          : 0n;
+      const selectedDepth = columns.depth?.get(selectedRow);
+      const total = columns.cumulative?.get(selectedRow);
+      return {selectionOffset, selectionCumulative, selectedDepth, total};
+    }, [columns, selectedRow]);
+
+    const {selectionOffset, selectionCumulative, selectedDepth, total} = selectionData;
+
+    // Memoize tsBounds - only changes when profileSource changes
+    const tsBounds = useMemo(() => boundsFromProfileSource(profileSource), [profileSource]);
+
+    // Memoize event handlers
+    const onMouseEnter = useCallback((): void => {
+      setHoveringRow(row);
+      window.dispatchEvent(
+        new CustomEvent(`flame-tooltip-update-${tooltipId}`, {
+          detail: {row},
+        })
+      );
+    }, [setHoveringRow, row, tooltipId]);
+
+    const onMouseLeave = useCallback((): void => {
+      setHoveringRow(undefined);
+      window.dispatchEvent(
+        new CustomEvent(`flame-tooltip-update-${tooltipId}`, {
+          detail: {row: null},
+        })
+      );
+    }, [setHoveringRow, tooltipId]);
+
+    const handleContextMenu = useCallback(
+      (e: React.MouseEvent): void => {
+        onContextMenu(e, row);
+      },
+      [onContextMenu, row]
+    );
+
+    // Early returns - all hooks must be called before this point
     // Hide frames beyond effective depth limit
     if (effectiveDepth !== undefined && depth > effectiveDepth) {
       return <></>;
     }
 
-    const selectionOffset =
-      valueOffsetColumn?.get(selectedRow) !== null &&
-      valueOffsetColumn?.get(selectedRow) !== undefined
-        ? BigInt(valueOffsetColumn?.get(selectedRow))
-        : 0n;
-    const selectionCumulative =
-      cumulativeColumn?.get(selectedRow) !== null ? BigInt(cumulativeColumn?.get(selectedRow)) : 0n;
     if (
       valueOffset + cumulative <= selectionOffset ||
       valueOffset >= selectionOffset + selectionCumulative
@@ -173,8 +230,6 @@ export const FlameNode = React.memo(
     }
 
     // Cumulative can be larger than total when a selection is made. All parents of the selection are likely larger, but we want to only show them as 100% in the graph.
-    const tsBounds = boundsFromProfileSource(profileSource);
-    const total = cumulativeColumn?.get(selectedRow);
     const totalRatio = cumulative > total ? 1 : Number(cumulative) / Number(total);
     const width: number = isFlameChart
       ? (Number(cumulative) / (Number(tsBounds[1]) - Number(tsBounds[0]))) * totalWidth
@@ -184,35 +239,12 @@ export const FlameNode = React.memo(
       return <></>;
     }
 
-    const selectedDepth = depthColumn?.get(selectedRow);
     const styles =
       selectedDepth !== undefined && selectedDepth > depth ? fadedFlameRectStyles : flameRectStyles;
 
-    const onMouseEnter = (): void => {
-      setHoveringRow(row);
-      window.dispatchEvent(
-        new CustomEvent(`flame-tooltip-update-${tooltipId}`, {
-          detail: {row},
-        })
-      );
-    };
-
-    const onMouseLeave = (): void => {
-      setHoveringRow(undefined);
-      window.dispatchEvent(
-        new CustomEvent(`flame-tooltip-update-${tooltipId}`, {
-          detail: {row: null},
-        })
-      );
-    };
-
-    const handleContextMenu = (e: React.MouseEvent): void => {
-      onContextMenu(e, row);
-    };
-
-    const ts = tsColumn !== null ? Number(tsColumn.get(row)) : 0;
+    const ts = columns.ts !== null ? Number(columns.ts.get(row)) : 0;
     const x =
-      isFlameChart && tsColumn !== null
+      isFlameChart && columns.ts !== null
         ? ((ts - Number(tsBounds[0])) / (Number(tsBounds[1]) - Number(tsBounds[0]))) * totalWidth
         : selectedDepth > depth
         ? 0

ui/packages/shared/profile/src/ProfileFlameGraph/FlameGraphArrow/index.tsx

@@ -25,7 +25,7 @@ import {Table, tableFromIPC} from 'apache-arrow';
 import {useContextMenu} from 'react-contexify';
 
 import {FlamegraphArrow} from '@parca/client';
-import {useParcaContext} from '@parca/components';
+import {FlameGraphSkeleton, SandwichFlameGraphSkeleton, useParcaContext} from '@parca/components';
 import {USER_PREFERENCES, useCurrentColorProfile, useUserPreference} from '@parca/hooks';
 import {ProfileType} from '@parca/parser';
 import {getColorForFeature, selectDarkMode, useAppSelector} from '@parca/store';
@@ -38,6 +38,7 @@ import ContextMenuWrapper, {ContextMenuWrapperRef} from './ContextMenuWrapper';
 import {FlameNode, RowHeight, colorByColors} from './FlameGraphNodes';
 import {MemoizedTooltip} from './MemoizedTooltip';
 import {TooltipProvider} from './TooltipContext';
+import {useBatchedRendering} from './useBatchedRendering';
 import {useScrollViewport} from './useScrollViewport';
 import {useVisibleNodes} from './useVisibleNodes';
 import {
@@ -136,6 +137,7 @@ export const FlameGraphArrow = memo(function FlameGraphArrow({
   isFlameChart = false,
   isRenderedAsFlamegraph = false,
   isInSandwichView = false,
+  isHalfScreen,
   tooltipId = 'default',
   maxFrameCount,
   isExpanded = false,
@@ -163,6 +165,7 @@ export const FlameGraphArrow = memo(function FlameGraphArrow({
   const svg = useRef(null);
   const containerRef = useRef<HTMLDivElement>(null);
   const renderStartTime = useRef<number>(0);
+  const hasInitialRenderCompleted = useRef(false);
 
   const [svgElement, setSvgElement] = useState<SVGSVGElement | null>(null);
 
@@ -291,6 +294,18 @@ export const FlameGraphArrow = memo(function FlameGraphArrow({
     effectiveDepth: deferredEffectiveDepth,
   });
 
+  // Add nodes in incremental batches to avoid blocking the UI
+  const {items: batchedNodes, isComplete: isBatchingComplete} = useBatchedRendering(visibleNodes, {
+    batchSize: 500,
+  });
+  if (isBatchingComplete) {
+    hasInitialRenderCompleted.current = true;
+  }
+
+  // Show skeleton only during initial load, not during scroll updates
+  const showSkeleton =
+    !hasInitialRenderCompleted.current && batchedNodes.length !== visibleNodes.length;
+
   useEffect(() => {
     if (perf?.markInteraction != null) {
       renderStartTime.current = performance.now();
@@ -327,12 +342,22 @@ export const FlameGraphArrow = memo(function FlameGraphArrow({
           isInSandwichView={isInSandwichView}
         />
         <MemoizedTooltip contextElement={svgElement} dockedMetainfo={dockedMetainfo} />
+        {showSkeleton && (
+          <div className="absolute inset-0 z-10">
+            {isRenderedAsFlamegraph ? (
+              <SandwichFlameGraphSkeleton isHalfScreen={isHalfScreen} isDarkMode={isDarkMode} />
+            ) : (
+              <FlameGraphSkeleton isHalfScreen={isHalfScreen} isDarkMode={isDarkMode} />
+            )}
+          </div>
+        )}
         <div
           ref={containerRef}
           className="overflow-auto scrollbar-thin scrollbar-thumb-gray-400 scrollbar-track-gray-100 dark:scrollbar-thumb-gray-600 dark:scrollbar-track-gray-800 will-change-transform scroll-smooth webkit-overflow-scrolling-touch contain"
           style={{
             width: width ?? '100%',
             contain: 'layout style paint',
+            visibility: !showSkeleton ? 'visible' : 'hidden',
           }}
         >
           <svg
@@ -342,7 +367,7 @@ export const FlameGraphArrow = memo(function FlameGraphArrow({
             preserveAspectRatio="xMinYMid"
             ref={svg}
           >
-            {visibleNodes.map(row => (
+            {batchedNodes.map(row => (
               <FlameNode
                 key={row}
                 table={table}

ui/packages/shared/profile/src/ProfileFlameGraph/FlameGraphArrow/useBatchedRendering.ts

@@ -0,0 +1,84 @@
+// Copyright 2022 The Parca Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import {useEffect, useRef, useState} from 'react';
+
+interface UseBatchedRenderingOptions {
+  batchSize?: number;
+  // Delay between batches in ms (0 = next animation frame)
+  batchDelay?: number;
+}
+
+interface UseBatchedRenderingResult<T> {
+  items: T[];
+  isComplete: boolean;
+}
+
+// useBatchedRendering - Helps in incrementally rendering items in batches to avoid UI blocking.
+export const useBatchedRendering = <T>(
+  items: T[],
+  options: UseBatchedRenderingOptions = {}
+): UseBatchedRenderingResult<T> => {
+  const {batchSize = 500, batchDelay = 0} = options;
+
+  const [renderedCount, setRenderedCount] = useState(0);
+  const itemsRef = useRef(items);
+  const rafRef = useRef<number | null>(null);
+  const timeoutRef = useRef<NodeJS.Timeout | null>(null);
+
+  useEffect(() => {
+    if (itemsRef.current !== items) {
+      itemsRef.current = items;
+      setRenderedCount(prev => {
+        if (items.length === 0) return 0;
+        // If new items were added (scrolling down), keep current progress
+        if (items.length > prev) return prev;
+        // If items reduced, cap to new length
+        return Math.min(prev, items.length);
+      });
+    }
+  }, [items]);
+
+  // Progressively render more items
+  useEffect(() => {
+    if (renderedCount === items.length) {
+      return;
+    }
+
+    const scheduleNextBatch = (): void => {
+      const incrementState = (): void => {
+        setRenderedCount(prev => Math.min(prev + batchSize, items.length));
+      };
+      if (batchDelay > 0) {
+        timeoutRef.current = setTimeout(incrementState, batchDelay);
+      } else {
+        rafRef.current = requestAnimationFrame(incrementState);
+      }
+    };
+    scheduleNextBatch();
+
+    return () => {
+      if (rafRef.current !== null) {
+        cancelAnimationFrame(rafRef.current);
+      }
+      if (timeoutRef.current !== null) {
+        clearTimeout(timeoutRef.current);
+      }
+    };
+  }, [renderedCount, items.length, batchSize, batchDelay]);
+
+  return {
+    items: items.slice(0, renderedCount),
+    isComplete: renderedCount === items.length,
+  };
+};