ui: Refactor Flame Graph color mapping logic

Author: yomete

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

@@ -130,8 +130,6 @@ export const FlameNode = React.memo(
     const colorAttribute =
       colorBy === 'filename' ? filename : colorBy === 'binary' ? mappingFile : null;
 
-    const colorsMap = colors;
-
     const hoveringName =
       hoveringRow !== undefined ? arrowToString(functionNameColumn?.get(hoveringRow)) : '';
     const shouldBeHighlighted =
@@ -142,7 +140,7 @@ export const FlameNode = React.memo(
       compareMode,
       cumulative,
       diff,
-      colorsMap,
+      colorsMap: colors,
       colorAttribute,
     });
 
@@ -296,7 +294,9 @@ export const FlameNode = React.memo(
       prevProps.hoveringRow === nextProps.hoveringRow &&
       prevProps.totalWidth === nextProps.totalWidth &&
       prevProps.height === nextProps.height &&
-      prevProps.effectiveDepth === nextProps.effectiveDepth
+      prevProps.effectiveDepth === nextProps.effectiveDepth &&
+      prevProps.colorBy === nextProps.colorBy &&
+      prevProps.colors === nextProps.colors
     );
   }
 );

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

@@ -21,15 +21,15 @@ import React, {
   useState,
 } from 'react';
 
-import {Dictionary, Table, Vector, tableFromIPC} from 'apache-arrow';
+import {Table, tableFromIPC} from 'apache-arrow';
 import {useContextMenu} from 'react-contexify';
 
 import {FlamegraphArrow} from '@parca/client';
-import {useParcaContext, useURLState} from '@parca/components';
+import {useParcaContext} from '@parca/components';
 import {USER_PREFERENCES, useCurrentColorProfile, useUserPreference} from '@parca/hooks';
 import {ProfileType} from '@parca/parser';
 import {getColorForFeature, selectDarkMode, useAppSelector} from '@parca/store';
-import {getLastItem, type ColorConfig} from '@parca/utilities';
+import {type ColorConfig} from '@parca/utilities';
 
 import {ProfileSource} from '../../ProfileSource';
 import {useProfileFilters} from '../../ProfileView/components/ProfileFilters/useProfileFilters';
@@ -38,12 +38,10 @@ import ContextMenuWrapper, {ContextMenuWrapperRef} from './ContextMenuWrapper';
 import {FlameNode, RowHeight, colorByColors} from './FlameGraphNodes';
 import {MemoizedTooltip} from './MemoizedTooltip';
 import {TooltipProvider} from './TooltipContext';
-import {useFilenamesList} from './useMappingList';
 import {useScrollViewport} from './useScrollViewport';
 import {useVisibleNodes} from './useVisibleNodes';
 import {
   CurrentPathFrame,
-  arrowToString,
   extractFeature,
   extractFilenameFeature,
   getCurrentPathFrameData,
@@ -84,6 +82,8 @@ interface FlameGraphArrowProps {
   setCurPath: (path: CurrentPathFrame[]) => void;
   isHalfScreen: boolean;
   mappingsListFromMetadata: string[];
+  filenamesListFromMetadata: string[];
+  colorBy: string;
   compareAbsolute: boolean;
   isFlameChart?: boolean;
   isRenderedAsFlamegraph?: boolean;
@@ -139,6 +139,9 @@ export const FlameGraphArrow = memo(function FlameGraphArrow({
   tooltipId = 'default',
   maxFrameCount,
   isExpanded = false,
+  mappingsListFromMetadata,
+  filenamesListFromMetadata,
+  colorBy,
 }: FlameGraphArrowProps): React.JSX.Element {
   const [highlightSimilarStacksPreference] = useUserPreference<boolean>(
     USER_PREFERENCES.HIGHLIGHT_SIMILAR_STACKS.key
@@ -169,49 +172,17 @@ export const FlameGraphArrow = memo(function FlameGraphArrow({
   const currentColorProfile = useCurrentColorProfile();
   const colorForSimilarNodes = currentColorProfile.colorForSimilarNodes;
 
-  const [colorBy, _] = useURLState('color_by');
-  const colorByValue = colorBy === undefined || colorBy === '' ? 'binary' : (colorBy as string);
-
-  const filenamesList = useFilenamesList(table);
-
-  const mappingsList = useMemo(() => {
-    // Read the mappings from the dictionary that contains all mapping strings.
-    // This is great, as might only have a dozen or so mappings,
-    // and don't need to read through all the rows (potentially thousands).
-    const mappingsDict: Vector<Dictionary> | null = table.getChild(FIELD_MAPPING_FILE);
-    const mappings =
-      mappingsDict?.data
-        .map(mapping => {
-          if (mapping.dictionary == null) {
-            return [];
-          }
-          const len = mapping.dictionary.length;
-          const entries: string[] = [];
-          for (let i = 0; i < len; i++) {
-            const fn = arrowToString(mapping.dictionary.get(i));
-            entries.push(getLastItem(fn) ?? '');
-          }
-          return entries;
-        })
-        .flat() ?? [];
-
-    // We add a EVERYTHING ELSE mapping to the list.
-    mappings.push('');
-
-    // We sort the mappings alphabetically to make sure that the order is always the same.
-    mappings.sort((a, b) => a.localeCompare(b));
-    return mappings;
-  }, [table]);
+  const colorByValue = colorBy === undefined || colorBy === '' ? 'binary' : colorBy;
 
   const filenameColors = useMemo(() => {
-    const colors = getFilenameColors(filenamesList, isDarkMode, currentColorProfile);
+    const colors = getFilenameColors(filenamesListFromMetadata, isDarkMode, currentColorProfile);
     return colors;
-  }, [isDarkMode, filenamesList, currentColorProfile]);
+  }, [isDarkMode, filenamesListFromMetadata, currentColorProfile]);
 
   const mappingColors = useMemo(() => {
-    const colors = getMappingColors(mappingsList, isDarkMode, currentColorProfile);
+    const colors = getMappingColors(mappingsListFromMetadata, isDarkMode, currentColorProfile);
     return colors;
-  }, [isDarkMode, mappingsList, currentColorProfile]);
+  }, [isDarkMode, mappingsListFromMetadata, currentColorProfile]);
 
   const colorByList = {
     filename: filenameColors,

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

@@ -31,9 +31,10 @@ import {capitalizeOnlyFirstLetter, divide} from '@parca/utilities';
 import {MergedProfileSource, ProfileSource} from '../ProfileSource';
 import DiffLegend from '../ProfileView/components/DiffLegend';
 import {useProfileViewContext} from '../ProfileView/context/ProfileViewContext';
+import {useProfileMetadata} from '../ProfileView/hooks/useProfileMetadata';
+import {useVisualizationState} from '../ProfileView/hooks/useVisualizationState';
 import {TimelineGuide} from '../TimelineGuide';
 import {FlameGraphArrow} from './FlameGraphArrow';
-import useMappingList from './FlameGraphArrow/useMappingList';
 import {CurrentPathFrame, boundsFromProfileSource} from './FlameGraphArrow/utils';
 
 const numberFormatter = new Intl.NumberFormat('en-US');
@@ -100,11 +101,13 @@ const ProfileFlameGraph = function ProfileFlameGraphNonMemo({
   tooltipId,
   maxFrameCount,
   isExpanded = false,
+  metadataLoading = false,
 }: ProfileFlameGraphProps): JSX.Element {
   const {onError, authenticationErrorMessage, isDarkMode, flamechartHelpText} = useParcaContext();
   const {compareMode} = useProfileViewContext();
   const [isLoading, setIsLoading] = useState<boolean>(true);
   const [flameChartRef, {height: flameChartHeight}] = useMeasure();
+  const {colorBy, setColorBy} = useVisualizationState();
 
   // Create local state for paths when in sandwich view to avoid URL updates
   const [localCurPathArrow, setLocalCurPathArrow] = useState<CurrentPathFrame[]>([]);
@@ -123,9 +126,12 @@ const ProfileFlameGraph = function ProfileFlameGraphNonMemo({
   // Determine which paths to use based on isInSandwichView flag
   const effectiveCurPathArrow = isInSandwichView ? localCurPathArrow : curPathArrow;
 
-  const mappingsList = useMappingList(metadataMappingFiles);
-
-  const [colorBy, setColorBy] = useURLState('color_by');
+  const {mappingsList, filenamesList} = useProfileMetadata({
+    flamegraphArrow: arrow,
+    metadataMappingFiles,
+    metadataLoading,
+    colorBy,
+  });
 
   // By default, we want delta profiles (CPU) to be relatively compared.
   // For non-delta profiles, like goroutines or memory, we want the profiles to be compared absolutely.
@@ -279,6 +285,7 @@ const ProfileFlameGraph = function ProfileFlameGraphNonMemo({
               profileType={profileType}
               isHalfScreen={isHalfScreen}
               mappingsListFromMetadata={mappingsList}
+              filenamesListFromMetadata={filenamesList}
               compareAbsolute={isCompareAbsolute}
               isFlameChart={isFlameChart}
               profileSource={profileSource}
@@ -287,6 +294,7 @@ const ProfileFlameGraph = function ProfileFlameGraphNonMemo({
               tooltipId={tooltipId}
               maxFrameCount={maxFrameCount}
               isExpanded={isExpanded}
+              colorBy={colorBy}
             />
           </div>
         </div>
@@ -302,7 +310,6 @@ const ProfileFlameGraph = function ProfileFlameGraphNonMemo({
     profileType,
     isHalfScreen,
     isDarkMode,
-    mappingsList,
     isCompareAbsolute,
     isFlameChart,
     profileSource,
@@ -316,6 +323,9 @@ const ProfileFlameGraph = function ProfileFlameGraphNonMemo({
     tooltipId,
     maxFrameCount,
     isExpanded,
+    mappingsList,
+    filenamesList,
+    colorBy,
   ]);
 
   useEffect(() => {