Flamegraph: Fix theme propagation (#76064)

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

@@ -1,7 +1,7 @@
 import { fireEvent, render, screen } from '@testing-library/react';
 import React from 'react';
 
-import { createDataFrame, createTheme } from '@grafana/data';
+import { createDataFrame } from '@grafana/data';
 
 import { ColorScheme } from '../types';
 
@@ -48,7 +48,6 @@ describe('FlameGraph', () => {
         onFocusPillClick={onFocusPillClick}
         onSandwichPillClick={onSandwichPillClick}
         colorScheme={ColorScheme.ValueBased}
-        getTheme={() => createTheme({ colors: { mode: 'dark' } })}
       />
     );
     return {

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

@@ -20,7 +20,6 @@ import { css } from '@emotion/css';
 import React, { MouseEvent as ReactMouseEvent, useCallback, useEffect, useMemo, useRef, useState } from 'react';
 import { useMeasure } from 'react-use';
 
-import { GrafanaTheme2 } from '@grafana/data';
 import { Icon } from '@grafana/ui';
 
 import { PIXELS_PER_LEVEL } from '../constants';
@@ -48,7 +47,6 @@ type Props = {
   onFocusPillClick: () => void;
   onSandwichPillClick: () => void;
   colorScheme: ColorScheme | ColorSchemeDiff;
-  getTheme: () => GrafanaTheme2;
 };
 
 const FlameGraph = ({
@@ -66,7 +64,6 @@ const FlameGraph = ({
   onFocusPillClick,
   onSandwichPillClick,
   colorScheme,
-  getTheme,
 }: Props) => {
   const styles = getStyles();
 
@@ -108,7 +105,6 @@ const FlameGraph = ({
     totalColorTicks: data.isDiffFlamegraph() ? totalProfileTicks : totalViewTicks,
     totalTicksRight: totalProfileTicksRight,
     wrapperWidth,
-    getTheme,
   });
 
   const onGraphClick = useCallback(
@@ -186,7 +182,6 @@ const FlameGraph = ({
   return (
     <div className={styles.graph}>
       <FlameGraphMetadata
-        getTheme={getTheme}
         data={data}
         focusedItem={focusedItemData}
         sandwichedLabel={sandwichItem}
@@ -220,13 +215,7 @@ const FlameGraph = ({
           />
         </div>
       </div>
-      <FlameGraphTooltip
+      <FlameGraphTooltip position={mousePosition} item={tooltipItem} data={data} totalTicks={totalViewTicks} />
-        getTheme={getTheme}
-        position={mousePosition}
-        item={tooltipItem}
-        data={data}
-        totalTicks={totalViewTicks}
-      />
       {clickedItemData && (
         <FlameGraphContextMenu
           itemData={clickedItemData}

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

@@ -2,8 +2,6 @@ import { render, screen } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
 import React from 'react';
 
-import { createTheme } from '@grafana/data';
-
 import FlameGraphMetadata from './FlameGraphMetadata';
 import { textToDataContainer } from './testHelpers';
 
@@ -23,7 +21,6 @@ function setup(props: Partial<React.ComponentProps<typeof FlameGraphMetadata>> =
       totalTicks={17}
       onFocusPillClick={onFocusPillClick}
       onSandwichPillClick={onSandwichPillClick}
-      getTheme={() => createTheme({ colors: { mode: 'dark' } })}
       {...props}
     />
   );

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

@@ -2,7 +2,7 @@ import { css } from '@emotion/css';
 import React, { ReactNode } from 'react';
 
 import { getValueFormat, GrafanaTheme2 } from '@grafana/data';
-import { Icon, IconButton } from '@grafana/ui';
+import { Icon, IconButton, useStyles2 } from '@grafana/ui';
 
 import { ClickedItemData } from '../types';
 
@@ -15,12 +15,11 @@ type Props = {
   onSandwichPillClick: () => void;
   focusedItem?: ClickedItemData;
   sandwichedLabel?: string;
-  getTheme: () => GrafanaTheme2;
 };
 
 const FlameGraphMetadata = React.memo(
-  ({ data, focusedItem, totalTicks, sandwichedLabel, onFocusPillClick, onSandwichPillClick, getTheme }: Props) => {
+  ({ data, focusedItem, totalTicks, sandwichedLabel, onFocusPillClick, onSandwichPillClick }: Props) => {
-    const styles = getStyles(getTheme());
+    const styles = useStyles2(getStyles);
     const parts: ReactNode[] = [];
     const ticksVal = getValueFormat('short')(totalTicks);
 

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

@@ -2,7 +2,7 @@ import { css } from '@emotion/css';
 import React from 'react';
 
 import { DisplayValue, getValueFormat, GrafanaTheme2 } from '@grafana/data';
-import { InteractiveTable, Portal, VizTooltipContainer } from '@grafana/ui';
+import { InteractiveTable, Portal, useStyles2, VizTooltipContainer } from '@grafana/ui';
 
 import { FlameGraphDataContainer, LevelItem } from './dataTransform';
 
@@ -11,11 +11,10 @@ type Props = {
   totalTicks: number;
   position?: { x: number; y: number };
   item?: LevelItem;
-  getTheme: () => GrafanaTheme2;
 };
 
-const FlameGraphTooltip = ({ data, item, totalTicks, position, getTheme }: Props) => {
+const FlameGraphTooltip = ({ data, item, totalTicks, position }: Props) => {
-  const styles = getStyles(getTheme());
+  const styles = useStyles2(getStyles);
 
   if (!(item && position)) {
     return null;

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

@@ -3,6 +3,7 @@ import { RefObject, useEffect, useMemo, useState } from 'react';
 import color from 'tinycolor2';
 
 import { GrafanaTheme2 } from '@grafana/data';
+import { useTheme2 } from '@grafana/ui';
 
 import {
   BAR_BORDER_WIDTH,
@@ -41,7 +42,6 @@ type RenderOptions = {
   totalTicksRight: number | undefined;
   colorScheme: ColorScheme | ColorSchemeDiff;
   focusedItemData?: ClickedItemData;
-  getTheme: () => GrafanaTheme2;
 };
 
 export function useFlameRender(options: RenderOptions) {
@@ -59,7 +59,6 @@ export function useFlameRender(options: RenderOptions) {
     totalTicksRight,
     colorScheme,
     focusedItemData,
-    getTheme,
   } = options;
   const foundLabels = useMemo(() => {
     if (search) {
@@ -79,7 +78,7 @@ export function useFlameRender(options: RenderOptions) {
   }, [search, data]);
 
   const ctx = useSetupCanvas(canvasRef, wrapperWidth, levels.length);
-  const theme = getTheme();
+  const theme = useTheme2();
 
   useEffect(() => {
     if (!ctx) {

packages/grafana-flamegraph/src/FlameGraphContainer.tsx

@@ -3,6 +3,7 @@ import React, { useCallback, useEffect, useMemo, useState } from 'react';
 import { useMeasure } from 'react-use';
 
 import { DataFrame, GrafanaTheme2 } from '@grafana/data';
+import { ThemeContext } from '@grafana/ui';
 
 import FlameGraph from './FlameGraph/FlameGraph';
 import { FlameGraphDataContainer } from './FlameGraph/dataTransform';
@@ -133,75 +134,76 @@ const FlameGraphContainer = ({
   }
 
   return (
-    <div ref={sizeRef} className={styles.container}>
+    // We add the theme context to bridge the gap if this is rendered in non grafana environment where the context
-      <FlameGraphHeader
+    // isn't already provided.
-        search={search}
+    <ThemeContext.Provider value={theme}>
-        setSearch={setSearch}
+      <div ref={sizeRef} className={styles.container}>
-        selectedView={selectedView}
+        <FlameGraphHeader
-        setSelectedView={(view) => {
+          search={search}
-          setSelectedView(view);
+          setSearch={setSearch}
-          onViewSelected?.(view);
+          selectedView={selectedView}
-        }}
+          setSelectedView={(view) => {
-        containerWidth={containerWidth}
+            setSelectedView(view);
-        onReset={() => {
+            onViewSelected?.(view);
-          resetFocus();
+          }}
-          resetSandwich();
+          containerWidth={containerWidth}
-        }}
+          onReset={() => {
-        textAlign={textAlign}
+            resetFocus();
-        onTextAlignChange={(align) => {
+            resetSandwich();
-          setTextAlign(align);
+          }}
-          onTextAlignSelected?.(align);
+          textAlign={textAlign}
-        }}
+          onTextAlignChange={(align) => {
-        showResetButton={Boolean(focusedItemData || sandwichItem)}
+            setTextAlign(align);
-        colorScheme={colorScheme}
+            onTextAlignSelected?.(align);
-        onColorSchemeChange={setColorScheme}
+          }}
-        stickyHeader={Boolean(stickyHeader)}
+          showResetButton={Boolean(focusedItemData || sandwichItem)}
-        extraHeaderElements={extraHeaderElements}
+          colorScheme={colorScheme}
-        vertical={vertical}
+          onColorSchemeChange={setColorScheme}
-        isDiffMode={Boolean(dataContainer.isDiffFlamegraph())}
+          stickyHeader={Boolean(stickyHeader)}
-        getTheme={getTheme}
+          extraHeaderElements={extraHeaderElements}
-      />
+          vertical={vertical}
+          isDiffMode={Boolean(dataContainer.isDiffFlamegraph())}
+        />
 
-      <div className={styles.body}>
+        <div className={styles.body}>
-        {selectedView !== SelectedView.FlameGraph && (
+          {selectedView !== SelectedView.FlameGraph && (
-          <FlameGraphTopTableContainer
+            <FlameGraphTopTableContainer
-            data={dataContainer}
+              data={dataContainer}
-            onSymbolClick={onSymbolClick}
+              onSymbolClick={onSymbolClick}
-            height={selectedView === SelectedView.TopTable ? 600 : undefined}
+              height={selectedView === SelectedView.TopTable ? 600 : undefined}
-            search={search}
+              search={search}
-            sandwichItem={sandwichItem}
+              sandwichItem={sandwichItem}
-            onSandwich={setSandwichItem}
+              onSandwich={setSandwichItem}
-            onSearch={setSearch}
+              onSearch={setSearch}
-            onTableSort={onTableSort}
+              onTableSort={onTableSort}
-            getTheme={getTheme}
+              vertical={vertical}
-            vertical={vertical}
+            />
-          />
+          )}
-        )}
 
-        {selectedView !== SelectedView.TopTable && (
+          {selectedView !== SelectedView.TopTable && (
-          <FlameGraph
+            <FlameGraph
-            getTheme={getTheme}
+              data={dataContainer}
-            data={dataContainer}
+              rangeMin={rangeMin}
-            rangeMin={rangeMin}
+              rangeMax={rangeMax}
-            rangeMax={rangeMax}
+              search={search}
-            search={search}
+              setRangeMin={setRangeMin}
-            setRangeMin={setRangeMin}
+              setRangeMax={setRangeMax}
-            setRangeMax={setRangeMax}
+              onItemFocused={(data) => setFocusedItemData(data)}
-            onItemFocused={(data) => setFocusedItemData(data)}
+              focusedItemData={focusedItemData}
-            focusedItemData={focusedItemData}
+              textAlign={textAlign}
-            textAlign={textAlign}
+              sandwichItem={sandwichItem}
-            sandwichItem={sandwichItem}
+              onSandwich={(label: string) => {
-            onSandwich={(label: string) => {
+                resetFocus();
-              resetFocus();
+                setSandwichItem(label);
-              setSandwichItem(label);
+              }}
-            }}
+              onFocusPillClick={resetFocus}
-            onFocusPillClick={resetFocus}
+              onSandwichPillClick={resetSandwich}
-            onSandwichPillClick={resetSandwich}
+              colorScheme={colorScheme}
-            colorScheme={colorScheme}
+            />
-          />
+          )}
-        )}
+        </div>
       </div>
-    </div>
+    </ThemeContext.Provider>
   );
 };
 

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

@@ -3,8 +3,6 @@ import { render, screen } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
 import React from 'react';
 
-import { createTheme } from '@grafana/data';
-
 import FlameGraphHeader from './FlameGraphHeader';
 import { ColorScheme, SelectedView } from './types';
 
@@ -28,7 +26,6 @@ describe('FlameGraphHeader', () => {
         showResetButton={true}
         colorScheme={ColorScheme.ValueBased}
         onColorSchemeChange={onSchemeChange}
-        getTheme={() => createTheme({ colors: { mode: 'dark' } })}
         stickyHeader={false}
         isDiffMode={false}
         {...props}

packages/grafana-flamegraph/src/FlameGraphHeader.tsx

@@ -4,7 +4,7 @@ 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 } from '@grafana/ui';
+import { Button, Dropdown, Input, Menu, RadioButtonGroup, useStyles2 } from '@grafana/ui';
 
 import { byPackageGradient, byValueGradient, diffColorBlindGradient, diffDefaultGradient } from './FlameGraph/colors';
 import { MIN_WIDTH_TO_SHOW_BOTH_TOPTABLE_AND_FLAMEGRAPH } from './constants';
@@ -25,7 +25,6 @@ type Props = {
   stickyHeader: boolean;
   vertical?: boolean;
   isDiffMode: boolean;
-  getTheme: () => GrafanaTheme2;
 
   extraHeaderElements?: React.ReactNode;
 };
@@ -46,9 +45,8 @@ const FlameGraphHeader = ({
   extraHeaderElements,
   vertical,
   isDiffMode,
-  getTheme,
 }: Props) => {
-  const styles = getStyles(getTheme(), stickyHeader);
+  const styles = useStyles2(getStyles, stickyHeader);
   const [localSearch, setLocalSearch] = useSearchInput(search, setSearch);
 
   const suffix =
@@ -95,12 +93,7 @@ const FlameGraphHeader = ({
             aria-label={'Reset focus and sandwich state'}
           />
         )}
-        <ColorSchemeButton
+        <ColorSchemeButton value={colorScheme} onChange={onColorSchemeChange} isDiffMode={isDiffMode} />
-          value={colorScheme}
-          onChange={onColorSchemeChange}
-          isDiffMode={isDiffMode}
-          getTheme={getTheme}
-        />
         <RadioButtonGroup<TextAlign>
           size="sm"
           disabled={selectedView === SelectedView.TopTable}
@@ -124,12 +117,11 @@ const FlameGraphHeader = ({
 type ColorSchemeButtonProps = {
   value: ColorScheme | ColorSchemeDiff;
   onChange: (colorScheme: ColorScheme | ColorSchemeDiff) => void;
-  getTheme: () => GrafanaTheme2;
   isDiffMode: boolean;
 };
 function ColorSchemeButton(props: ColorSchemeButtonProps) {
   // TODO: probably create separate getStyles
-  const styles = getStyles(props.getTheme(), false);
+  const styles = useStyles2(getStyles, false);
   let menu = (
     <Menu>
       <Menu.Item label="By package name" onClick={() => props.onChange(ColorScheme.PackageBased)} />

packages/grafana-flamegraph/src/TopTable/FlameGraphTopTableContainer.test.tsx

@@ -2,7 +2,7 @@ import { render, screen } from '@testing-library/react';
 import userEvents from '@testing-library/user-event';
 import React from 'react';
 
-import { createDataFrame, createTheme } from '@grafana/data';
+import { createDataFrame } from '@grafana/data';
 
 import { FlameGraphDataContainer } from '../FlameGraph/dataTransform';
 import { data } from '../FlameGraph/testData/dataNestedSet';
@@ -22,7 +22,6 @@ describe('FlameGraphTopTableContainer', () => {
         onSymbolClick={jest.fn()}
         onSearch={onSearch}
         onSandwich={onSandwich}
-        getTheme={() => createTheme({ colors: { mode: 'dark' } })}
       />
     );
 

packages/grafana-flamegraph/src/TopTable/FlameGraphTopTableContainer.tsx

@@ -18,6 +18,8 @@ import {
   TableCustomCellOptions,
   TableFieldOptions,
   TableSortByFieldState,
+  useStyles2,
+  useTheme2,
 } from '@grafana/ui';
 
 import { FlameGraphDataContainer } from '../FlameGraph/dataTransform';
@@ -33,23 +35,11 @@ type Props = {
   onSearch: (str: string) => void;
   onSandwich: (str?: string) => void;
   onTableSort?: (sort: string) => void;
-  getTheme: () => GrafanaTheme2;
   vertical?: boolean;
 };
 
 const FlameGraphTopTableContainer = React.memo(
-  ({
+  ({ data, onSymbolClick, height, search, onSearch, sandwichItem, onSandwich, onTableSort, vertical }: Props) => {
-    data,
-    onSymbolClick,
-    height,
-    search,
-    onSearch,
-    sandwichItem,
-    onSandwich,
-    onTableSort,
-    getTheme,
-    vertical,
-  }: Props) => {
     const table = useMemo(() => {
       // Group the data by label, we show only one row per label and sum the values
       // TODO: should be by filename + funcName + linenumber?
@@ -73,7 +63,9 @@ const FlameGraphTopTableContainer = React.memo(
     // so we don't show potentially thousands of rows at once which can hinder performance (the table is virtualized
     // so with some max height it handles it fine)
     const tableHeight = vertical ? Math.min(Object.keys(table).length * rowHeight, 800) : 0;
-    const styles = getStyles(tableHeight, getTheme());
+
+    const styles = useStyles2(getStyles, tableHeight);
+    const theme = useTheme2();
 
     const [sort, setSort] = useState<TableSortByFieldState[]>([{ displayName: 'Self', desc: true }]);
 
@@ -92,7 +84,7 @@ const FlameGraphTopTableContainer = React.memo(
               onSymbolClick,
               onSearch,
               onSandwich,
-              getTheme,
+              theme,
               search,
               sandwichItem
             );
@@ -126,7 +118,7 @@ function buildTableDataFrame(
   onSymbolClick: (str: string) => void,
   onSearch: (str: string) => void,
   onSandwich: (str?: string) => void,
-  getTheme: () => GrafanaTheme2,
+  theme: GrafanaTheme2,
   search?: string,
   sandwichItem?: string
 ): DataFrame {
@@ -218,7 +210,7 @@ function buildTableDataFrame(
       overrides: [],
     },
     replaceVariables: (value: string) => value,
-    theme: getTheme(),
+    theme,
   });
 
   return dataFrames[0];
@@ -327,7 +319,7 @@ function ActionCell(props: ActionCellProps) {
   );
 }
 
-const getStyles = (height: number, theme: GrafanaTheme2) => {
+const getStyles = (theme: GrafanaTheme2, height: number) => {
   return {
     topTableContainer: css`
       label: topTableContainer;