FlameGraph: adds ability to add context menu items (#81675)

* pyroscope: adds ability to add context menu items

* moves things around

* removes console.log

* improvements

* Change the extra context button API shape

* Add test

* lint

---------

Co-authored-by: Andrej Ocenas <mr.ocenas@gmail.com>

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

@@ -3,7 +3,7 @@ import React from 'react';
 
 import { createDataFrame } from '@grafana/data';
 
-import { ColorScheme } from '../types';
+import { ColorScheme, SelectedView } from '../types';
 
 import FlameGraph from './FlameGraph';
 import { FlameGraphDataContainer } from './dataTransform';
@@ -23,7 +23,7 @@ jest.mock('react-use', () => {
 });
 
 describe('FlameGraph', () => {
-  function setup() {
+  function setup(props?: Partial<React.ComponentProps<typeof FlameGraph>>) {
     const flameGraphData = createDataFrame(data);
     const container = new FlameGraphDataContainer(flameGraphData, { collapsing: true });
 
@@ -47,6 +47,9 @@ describe('FlameGraph', () => {
         onFocusPillClick={onFocusPillClick}
         onSandwichPillClick={onSandwichPillClick}
         colorScheme={ColorScheme.ValueBased}
+        selectedView={SelectedView.FlameGraph}
+        search={''}
+        {...props}
       />
     );
     return {
@@ -80,18 +83,31 @@ describe('FlameGraph', () => {
     expect(screen.getByText('16.5 Bil | 16.5 Bil samples (Count)')).toBeDefined();
   });
 
-  it('should render context menu', async () => {
+  it('should render context menu + extra items', async () => {
     const event = new MouseEvent('click', { bubbles: true });
     Object.defineProperty(event, 'offsetX', { get: () => 10 });
     Object.defineProperty(event, 'offsetY', { get: () => 10 });
     Object.defineProperty(HTMLCanvasElement.prototype, 'clientWidth', { configurable: true, value: 500 });
 
-    setup();
+    setup({
+      getExtraContextMenuButtons: (clickedItemData, data, state) => {
+        expect(clickedItemData).toMatchObject({ posX: 0, posY: 0, label: 'total' });
+        expect(data.length).toEqual(1101);
+        expect(state).toEqual({
+          selectedView: SelectedView.FlameGraph,
+          isDiff: false,
+          search: '',
+          collapseConfig: undefined,
+        });
+        return [{ label: 'test extra item', icon: 'eye', onClick: () => {} }];
+      },
+    });
     const canvas = screen.getByTestId('flameGraph') as HTMLCanvasElement;
     expect(canvas).toBeInTheDocument();
     expect(screen.queryByTestId('contextMenu')).not.toBeInTheDocument();
 
     fireEvent(canvas, event);
     expect(screen.getByTestId('contextMenu')).toBeInTheDocument();
+    expect(screen.getByText('test extra item')).toBeInTheDocument();
   });
 });

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

@@ -22,9 +22,10 @@ import React, { useEffect, useState } from 'react';
 import { Icon } from '@grafana/ui';
 
 import { PIXELS_PER_LEVEL } from '../constants';
-import { ClickedItemData, ColorScheme, ColorSchemeDiff, TextAlign } from '../types';
+import { ClickedItemData, ColorScheme, ColorSchemeDiff, SelectedView, TextAlign } from '../types';
 
 import FlameGraphCanvas from './FlameGraphCanvas';
+import { GetExtraContextMenuButtonsFunction } from './FlameGraphContextMenu';
 import FlameGraphMetadata from './FlameGraphMetadata';
 import { CollapsedMap, FlameGraphDataContainer, LevelItem } from './dataTransform';
 
@@ -45,7 +46,10 @@ type Props = {
   onSandwichPillClick: () => void;
   colorScheme: ColorScheme | ColorSchemeDiff;
   showFlameGraphOnly?: boolean;
+  getExtraContextMenuButtons?: GetExtraContextMenuButtonsFunction;
   collapsing?: boolean;
+  selectedView: SelectedView;
+  search: string;
 };
 
 const FlameGraph = ({
@@ -64,7 +68,10 @@ const FlameGraph = ({
   onSandwichPillClick,
   colorScheme,
   showFlameGraphOnly,
+  getExtraContextMenuButtons,
   collapsing,
+  selectedView,
+  search,
 }: Props) => {
   const styles = getStyles();
 
@@ -122,7 +129,10 @@ const FlameGraph = ({
     showFlameGraphOnly,
     collapsedMap,
     setCollapsedMap,
+    getExtraContextMenuButtons,
     collapsing,
+    search,
+    selectedView,
   };
   const canvas = levelsCallers ? (
     <>

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

@@ -3,9 +3,9 @@ import React, { MouseEvent as ReactMouseEvent, useCallback, useEffect, useRef, u
 import { useMeasure } from 'react-use';
 
 import { PIXELS_PER_LEVEL } from '../constants';
-import { ClickedItemData, ColorScheme, ColorSchemeDiff, TextAlign } from '../types';
+import { ClickedItemData, ColorScheme, ColorSchemeDiff, SelectedView, TextAlign } from '../types';
 
-import FlameGraphContextMenu from './FlameGraphContextMenu';
+import FlameGraphContextMenu, { GetExtraContextMenuButtonsFunction } from './FlameGraphContextMenu';
 import FlameGraphTooltip from './FlameGraphTooltip';
 import { CollapseConfig, CollapsedMap, FlameGraphDataContainer, LevelItem } from './dataTransform';
 import { getBarX, useFlameRender } from './rendering';
@@ -37,6 +37,10 @@ type Props = {
   collapsedMap: CollapsedMap;
   setCollapsedMap: (collapsedMap: CollapsedMap) => void;
   collapsing?: boolean;
+  getExtraContextMenuButtons?: GetExtraContextMenuButtonsFunction;
+
+  selectedView: SelectedView;
+  search: string;
 };
 
 const FlameGraphCanvas = ({
@@ -61,6 +65,9 @@ const FlameGraphCanvas = ({
   collapsedMap,
   setCollapsedMap,
   collapsing,
+  getExtraContextMenuButtons,
+  selectedView,
+  search,
 }: Props) => {
   const styles = getStyles();
 
@@ -186,6 +193,7 @@ const FlameGraphCanvas = ({
       />
       {!showFlameGraphOnly && clickedItemData && (
         <FlameGraphContextMenu
+          data={data}
           itemData={clickedItemData}
           collapsing={collapsing}
           collapseConfig={collapsedMap.get(clickedItemData.item)}
@@ -214,6 +222,9 @@ const FlameGraphCanvas = ({
           }}
           allGroupsCollapsed={Array.from(collapsedMap.values()).every((i) => i.collapsed)}
           allGroupsExpanded={Array.from(collapsedMap.values()).every((i) => !i.collapsed)}
+          getExtraContextMenuButtons={getExtraContextMenuButtons}
+          selectedView={selectedView}
+          search={search}
         />
       )}
     </div>

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

@@ -1,12 +1,26 @@
 import React from 'react';
 
-import { MenuItem, MenuGroup, ContextMenu } from '@grafana/ui';
+import { DataFrame } from '@grafana/data';
+import { MenuItem, MenuGroup, ContextMenu, IconName } from '@grafana/ui';
 
-import { ClickedItemData } from '../types';
+import { ClickedItemData, SelectedView } from '../types';
 
-import { CollapseConfig } from './dataTransform';
+import { CollapseConfig, FlameGraphDataContainer } from './dataTransform';
+
+export type GetExtraContextMenuButtonsFunction = (
+  clickedItemData: ClickedItemData,
+  data: DataFrame,
+  state: { selectedView: SelectedView; isDiff: boolean; search: string; collapseConfig?: CollapseConfig }
+) => ExtraContextMenuButton[];
+
+export type ExtraContextMenuButton = {
+  label: string;
+  icon: IconName;
+  onClick: () => void;
+};
 
 type Props = {
+  data: FlameGraphDataContainer;
   itemData: ClickedItemData;
   onMenuItemClick: () => void;
   onItemFocus: () => void;
@@ -15,13 +29,17 @@ type Props = {
   onCollapseGroup: () => void;
   onExpandAllGroups: () => void;
   onCollapseAllGroups: () => void;
+  getExtraContextMenuButtons?: GetExtraContextMenuButtonsFunction;
   collapseConfig?: CollapseConfig;
   collapsing?: boolean;
   allGroupsCollapsed?: boolean;
   allGroupsExpanded?: boolean;
+  selectedView: SelectedView;
+  search: string;
 };
 
 const FlameGraphContextMenu = ({
+  data,
   itemData,
   onMenuItemClick,
   onItemFocus,
@@ -31,11 +49,21 @@ const FlameGraphContextMenu = ({
   onCollapseGroup,
   onExpandAllGroups,
   onCollapseAllGroups,
+  getExtraContextMenuButtons,
   collapsing,
   allGroupsExpanded,
   allGroupsCollapsed,
+  selectedView,
+  search,
 }: Props) => {
   function renderItems() {
+    const extraButtons =
+      getExtraContextMenuButtons?.(itemData, data.data, {
+        selectedView,
+        isDiff: data.isDiffFlamegraph(),
+        search,
+        collapseConfig,
+      }) || [];
     return (
       <>
         <MenuItem
@@ -63,7 +91,9 @@ const FlameGraphContextMenu = ({
             onMenuItemClick();
           }}
         />
-
+        {extraButtons.map(({ label, icon, onClick }) => {
+          return <MenuItem label={label} icon={icon} onClick={() => onClick()} key={label} />;
+        })}
         {collapsing && (
           <MenuGroup label={'Grouping'}>
             {collapseConfig ? (

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

@@ -272,7 +272,7 @@ export class FlameGraphDataContainer {
   }
 
   isDiffFlamegraph() {
-    return this.valueRightField && this.selfRightField;
+    return Boolean(this.valueRightField && this.selfRightField);
   }
 
   getLabel(index: number) {

packages/grafana-flamegraph/src/FlameGraphContainer.tsx

@@ -7,6 +7,7 @@ import { DataFrame, GrafanaTheme2 } from '@grafana/data';
 import { ThemeContext } from '@grafana/ui';
 
 import FlameGraph from './FlameGraph/FlameGraph';
+import { GetExtraContextMenuButtonsFunction } from './FlameGraph/FlameGraphContextMenu';
 import { FlameGraphDataContainer } from './FlameGraph/dataTransform';
 import FlameGraphHeader from './FlameGraphHeader';
 import FlameGraphTopTableContainer from './TopTable/FlameGraphTopTableContainer';
@@ -52,6 +53,11 @@ export type Props = {
    */
   extraHeaderElements?: React.ReactNode;
 
+  /**
+   * Extra buttons that will be shown in the context menu when user clicks on a Node.
+   */
+  getExtraContextMenuButtons?: GetExtraContextMenuButtonsFunction;
+
   /**
    * If true the flamegraph will be rendered on top of the table.
    */
@@ -80,6 +86,7 @@ const FlameGraphContainer = ({
   vertical,
   showFlameGraphOnly,
   disableCollapsing,
+  getExtraContextMenuButtons,
 }: Props) => {
   const [focusedItemData, setFocusedItemData] = useState<ClickedItemData>();
 
@@ -169,6 +176,9 @@ const FlameGraphContainer = ({
       colorScheme={colorScheme}
       showFlameGraphOnly={showFlameGraphOnly}
       collapsing={!disableCollapsing}
+      getExtraContextMenuButtons={getExtraContextMenuButtons}
+      selectedView={selectedView}
+      search={search}
     />
   );
 
@@ -238,7 +248,7 @@ const FlameGraphContainer = ({
             stickyHeader={Boolean(stickyHeader)}
             extraHeaderElements={extraHeaderElements}
             vertical={vertical}
-            isDiffMode={Boolean(dataContainer.isDiffFlamegraph())}
+            isDiffMode={dataContainer.isDiffFlamegraph()}
           />
         )}
 

packages/grafana-flamegraph/src/types.ts

@@ -1,5 +1,9 @@
 import { LevelItem } from './FlameGraph/dataTransform';
 
+export { type FlameGraphDataContainer } from './FlameGraph/dataTransform';
+
+export { type ExtraContextMenuButton } from './FlameGraph/FlameGraphContextMenu';
+
 export type ClickedItemData = {
   posX: number;
   posY: number;