diff --git a/packages/grafana-flamegraph/src/CallTree/FlameGraphCallTreeContainer.tsx b/packages/grafana-flamegraph/src/CallTree/FlameGraphCallTreeContainer.tsx index 05eebe1056b..0a304e1d435 100644 --- a/packages/grafana-flamegraph/src/CallTree/FlameGraphCallTreeContainer.tsx +++ b/packages/grafana-flamegraph/src/CallTree/FlameGraphCallTreeContainer.tsx @@ -1,5 +1,5 @@ import { css, cx } from '@emotion/css'; -import { memo, useEffect, useMemo, useState } from 'react'; +import { memo, useEffect, useMemo, useRef, useState } from 'react'; import { useTable, useSortBy, useExpanded, Column, Row, UseExpandedRowProps } from 'react-table'; import AutoSizer from 'react-virtualized-auto-sizer'; import { useDebounce, usePrevious } from 'react-use'; @@ -40,6 +40,9 @@ const FlameGraphCallTreeContainer = memo( // Color scheme state const [colorScheme, setColorScheme] = useState(initialColorScheme); + // Focus state - track which node is focused + const [focusedNodeId, setFocusedNodeId] = useState(undefined); + // Update color scheme when prop changes useEffect(() => { setColorScheme(initialColorScheme); @@ -60,11 +63,38 @@ const FlameGraphCallTreeContainer = memo( const effectiveMatchedLabels = searchMatchedLabels || matchedLabels; // Build and filter nodes - const { nodes, matchingIds } = useMemo(() => { + const { nodes, matchingIds, focusedNode } = useMemo(() => { const allNodes = buildAllCallTreeNodes(data); - const { visibleNodes, matchingNodeIds } = filterCallTree(allNodes, effectiveMatchedLabels); - return { nodes: visibleNodes, matchingIds: matchingNodeIds }; - }, [data, effectiveMatchedLabels]); + + // If there's a focused node, find it and use its subtree + let nodesToUse = allNodes; + let focused: CallTreeNode | undefined; + + if (focusedNodeId) { + // Find the focused node in the tree + const findNode = (nodes: CallTreeNode[], id: string): CallTreeNode | undefined => { + for (const node of nodes) { + if (node.id === id) { + return node; + } + if (node.subRows) { + const found = findNode(node.subRows, id); + if (found) return found; + } + } + return undefined; + }; + + focused = findNode(allNodes, focusedNodeId); + if (focused) { + // Use the focused node as the root + nodesToUse = [focused]; + } + } + + const { visibleNodes, matchingNodeIds } = filterCallTree(nodesToUse, effectiveMatchedLabels); + return { nodes: visibleNodes, matchingIds: matchingNodeIds, focusedNode: focused }; + }, [data, effectiveMatchedLabels, focusedNodeId]); // Calculate expanded state based on search const calculatedExpanded = useMemo(() => { @@ -75,6 +105,13 @@ const FlameGraphCallTreeContainer = memo( firstNode: nodes[0]?.id, expandedKeys: Object.keys(baseExpanded).filter(k => baseExpanded[k]), }); + + // If there's a focused node, expand it using react-table's row ID + // The focused node becomes the root, so its react-table ID is "0" + if (focusedNodeId && nodes.length > 0 && nodes[0].id === focusedNodeId && nodes[0].hasChildren) { + baseExpanded['0'] = true; + } + if (effectiveMatchedLabels && effectiveMatchedLabels.size > 0) { const matchExpanded = getExpandedStateForMatches(nodes, matchingIds); console.log('Adding match expansion:', { @@ -86,20 +123,47 @@ const FlameGraphCallTreeContainer = memo( return { ...baseExpanded, ...matchExpanded }; } return baseExpanded; - }, [nodes, effectiveMatchedLabels, matchingIds]); + }, [nodes, effectiveMatchedLabels, matchingIds, focusedNodeId]); // Create a key that changes when expansion should reset, forcing table remount const tableKey = useMemo(() => { // Include matchingIds in the key so table remounts when search results change const expandedKeys = Object.keys(calculatedExpanded).filter(k => calculatedExpanded[k]).sort().join(','); - return `table-${expandedKeys}`; - }, [calculatedExpanded]); + const focusPart = focusedNodeId ? `-focus-${focusedNodeId}` : ''; + return `table-${expandedKeys}${focusPart}`; + }, [calculatedExpanded, focusedNodeId]); + + // Callback to handle recursive expand/collapse - will be populated after tableInstance is created + const handleToggleExpandRef = useRef<((node: CallTreeNode, rowId: string, isExpanded: boolean) => void) | null>(null); // Define columns const columns: Column[] = useMemo(() => { if (data.isDiffFlamegraph()) { return [ + { + Header: '', + id: 'actions', + Cell: ({ row }: { row: Row & UseExpandedRowProps }) => ( + { + setFocusedNodeId(nodeId); + setLocalSearch(''); + }} + onToggleExpand={(isExpanded) => { + if (handleToggleExpandRef.current) { + handleToggleExpandRef.current(row.original, row.id, isExpanded); + } + }} + isFocusedRoot={row.id === '0'} + styles={styles} + /> + ), + width: 50, + minWidth: 50, + disableSortBy: true, + }, { Header: 'Function', accessor: 'label', @@ -157,6 +221,29 @@ const FlameGraphCallTreeContainer = memo( ]; } else { return [ + { + Header: '', + id: 'actions', + Cell: ({ row }: { row: Row & UseExpandedRowProps }) => ( + { + setFocusedNodeId(nodeId); + setLocalSearch(''); + }} + onToggleExpand={(isExpanded) => { + if (handleToggleExpandRef.current) { + handleToggleExpandRef.current(row.original, row.id, isExpanded); + } + }} + isFocusedRoot={row.id === '0'} + styles={styles} + /> + ), + width: 50, + minWidth: 50, + disableSortBy: true, + }, { Header: 'Function', accessor: 'label', @@ -224,7 +311,10 @@ const FlameGraphCallTreeContainer = memo( }, ]; } - }, [data, effectiveMatchedLabels, matchingIds, onSymbolClick, colorScheme, theme, styles]); + }, [data, effectiveMatchedLabels, matchingIds, onSymbolClick, colorScheme, theme, styles, setLocalSearch]); + + // toggleRowExpanded is used in the Cell renderers but doesn't need to be in the dependencies + // because it's accessed at render time, not definition time // Setup table instance with expand and sort // Using initialState - expansion resets when data changes so new initialState takes effect @@ -244,7 +334,22 @@ const FlameGraphCallTreeContainer = memo( useExpanded ); - const { getTableProps, getTableBodyProps, headerGroups, rows, prepareRow, state, toggleAllRowsExpanded } = tableInstance; + const { getTableProps, getTableBodyProps, headerGroups, rows, prepareRow, state, toggleAllRowsExpanded, toggleRowExpanded } = tableInstance; + + // Set up the recursive expand/collapse handler now that we have toggleRowExpanded + // We need to find the actual row ID that react-table uses, which may differ from node.id + handleToggleExpandRef.current = (node: CallTreeNode, rowId: string, isExpanded: boolean) => { + const toggleRecursive = (rId: string, n: CallTreeNode, expand: boolean) => { + toggleRowExpanded(rId, expand); + if (n.subRows) { + n.subRows.forEach((child, index) => { + const childRowId = `${rId}.${index}`; + toggleRecursive(childRowId, child, expand); + }); + } + }; + toggleRecursive(rowId, node, !isExpanded); + }; console.log('Table state:', { visibleRows: rows.length, @@ -290,6 +395,21 @@ const FlameGraphCallTreeContainer = memo( /> + {focusedNode && ( +
+
+ )} +
@@ -364,11 +484,12 @@ const FlameGraphCallTreeContainer = memo( {row.cells.map((cell) => { const { key: cellKey, ...cellProps } = cell.getCellProps(); const isValueColumn = cell.column.id === 'self' || cell.column.id === 'total'; + const isActionsColumn = cell.column.id === 'actions'; return ( {cell.render('Cell', { rowIndex })} @@ -490,6 +611,57 @@ function getRowBackgroundColor( // Cell Components +function ActionsCell({ + row, + onFocus, + onToggleExpand, + isFocusedRoot, + styles, +}: { + row: Row & UseExpandedRowProps; + onFocus: (nodeId: string) => void; + onToggleExpand: (isExpanded: boolean) => void; + isFocusedRoot: boolean; + styles: any; +}) { + const hasChildren = row.original.hasChildren; + const isExpanded = row.isExpanded; + + return ( +
+
+ ); +} + function FunctionCellWithExpander({ row, value, @@ -874,6 +1046,41 @@ function getStyles(theme: GrafanaTheme2) { minWidth: '2px', borderRadius: '2px', }), + actionsCell: css({ + display: 'flex', + alignItems: 'center', + gap: theme.spacing(0.5), + height: '20px', + }), + actionsColumnCell: css({ + backgroundColor: theme.colors.background.secondary, + '&:hover': { + backgroundColor: theme.colors.background.secondary, + }, + }), + actionButton: css({ + padding: 0, + minWidth: 'auto', + height: '20px', + color: theme.colors.text.secondary, + '&:hover': { + color: theme.colors.text.primary, + }, + }), + focusedItem: css({ + display: 'flex', + alignItems: 'center', + gap: theme.spacing(1), + padding: `${theme.spacing(0.5)} ${theme.spacing(1)}`, + backgroundColor: theme.colors.background.secondary, + borderRadius: theme.shape.radius.default, + border: `1px solid ${theme.colors.border.weak}`, + }), + focusedItemLabel: css({ + fontSize: theme.typography.bodySmall.fontSize, + color: theme.colors.text.primary, + fontWeight: theme.typography.fontWeightMedium, + }), }; }