Flamegraph: Add collapse and expand group buttons to toolbar (#87395)

packages/grafana-flamegraph/src/FlameGraph/FlameGraph.test.tsx

@@ -49,6 +49,8 @@ describe('FlameGraph', () => {
         colorScheme={ColorScheme.ValueBased}
         selectedView={SelectedView.FlameGraph}
         search={''}
+        collapsedMap={container.getCollapsedMap()}
+        setCollapsedMap={() => {}}
         {...props}
       />
     );

packages/grafana-flamegraph/src/FlameGraph/FlameGraph.tsx

@@ -50,6 +50,8 @@ type Props = {
   collapsing?: boolean;
   selectedView: SelectedView;
   search: string;
+  collapsedMap: CollapsedMap;
+  setCollapsedMap: (collapsedMap: CollapsedMap) => void;
 };
 
 const FlameGraph = ({
@@ -72,10 +74,11 @@ const FlameGraph = ({
   collapsing,
   selectedView,
   search,
+  collapsedMap,
+  setCollapsedMap,
 }: Props) => {
   const styles = getStyles();
 
-  const [collapsedMap, setCollapsedMap] = useState<CollapsedMap>(new Map());
   const [levels, setLevels] = useState<LevelItem[][]>();
   const [levelsCallers, setLevelsCallers] = useState<LevelItem[][]>();
   const [totalProfileTicks, setTotalProfileTicks] = useState<number>(0);
@@ -84,8 +87,6 @@ const FlameGraph = ({
 
   useEffect(() => {
     if (data) {
-      setCollapsedMap(data.getCollapsedMap());
-
       let levels = data.getLevels();
       let totalProfileTicks = levels.length ? levels[0][0].value : 0;
       let totalProfileTicksRight = levels.length ? levels[0][0].valueRight : undefined;

packages/grafana-flamegraph/src/FlameGraph/FlameGraphCanvas.tsx

@@ -7,7 +7,7 @@ import { ClickedItemData, ColorScheme, ColorSchemeDiff, SelectedView, TextAlign
 
 import FlameGraphContextMenu, { GetExtraContextMenuButtonsFunction } from './FlameGraphContextMenu';
 import FlameGraphTooltip from './FlameGraphTooltip';
-import { CollapseConfig, CollapsedMap, FlameGraphDataContainer, LevelItem } from './dataTransform';
+import { CollapsedMap, FlameGraphDataContainer, LevelItem } from './dataTransform';
 import { getBarX, useFlameRender } from './rendering';
 
 type Props = {
@@ -209,16 +209,16 @@ const FlameGraphCanvas = ({
             onSandwich(data.getLabel(clickedItemData.item.itemIndexes[0]));
           }}
           onExpandGroup={() => {
-            setCollapsedMap(setCollapsedStatus(collapsedMap, clickedItemData.item, false));
+            setCollapsedMap(collapsedMap.setCollapsedStatus(clickedItemData.item, false));
           }}
           onCollapseGroup={() => {
-            setCollapsedMap(setCollapsedStatus(collapsedMap, clickedItemData.item, true));
+            setCollapsedMap(collapsedMap.setCollapsedStatus(clickedItemData.item, true));
           }}
           onExpandAllGroups={() => {
-            setCollapsedMap(setAllCollapsedStatus(collapsedMap, false));
+            setCollapsedMap(collapsedMap.setAllCollapsedStatus(false));
           }}
           onCollapseAllGroups={() => {
-            setCollapsedMap(setAllCollapsedStatus(collapsedMap, true));
+            setCollapsedMap(collapsedMap.setAllCollapsedStatus(true));
           }}
           allGroupsCollapsed={Array.from(collapsedMap.values()).every((i) => i.collapsed)}
           allGroupsExpanded={Array.from(collapsedMap.values()).every((i) => !i.collapsed)}
@@ -231,27 +231,6 @@ const FlameGraphCanvas = ({
   );
 };
 
-function setCollapsedStatus(collapsedMap: CollapsedMap, item: LevelItem, collapsed: boolean) {
-  const newMap = new Map(collapsedMap);
-  const collapsedConfig = collapsedMap.get(item)!;
-  const newConfig = { ...collapsedConfig, collapsed };
-  for (const item of collapsedConfig.items) {
-    newMap.set(item, newConfig);
-  }
-  return newMap;
-}
-
-function setAllCollapsedStatus(collapsedMap: CollapsedMap, collapsed: boolean) {
-  const newMap = new Map(collapsedMap);
-  for (const item of collapsedMap.keys()) {
-    const collapsedConfig = collapsedMap.get(item)!;
-    const newConfig = { ...collapsedConfig, collapsed };
-    newMap.set(item, newConfig);
-  }
-
-  return newMap;
-}
-
 const getStyles = () => ({
   graph: css({
     label: 'graph',
@@ -291,7 +270,7 @@ export const convertPixelCoordinatesToBarCoordinates = (
   pixelsPerTick: number,
   totalTicks: number,
   rangeMin: number,
-  collapsedMap: Map<LevelItem, CollapseConfig>
+  collapsedMap: CollapsedMap
 ): LevelItem | undefined => {
   let next: LevelItem | undefined = root;
   let currentLevel = direction === 'children' ? 0 : depth - 1;

packages/grafana-flamegraph/src/FlameGraph/dataTransform.test.ts

@@ -1,6 +1,12 @@
 import { createDataFrame, FieldType } from '@grafana/data';
 
-import { CollapsedMapContainer, FlameGraphDataContainer, LevelItem, nestedSetToLevels } from './dataTransform';
+import {
+  CollapsedMapBuilder,
+  FlameGraphDataContainer,
+  LevelItem,
+  nestedSetToLevels,
+  CollapsedMap,
+} from './dataTransform';
 import { textToDataContainer } from './testHelpers';
 
 describe('nestedSetToLevels', () => {
@@ -126,7 +132,7 @@ describe('CollapsedMapContainer', () => {
   };
 
   it('groups items if they are within value threshold', () => {
-    const container = new CollapsedMapContainer();
+    const container = new CollapsedMapBuilder();
 
     const child2: LevelItem = {
       ...defaultItem,
@@ -147,13 +153,13 @@ describe('CollapsedMapContainer', () => {
 
     container.addItem(child1, parent);
     container.addItem(child2, child1);
-    expect(container.getMap().get(child1)).toMatchObject({ collapsed: true, items: [parent, child1, child2] });
-    expect(container.getMap().get(child2)).toMatchObject({ collapsed: true, items: [parent, child1, child2] });
-    expect(container.getMap().get(parent)).toMatchObject({ collapsed: true, items: [parent, child1, child2] });
+    expect(container.getCollapsedMap().get(child1)).toMatchObject({ collapsed: true, items: [parent, child1, child2] });
+    expect(container.getCollapsedMap().get(child2)).toMatchObject({ collapsed: true, items: [parent, child1, child2] });
+    expect(container.getCollapsedMap().get(parent)).toMatchObject({ collapsed: true, items: [parent, child1, child2] });
   });
 
   it("doesn't group items if they are outside value threshold", () => {
-    const container = new CollapsedMapContainer();
+    const container = new CollapsedMapBuilder();
 
     const parent: LevelItem = {
       ...defaultItem,
@@ -166,11 +172,11 @@ describe('CollapsedMapContainer', () => {
     };
 
     container.addItem(child, parent);
-    expect(container.getMap().size).toBe(0);
+    expect(container.getCollapsedMap().size()).toBe(0);
   });
 
   it("doesn't group items if parent has multiple children", () => {
-    const container = new CollapsedMapContainer();
+    const container = new CollapsedMapBuilder();
 
     const child1: LevelItem = {
       ...defaultItem,
@@ -190,6 +196,52 @@ describe('CollapsedMapContainer', () => {
     };
 
     container.addItem(child1, parent);
-    expect(container.getMap().size).toBe(0);
+    expect(container.getCollapsedMap().size()).toBe(0);
+  });
+});
+
+describe('CollapsedMap', () => {
+  function getMap() {
+    const container = textToDataContainer(`
+      [0///////////]
+      [1][3//][6///]
+      [2]     [9/]
+    `)!;
+
+    const items = container.getLevels();
+
+    return {
+      map: new CollapsedMap(
+        new Map([
+          [items[1][0], { items: [items[1][0], items[2][0]], collapsed: false }],
+          [items[1][2], { items: [items[1][2], items[2][1]], collapsed: true }],
+        ])
+      ),
+      items,
+    };
+  }
+
+  it('collapses and expands single item', () => {
+    const { map: collapsedMap, items } = getMap();
+    let newMap = collapsedMap.setCollapsedStatus(items[1][0], true);
+    expect(collapsedMap.get(items[1][0])?.collapsed).toBe(false);
+    expect(newMap.get(items[1][0])?.collapsed).toBe(true);
+
+    newMap = collapsedMap.setCollapsedStatus(items[1][2], false);
+    expect(collapsedMap.get(items[1][2])?.collapsed).toBe(true);
+    expect(newMap.get(items[1][2])?.collapsed).toBe(false);
+  });
+
+  it('collapses and expands all items', () => {
+    const { map: collapsedMap, items } = getMap();
+    let newMap = collapsedMap.setAllCollapsedStatus(true);
+    expect(collapsedMap.get(items[1][0])?.collapsed).toBe(false);
+    expect(newMap.get(items[1][0])?.collapsed).toBe(true);
+    expect(newMap.get(items[1][2])?.collapsed).toBe(true);
+
+    newMap = collapsedMap.setAllCollapsedStatus(false);
+    expect(collapsedMap.get(items[1][2])?.collapsed).toBe(true);
+    expect(newMap.get(items[1][0])?.collapsed).toBe(false);
+    expect(newMap.get(items[1][2])?.collapsed).toBe(false);
   });
 });

packages/grafana-flamegraph/src/FlameGraph/dataTransform.ts

@@ -91,7 +91,7 @@ export function nestedSetToLevels(
     levels[currentLevel].push(newItem);
   }
 
-  const collapsedMapContainer = new CollapsedMapContainer(options?.collapsingThreshold);
+  const collapsedMapContainer = new CollapsedMapBuilder(options?.collapsingThreshold);
   if (options?.collapsing) {
     // We collapse similar items here, where it seems like parent and child are the same thing and so the distinction
     // isn't that important. We create a map of items that should be collapsed together. We need to do it with complete
@@ -99,11 +99,66 @@ export function nestedSetToLevels(
     collapsedMapContainer.addTree(levels[0][0]);
   }
 
-  return [levels, uniqueLabels, collapsedMapContainer.getMap()];
+  return [levels, uniqueLabels, collapsedMapContainer.getCollapsedMap()];
 }
 
-export type CollapsedMap = Map<LevelItem, CollapseConfig>;
-export class CollapsedMapContainer {
+/**
+ * Small wrapper around the map of items that should be visually collapsed in the flame graph. Reason this is a wrapper
+ * is that we want to make sure that when this is in the state we don't update the map directly but create a new map
+ * and to have a place for the methods to collapse/expand either single item or all the items.
+ */
+export class CollapsedMap {
+  // The levelItem used as a key is the item that will always be rendered in the flame graph. The config.items are all
+  // the items that are in the group and if the config.collapsed is true they will be hidden.
+  private map: Map<LevelItem, CollapseConfig> = new Map();
+
+  constructor(map?: Map<LevelItem, CollapseConfig>) {
+    this.map = map || new Map();
+  }
+
+  get(item: LevelItem) {
+    return this.map.get(item);
+  }
+
+  keys() {
+    return this.map.keys();
+  }
+
+  values() {
+    return this.map.values();
+  }
+
+  size() {
+    return this.map.size;
+  }
+
+  setCollapsedStatus(item: LevelItem, collapsed: boolean) {
+    const newMap = new Map(this.map);
+    const collapsedConfig = this.map.get(item)!;
+    const newConfig = { ...collapsedConfig, collapsed };
+    for (const item of collapsedConfig.items) {
+      newMap.set(item, newConfig);
+    }
+    return new CollapsedMap(newMap);
+  }
+
+  setAllCollapsedStatus(collapsed: boolean) {
+    const newMap = new Map(this.map);
+    for (const item of this.map.keys()) {
+      const collapsedConfig = this.map.get(item)!;
+      const newConfig = { ...collapsedConfig, collapsed };
+      newMap.set(item, newConfig);
+    }
+
+    return new CollapsedMap(newMap);
+  }
+}
+
+/**
+ * Similar to CollapsedMap but this one is mutable and used during transformation of the dataFrame data into structure
+ * we use for rendering. This should not be passed to the React components.
+ */
+export class CollapsedMapBuilder {
   private map = new Map();
   private threshold = 0.99;
 
@@ -128,10 +183,10 @@ export class CollapsedMapContainer {
     }
   }
 
+  // The heuristics here is pretty simple right now. Just check if it's single child and if we are within threshold.
+  // We assume items with small self just aren't too important while we cannot really collapse items with siblings
+  // as it's not clear what to do with said sibling.
   addItem(item: LevelItem, parent?: LevelItem) {
-    // The heuristics here is pretty simple right now. Just check if it's single child and if we are within threshold.
-    // We assume items with small self just aren't too important while we cannot really collapse items with siblings
-    // as it's not clear what to do with said sibling.
     if (parent && item.value > parent.value * this.threshold && parent.children.length === 1) {
       if (this.map.has(parent)) {
         const config = this.map.get(parent)!;
@@ -145,8 +200,8 @@ export class CollapsedMapContainer {
     }
   }
 
-  getMap() {
-    return new Map(this.map);
+  getCollapsedMap() {
+    return new CollapsedMap(this.map);
   }
 }
 
@@ -225,7 +280,7 @@ export class FlameGraphDataContainer {
 
   private levels: LevelItem[][] | undefined;
   private uniqueLabelsMap: Record<string, LevelItem[]> | undefined;
-  private collapsedMap: Map<LevelItem, CollapseConfig> | undefined;
+  private collapsedMap: CollapsedMap | undefined;
 
   constructor(data: DataFrame, options: Options, theme: GrafanaTheme2 = createTheme()) {
     this.data = data;

packages/grafana-flamegraph/src/FlameGraphContainer.test.tsx

@@ -1,6 +1,6 @@
 import { render, screen, waitFor } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
-import React from 'react';
+import React, { useCallback } from 'react';
 
 import { createDataFrame, createTheme } from '@grafana/data';
 
@@ -29,7 +29,8 @@ describe('FlameGraphContainer', () => {
       },
     };
 
-    return <FlameGraphContainer data={flameGraphData} getTheme={() => createTheme({ colors: { mode: 'dark' } })} />;
+    const getTheme = useCallback(() => createTheme({ colors: { mode: 'dark' } }), []);
+    return <FlameGraphContainer data={flameGraphData} getTheme={getTheme} />;
   };
 
   it('should render without error', async () => {

packages/grafana-flamegraph/src/FlameGraphContainer.tsx

@@ -8,7 +8,7 @@ import { ThemeContext } from '@grafana/ui';
 
 import FlameGraph from './FlameGraph/FlameGraph';
 import { GetExtraContextMenuButtonsFunction } from './FlameGraph/FlameGraphContextMenu';
-import { FlameGraphDataContainer } from './FlameGraph/dataTransform';
+import { CollapsedMap, FlameGraphDataContainer } from './FlameGraph/dataTransform';
 import FlameGraphHeader from './FlameGraphHeader';
 import FlameGraphTopTableContainer from './TopTable/FlameGraphTopTableContainer';
 import { MIN_WIDTH_TO_SHOW_BOTH_TOPTABLE_AND_FLAMEGRAPH } from './constants';
@@ -98,14 +98,17 @@ const FlameGraphContainer = ({
   const [textAlign, setTextAlign] = useState<TextAlign>('left');
   // This is a label of the item because in sandwich view we group all items by label and present a merged graph
   const [sandwichItem, setSandwichItem] = useState<string>();
+  const [collapsedMap, setCollapsedMap] = useState(new CollapsedMap());
 
-  const theme = getTheme();
-
+  const theme = useMemo(() => getTheme(), [getTheme]);
   const dataContainer = useMemo((): FlameGraphDataContainer | undefined => {
     if (!data) {
       return;
     }
-    return new FlameGraphDataContainer(data, { collapsing: !disableCollapsing }, theme);
+
+    const container = new FlameGraphDataContainer(data, { collapsing: !disableCollapsing }, theme);
+    setCollapsedMap(container.getCollapsedMap());
+    return container;
   }, [data, theme, disableCollapsing]);
   const [colorScheme, setColorScheme] = useColorScheme(dataContainer);
   const styles = getStyles(theme);
@@ -179,6 +182,8 @@ const FlameGraphContainer = ({
       getExtraContextMenuButtons={getExtraContextMenuButtons}
       selectedView={selectedView}
       search={search}
+      collapsedMap={collapsedMap}
+      setCollapsedMap={setCollapsedMap}
     />
   );
 
@@ -250,6 +255,8 @@ const FlameGraphContainer = ({
             extraHeaderElements={extraHeaderElements}
             vertical={vertical}
             isDiffMode={dataContainer.isDiffFlamegraph()}
+            setCollapsedMap={setCollapsedMap}
+            collapsedMap={collapsedMap}
           />
         )}
 

packages/grafana-flamegraph/src/FlameGraphHeader.test.tsx

@@ -3,6 +3,7 @@ import { render, screen } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
 import React from 'react';
 
+import { CollapsedMap } from './FlameGraph/dataTransform';
 import FlameGraphHeader from './FlameGraphHeader';
 import { ColorScheme, SelectedView } from './types';
 
@@ -28,6 +29,8 @@ describe('FlameGraphHeader', () => {
         onColorSchemeChange={onSchemeChange}
         stickyHeader={false}
         isDiffMode={false}
+        setCollapsedMap={() => {}}
+        collapsedMap={new CollapsedMap()}
         {...props}
       />
     );

packages/grafana-flamegraph/src/FlameGraphHeader.tsx

@@ -4,9 +4,10 @@ import useDebounce from 'react-use/lib/useDebounce';
 import usePrevious from 'react-use/lib/usePrevious';
 
 import { GrafanaTheme2, SelectableValue } from '@grafana/data';
-import { Button, Dropdown, Input, Menu, RadioButtonGroup, useStyles2 } from '@grafana/ui';
+import { Button, ButtonGroup, Dropdown, Input, Menu, RadioButtonGroup, useStyles2 } from '@grafana/ui';
 
 import { byPackageGradient, byValueGradient, diffColorBlindGradient, diffDefaultGradient } from './FlameGraph/colors';
+import { CollapsedMap } from './FlameGraph/dataTransform';
 import { MIN_WIDTH_TO_SHOW_BOTH_TOPTABLE_AND_FLAMEGRAPH } from './constants';
 import { ColorScheme, ColorSchemeDiff, SelectedView, TextAlign } from './types';
 
@@ -25,6 +26,8 @@ type Props = {
   stickyHeader: boolean;
   vertical?: boolean;
   isDiffMode: boolean;
+  setCollapsedMap: (collapsedMap: CollapsedMap) => void;
+  collapsedMap: CollapsedMap;
 
   extraHeaderElements?: React.ReactNode;
 };
@@ -45,6 +48,8 @@ const FlameGraphHeader = ({
   extraHeaderElements,
   vertical,
   isDiffMode,
+  setCollapsedMap,
+  collapsedMap,
 }: Props) => {
   const styles = useStyles2(getStyles);
   const [localSearch, setLocalSearch] = useSearchInput(search, setSearch);
@@ -94,6 +99,32 @@ const FlameGraphHeader = ({
           />
         )}
         <ColorSchemeButton value={colorScheme} onChange={onColorSchemeChange} isDiffMode={isDiffMode} />
+        <ButtonGroup className={styles.buttonSpacing}>
+          <Button
+            variant={'secondary'}
+            fill={'outline'}
+            size={'sm'}
+            tooltip={'Expand all groups'}
+            onClick={() => {
+              setCollapsedMap(collapsedMap.setAllCollapsedStatus(false));
+            }}
+            aria-label={'Expand all groups'}
+            icon={'angle-double-down'}
+            disabled={selectedView === SelectedView.TopTable}
+          />
+          <Button
+            variant={'secondary'}
+            fill={'outline'}
+            size={'sm'}
+            tooltip={'Collapse all groups'}
+            onClick={() => {
+              setCollapsedMap(collapsedMap.setAllCollapsedStatus(true));
+            }}
+            aria-label={'Collapse all groups'}
+            icon={'angle-double-up'}
+            disabled={selectedView === SelectedView.TopTable}
+          />
+        </ButtonGroup>
         <RadioButtonGroup<TextAlign>
           size="sm"
           disabled={selectedView === SelectedView.TopTable}