ref(contextmenu): simplify the hook API (#35045)

* ref(contextmenu): split out context menu components

* ref(contextmenu): simplify the hook usage

* ref(contextmenu): remove an effect

* ref(profiling): add back callbacks

* feat(profiling): zoom to frame from call tree table (#35060)

* feat(profiling): add zoom into

* Update static/app/components/profiling/FrameStack/frameStackContextMenu.tsx

Co-authored-by: Tony Xiao <txiao@sentry.io>

* Update static/app/components/profiling/FrameStack/frameStackContextMenu.tsx

Co-authored-by: Tony Xiao <txiao@sentry.io>

* feat(profiling): import t

* feat(profiling): add keyboard nav to table (#35148)

Co-authored-by: Tony Xiao <txiao@sentry.io>

Co-authored-by: Tony Xiao <txiao@sentry.io>

static/app/components/profiling/FrameStack/frameStack.tsx

@@ -105,7 +105,12 @@ function FrameStack(props: FrameStackProps) {
         </li>
         <li style={{flex: '1 1 100%', cursor: 'ns-resize'}} onMouseDown={onMouseDown} />
       </FrameTabs>
-      <FrameStackTable {...props} roots={roots ?? []} referenceNode={selectedNode} />
+      <FrameStackTable
+        {...props}
+        roots={roots ?? []}
+        referenceNode={selectedNode}
+        canvasPoolManager={props.canvasPoolManager}
+      />
     </FrameDrawer>
   ) : null;
 }

static/app/components/profiling/FrameStack/frameStackContextMenu.tsx

@@ -0,0 +1,43 @@
+import {Fragment} from 'react';
+
+import {
+  ProfilingContextMenu,
+  ProfilingContextMenuGroup,
+  ProfilingContextMenuHeading,
+  ProfilingContextMenuItem,
+  ProfilingContextMenuLayer,
+} from 'sentry/components/profiling/ProfilingContextMenu/profilingContextMenu';
+import {t} from 'sentry/locale';
+import {useContextMenu} from 'sentry/utils/profiling/hooks/useContextMenu';
+
+interface FrameStackContextMenuProps {
+  contextMenu: ReturnType<typeof useContextMenu>;
+  onZoomIntoNodeClick: (evt: React.MouseEvent<HTMLDivElement>) => void;
+}
+
+export function FrameStackContextMenu(props: FrameStackContextMenuProps) {
+  return props.contextMenu.open ? (
+    <Fragment>
+      <ProfilingContextMenuLayer onClick={() => props.contextMenu.setOpen(false)} />
+      <ProfilingContextMenu
+        {...props.contextMenu.getMenuProps()}
+        style={{
+          position: 'absolute',
+          left: props.contextMenu.position?.left ?? -9999,
+          top: props.contextMenu.position?.top ?? -9999,
+          maxHeight: props.contextMenu.containerCoordinates?.height ?? 'auto',
+        }}
+      >
+        <ProfilingContextMenuGroup>
+          <ProfilingContextMenuHeading>{t('Flamegraph')}</ProfilingContextMenuHeading>
+          <ProfilingContextMenuItem
+            {...props.contextMenu.getMenuItemProps()}
+            onClick={props.onZoomIntoNodeClick}
+          >
+            {t('Scope view to this node')}
+          </ProfilingContextMenuItem>
+        </ProfilingContextMenuGroup>
+      </ProfilingContextMenu>
+    </Fragment>
+  ) : null;
+}

static/app/components/profiling/FrameStack/frameStackTable.tsx

@@ -4,12 +4,15 @@ import styled from '@emotion/styled';
 import {IconArrow} from 'sentry/icons';
 import {t} from 'sentry/locale';
 import space from 'sentry/styles/space';
+import {CanvasPoolManager} from 'sentry/utils/profiling/canvasScheduler';
 import {FlamegraphFrame} from 'sentry/utils/profiling/flamegraphFrame';
+import {useContextMenu} from 'sentry/utils/profiling/hooks/useContextMenu';
 import {useVirtualizedTree} from 'sentry/utils/profiling/hooks/useVirtualizedTree/useVirtualizedTree';
 import {VirtualizedTreeNode} from 'sentry/utils/profiling/hooks/useVirtualizedTree/VirtualizedTreeNode';
 import {FlamegraphRenderer} from 'sentry/utils/profiling/renderers/flamegraphRenderer';
 
 import {FrameCallersTableCell} from './frameStack';
+import {FrameStackContextMenu} from './frameStackContextMenu';
 import {FrameStackTableRow} from './frameStackTableRow';
 
 function makeSortFunction(
@@ -68,6 +71,7 @@ function makeSortFunction(
 }
 
 interface FrameStackTableProps {
+  canvasPoolManager: CanvasPoolManager;
   flamegraphRenderer: FlamegraphRenderer;
   referenceNode: FlamegraphFrame;
   roots: FlamegraphFrame[];
@@ -76,6 +80,7 @@ interface FrameStackTableProps {
 export function FrameStackTable({
   roots,
   flamegraphRenderer,
+  canvasPoolManager,
   referenceNode,
 }: FrameStackTableProps) {
   const scrollContainerRef = useRef<HTMLDivElement | null>(null);
@@ -90,8 +95,11 @@ export function FrameStackTable({
 
   const {
     items,
+    tabIndexKey,
     scrollContainerStyles,
     containerStyles,
+    handleRowClick,
+    handleRowKeyDown,
     handleExpandTreeNode,
     handleSortingChange,
     handleScroll,
@@ -116,6 +124,19 @@ export function FrameStackTable({
     [sort, direction, handleSortingChange]
   );
 
+  const [clickedContextMenuNode, setClickedContextMenuClose] =
+    useState<VirtualizedTreeNode<FlamegraphFrame> | null>(null);
+
+  const contextMenu = useContextMenu({container: scrollContainerRef.current});
+
+  const handleZoomIntoNodeClick = useCallback(() => {
+    if (!clickedContextMenuNode) {
+      return;
+    }
+
+    canvasPoolManager.dispatch('zoomIntoFrame', [clickedContextMenuNode.node]);
+  }, [canvasPoolManager, clickedContextMenuNode]);
+
   return (
     <FrameBar>
       <FrameCallersTable>
@@ -145,21 +166,34 @@ export function FrameStackTable({
             </TableHeaderButton>
           </FrameNameCell>
         </FrameCallersTableHeader>
+        <FrameStackContextMenu
+          onZoomIntoNodeClick={handleZoomIntoNodeClick}
+          contextMenu={contextMenu}
+        />
         <div
           ref={scrollContainerRef}
           style={scrollContainerStyles}
           onScroll={handleScroll}
+          onContextMenu={contextMenu.handleContextMenu}
         >
           <div style={containerStyles}>
             {items.map(r => {
               return (
                 <FrameStackTableRow
                   key={r.key}
+                  ref={n => (r.ref = n)}
                   node={r.item}
                   style={r.styles}
                   referenceNode={referenceNode}
+                  tabIndex={tabIndexKey === r.key ? 0 : 1}
                   flamegraphRenderer={flamegraphRenderer}
-                  handleExpandedClick={handleExpandTreeNode}
+                  onClick={() => handleRowClick(r.key)}
+                  onExpandClick={handleExpandTreeNode}
+                  onKeyDown={evt => handleRowKeyDown(r.key, evt)}
+                  onContextMenu={evt => {
+                    setClickedContextMenuClose(r.item);
+                    contextMenu.handleContextMenu(evt);
+                  }}
                 />
               );
             })}

static/app/components/profiling/FrameStack/frameStackTableRow.tsx

@@ -1,4 +1,4 @@
-import {useCallback, useMemo} from 'react';
+import {forwardRef, useCallback, useMemo} from 'react';
 import styled from '@emotion/styled';
 
 import {IconSettings, IconUser} from 'sentry/icons';
@@ -20,81 +20,109 @@ function computeRelativeWeight(base: number, value: number) {
 
 interface FrameStackTableRowProps {
   flamegraphRenderer: FlamegraphRenderer;
-  handleExpandedClick: (
+  node: VirtualizedTreeNode<FlamegraphFrame>;
+  onClick: React.MouseEventHandler<HTMLDivElement>;
+  onContextMenu: React.MouseEventHandler<HTMLDivElement>;
+  onExpandClick: (
     node: VirtualizedTreeNode<FlamegraphFrame>,
     opts?: {expandChildren: boolean}
   ) => void;
-  node: VirtualizedTreeNode<FlamegraphFrame>;
+  onKeyDown: React.KeyboardEventHandler<HTMLDivElement>;
   referenceNode: FlamegraphFrame;
   style: React.CSSProperties;
+  tabIndex: number;
 }
 
-export function FrameStackTableRow({
-  node,
-  flamegraphRenderer,
-  referenceNode,
-  handleExpandedClick,
-  style,
-}: FrameStackTableRowProps) {
-  const colorString = useMemo(() => {
-    return formatColorForFrame(node.node, flamegraphRenderer);
-  }, [node, flamegraphRenderer]);
-
-  const handleExpanding = useCallback(
-    (evt: React.MouseEvent) => {
-      handleExpandedClick(node, {expandChildren: evt.metaKey});
+export const FrameStackTableRow = forwardRef<HTMLDivElement, FrameStackTableRowProps>(
+  (
+    {
+      node,
+      flamegraphRenderer,
+      referenceNode,
+      onExpandClick,
+      onContextMenu,
+      tabIndex,
+      onKeyDown,
+      onClick,
+      style,
     },
-    [node, handleExpandedClick]
-  );
-
-  return (
-    <FrameCallersRow style={style}>
-      <FrameCallersTableCell textAlign="right">
-        {flamegraphRenderer.flamegraph.formatter(node.node.node.selfWeight)}
-        <Weight
-          weight={computeRelativeWeight(
-            referenceNode.node.totalWeight,
-            node.node.node.selfWeight
-          )}
-        />
-      </FrameCallersTableCell>
-      <FrameCallersTableCell noPadding textAlign="right">
-        <FrameWeightTypeContainer>
-          <FrameWeightContainer>
-            {flamegraphRenderer.flamegraph.formatter(node.node.node.totalWeight)}
-            <Weight
-              weight={computeRelativeWeight(
-                referenceNode.node.totalWeight,
-                node.node.node.totalWeight
-              )}
-            />
-          </FrameWeightContainer>
-          <FrameTypeIndicator>
-            {node.node.node.frame.is_application ? (
-              <IconUser size="xs" />
-            ) : (
-              <IconSettings size="xs" />
-            )}
-          </FrameTypeIndicator>
-        </FrameWeightTypeContainer>
-      </FrameCallersTableCell>
-      <FrameCallersTableCell
-        // We stretch this table to 100% width.
-        style={{paddingLeft: node.depth * 14 + 8, width: '100%'}}
+    ref
+  ) => {
+    const colorString = useMemo(() => {
+      return formatColorForFrame(node.node, flamegraphRenderer);
+    }, [node, flamegraphRenderer]);
+
+    const handleExpanding = useCallback(
+      (evt: React.MouseEvent) => {
+        evt.stopPropagation();
+        onExpandClick(node, {expandChildren: evt.metaKey});
+      },
+      [node, onExpandClick]
+    );
+
+    return (
+      <FrameCallersRow
+        ref={ref}
+        style={style}
+        onContextMenu={onContextMenu}
+        tabIndex={tabIndex}
+        isSelected={tabIndex === 0}
+        onKeyDown={onKeyDown}
+        onClick={onClick}
       >
-        <FrameNameContainer>
-          <FrameColorIndicator backgroundColor={colorString} />
-          <FrameChildrenIndicator onClick={handleExpanding} open={node.expanded}>
-            {node.node.children.length > 0 ? '\u203A' : null}
-          </FrameChildrenIndicator>
-          <FrameName>{node.node.frame.name}</FrameName>
-        </FrameNameContainer>
-      </FrameCallersTableCell>
-    </FrameCallersRow>
-  );
-}
+        <FrameCallersTableCell textAlign="right">
+          {flamegraphRenderer.flamegraph.formatter(node.node.node.selfWeight)}
+          <Weight
+            isSelected={tabIndex === 0}
+            weight={computeRelativeWeight(
+              referenceNode.node.totalWeight,
+              node.node.node.selfWeight
+            )}
+          />
+        </FrameCallersTableCell>
+        <FrameCallersTableCell noPadding textAlign="right">
+          <FrameWeightTypeContainer>
+            <FrameWeightContainer>
+              {flamegraphRenderer.flamegraph.formatter(node.node.node.totalWeight)}
+              <Weight
+                isSelected={tabIndex === 0}
+                weight={computeRelativeWeight(
+                  referenceNode.node.totalWeight,
+                  node.node.node.totalWeight
+                )}
+              />
+            </FrameWeightContainer>
+            <FrameTypeIndicator isSelected={tabIndex === 0}>
+              {node.node.node.frame.is_application ? (
+                <IconUser size="xs" />
+              ) : (
+                <IconSettings size="xs" />
+              )}
+            </FrameTypeIndicator>
+          </FrameWeightTypeContainer>
+        </FrameCallersTableCell>
+        <FrameCallersTableCell
+          // We stretch this table to 100% width.
+          style={{paddingLeft: node.depth * 14 + 8, width: '100%'}}
+        >
+          <FrameNameContainer>
+            <FrameColorIndicator backgroundColor={colorString} />
+            <FrameChildrenIndicator
+              tabIndex={-1}
+              onClick={handleExpanding}
+              open={node.expanded}
+            >
+              {node.node.children.length > 0 ? '\u203A' : null}
+            </FrameChildrenIndicator>
+            <FrameName>{node.node.frame.name}</FrameName>
+          </FrameNameContainer>
+        </FrameCallersTableCell>
+      </FrameCallersRow>
+    );
+  }
+);
 
-const Weight = styled((props: {weight: number}) => {
+const Weight = styled((props: {isSelected: boolean; weight: number}) => {
   const {weight, ...rest} = props;
   return (
     <div {...rest}>
@@ -105,23 +133,26 @@ const Weight = styled((props: {weight: number}) => {
 })`
   display: inline-block;
   min-width: 7ch;
-  color: ${props => props.theme.subText};
+  color: ${p => (p.isSelected ? p.theme.white : p.theme.subText)};
+  opacity: ${p => (p.isSelected ? 0.8 : 1)};
 `;
 
 const FrameWeightTypeContainer = styled('div')`
   display: flex;
   align-items: center;
   justify-content: flex-end;
+  position: relative;
 `;
 
-const FrameTypeIndicator = styled('div')`
+const FrameTypeIndicator = styled('div')<{isSelected: boolean}>`
   flex-shrink: 0;
   width: 26px;
   height: 12px;
   display: flex;
   align-items: center;
   justify-content: center;
-  color: ${p => p.theme.subText};
+  color: ${p => (p.isSelected ? p.theme.white : p.theme.subText)};
+  opacity: ${p => (p.isSelected ? 0.8 : 1)};
 `;
 
 const FrameWeightContainer = styled('div')`
@@ -144,12 +175,19 @@ const BackgroundWeightBar = styled('div')`
   width: 100%;
 `;
 
-const FrameCallersRow = styled('div')`
+const FrameCallersRow = styled('div')<{isSelected: boolean}>`
   display: flex;
   width: 100%;
 
+  background-color: ${p => (p.isSelected ? p.theme.blue300 : 'transparent')};
+  color: ${p => (p.isSelected ? p.theme.white : 'inherit')};
+
   &:hover {
-    background-color: ${p => p.theme.surface400};
+    background-color: ${p => (p.isSelected ? p.theme.blue300 : p.theme.blue100)};
+  }
+
+  &:focus {
+    outline: none;
   }
 `;
 

static/app/components/profiling/ProfilingContextMenu/profilingContextMenu.tsx

@@ -0,0 +1,212 @@
+import {forwardRef} from 'react';
+import styled from '@emotion/styled';
+
+import {IconCheckmark} from 'sentry/icons';
+import space from 'sentry/styles/space';
+
+interface MenuProps
+  extends React.DetailedHTMLProps<React.HTMLAttributes<HTMLDivElement>, HTMLDivElement> {
+  children: React.ReactNode;
+}
+
+const Menu = styled(
+  forwardRef((props: MenuProps, ref: React.Ref<HTMLDivElement> | undefined) => {
+    return <div ref={ref} role="menu" {...props} />;
+  })
+)`
+  position: absolute;
+  font-size: ${p => p.theme.fontSizeMedium};
+  z-index: ${p => p.theme.zIndex.dropdown};
+  background: ${p => p.theme.backgroundElevated};
+  border: 1px solid ${p => p.theme.border};
+  border-radius: ${p => p.theme.borderRadius};
+  box-shadow: ${p => p.theme.dropShadowHeavy};
+  width: auto;
+  min-width: 164px;
+  overflow: auto;
+  padding-bottom: ${space(0.5)};
+`;
+
+export {Menu as ProfilingContextMenu};
+
+const MenuContentContainer = styled('div')`
+  cursor: pointer;
+  display: flex;
+  align-items: center;
+  font-weight: normal;
+  padding: 0 ${space(1)};
+  border-radius: ${p => p.theme.borderRadius};
+  box-sizing: border-box;
+  background: ${p => (p.tabIndex === 0 ? p.theme.hover : undefined)};
+
+  &:focus {
+    color: ${p => p.theme.textColor};
+    background: ${p => p.theme.hover};
+    outline: none;
+  }
+`;
+
+const MenuItemCheckboxLabel = styled('label')`
+  display: flex;
+  align-items: center;
+  font-weight: normal;
+  margin: 0;
+  cursor: pointer;
+  flex: 1 1 100%;
+`;
+
+interface MenuItemCheckboxProps
+  extends React.DetailedHTMLProps<React.HTMLAttributes<HTMLDivElement>, HTMLDivElement> {
+  checked?: boolean;
+}
+
+const MenuItemCheckbox = forwardRef(
+  (props: MenuItemCheckboxProps, ref: React.Ref<HTMLDivElement> | undefined) => {
+    const {children, checked, ...rest} = props;
+
+    return (
+      <MenuContentOuterContainer>
+        <MenuContentContainer ref={ref} role="menuitem" {...rest}>
+          <MenuItemCheckboxLabel>
+            <MenuLeadingItem>
+              <Input type="checkbox" checked={checked} onChange={() => void 0} />
+              <IconCheckmark />
+            </MenuLeadingItem>
+            <MenuContent>{children}</MenuContent>
+          </MenuItemCheckboxLabel>
+        </MenuContentContainer>
+      </MenuContentOuterContainer>
+    );
+  }
+);
+
+export {MenuItemCheckbox as ProfilingContextMenuItemCheckbox};
+
+const MenuLeadingItem = styled('div')`
+  display: flex;
+  align-items: center;
+  height: 1.4em;
+  width: 1em;
+  gap: ${space(1)};
+  padding: ${space(1)} 0;
+  position: relative;
+`;
+
+const MenuContent = styled('div')`
+  position: relative;
+  width: 100%;
+  display: flex;
+  gap: ${space(0.5)};
+  justify-content: space-between;
+  padding: ${space(0.5)} 0;
+  margin-left: ${space(0.5)};
+  text-transform: capitalize;
+
+  margin-bottom: 0;
+  line-height: 1.4;
+  white-space: nowrap;
+  overflow: hidden;
+  text-overflow: ellipsis;
+`;
+
+const Input = styled('input')`
+  position: absolute;
+  opacity: 0;
+  cursor: pointer;
+  height: 0;
+  padding-right: ${space(1)};
+
+  & + svg {
+    position: absolute;
+    left: 50%;
+    top: 50%;
+    transform: translate(-50%, -50%);
+    width: 1em;
+    height: 1.4em;
+    display: none;
+  }
+
+  &:checked + svg {
+    display: block;
+  }
+`;
+
+interface MenuItemProps
+  extends React.DetailedHTMLProps<React.HTMLAttributes<HTMLDivElement>, HTMLDivElement> {
+  children: React.ReactNode;
+}
+
+const MenuItem = styled(
+  forwardRef((props: MenuItemProps, ref: React.Ref<HTMLDivElement> | undefined) => {
+    const {children, ...rest} = props;
+    return (
+      <MenuContentOuterContainer>
+        <MenuContentContainer ref={ref} role="menuitem" {...rest}>
+          <MenuContent>{children}</MenuContent>
+        </MenuContentContainer>
+      </MenuContentOuterContainer>
+    );
+  })
+)`
+  cursor: pointer;
+  color: ${p => p.theme.textColor};
+  background: transparent;
+  padding: 0 ${space(0.5)};
+
+  &:focus {
+    outline: none;
+  }
+
+  &:active: {
+    background: transparent;
+  }
+`;
+
+export {MenuItem as ProfilingContextMenuItem};
+
+const MenuContentOuterContainer = styled('div')`
+  padding: 0 ${space(0.5)};
+`;
+
+const MenuGroup = styled('div')`
+  padding-top: 0;
+  padding-bottom: ${space(1)};
+
+  &:last-of-type {
+    padding-bottom: 0;
+  }
+`;
+
+export {MenuGroup as ProfilingContextMenuGroup};
+
+interface MenuHeadingProps
+  extends React.DetailedHTMLProps<React.HTMLAttributes<HTMLDivElement>, HTMLDivElement> {
+  children: React.ReactNode;
+}
+
+const MenuHeading = styled((props: MenuHeadingProps) => {
+  const {children, ...rest} = props;
+  return <div {...rest}>{children}</div>;
+})`
+  text-transform: uppercase;
+  line-height: 1.5;
+  font-weight: 600;
+  color: ${p => p.theme.subText};
+  margin-bottom: 0;
+  cursor: default;
+  font-size: 75%;
+  padding: ${space(0.5)} ${space(1.5)};
+`;
+
+export {MenuHeading as ProfilingContextMenuHeading};
+
+const Layer = styled('div')`
+  width: 100%;
+  height: 100%;
+  position: absolute;
+  left: 0;
+  top: 0;
+  z-index: ${p => p.theme.zIndex.dropdown - 1};
+`;
+
+export {Layer as ProfilingContextMenuLayer};

static/app/components/profiling/flamegraphOptionsContextMenu.tsx

@@ -1,10 +1,6 @@
-import {forwardRef, Fragment, useEffect, useState} from 'react';
-import styled from '@emotion/styled';
+import {Fragment} from 'react';
 
-import {IconCheckmark} from 'sentry/icons';
 import {t} from 'sentry/locale';
-import space from 'sentry/styles/space';
-import {clamp} from 'sentry/utils/profiling/colors/utils';
 import {
   FlamegraphAxisOptions,
   FlamegraphColorCodings,
@@ -12,188 +8,15 @@ import {
   FlamegraphViewOptions,
 } from 'sentry/utils/profiling/flamegraph/flamegraphStateProvider/flamegraphPreferences';
 import {useFlamegraphPreferences} from 'sentry/utils/profiling/flamegraph/useFlamegraphPreferences';
-import {Rect} from 'sentry/utils/profiling/gl/utils';
 import {useContextMenu} from 'sentry/utils/profiling/hooks/useContextMenu';
 
-interface MenuProps
-  extends React.DetailedHTMLProps<React.HTMLAttributes<HTMLDivElement>, HTMLDivElement> {
-  children: React.ReactNode;
-}
-
-const Menu = styled(
-  forwardRef((props: MenuProps, ref: React.Ref<HTMLDivElement> | undefined) => {
-    return <div ref={ref} role="menu" {...props} />;
-  })
-)`
-  position: absolute;
-  font-size: ${p => p.theme.fontSizeMedium};
-  z-index: ${p => p.theme.zIndex.dropdown};
-  background: ${p => p.theme.backgroundElevated};
-  border: 1px solid ${p => p.theme.border};
-  border-radius: ${p => p.theme.borderRadius};
-  box-shadow: ${p => p.theme.dropShadowHeavy};
-  width: auto;
-  overflow: auto;
-
-  &:focus {
-    outline: none;
-  }
-`;
-
-interface MenuItemCheckboxProps
-  extends React.DetailedHTMLProps<React.HTMLAttributes<HTMLDivElement>, HTMLDivElement> {
-  checked?: boolean;
-}
-
-const MenuLeadingItem = styled('div')`
-  display: flex;
-  align-items: center;
-  height: 1.4em;
-  width: 1em;
-  gap: ${space(1)};
-  padding: ${space(1)} 0;
-  position: relative;
-`;
-
-const MenuContent = styled('div')`
-  position: relative;
-  width: 100%;
-  display: flex;
-  gap: ${space(0.5)};
-  justify-content: space-between;
-  padding: ${space(0.5)} 0;
-  margin-left: ${space(0.5)};
-  text-transform: capitalize;
-
-  margin-bottom: 0;
-  line-height: 1.4;
-  white-space: nowrap;
-  overflow: hidden;
-  text-overflow: ellipsis;
-`;
-
-const Input = styled('input')`
-  position: absolute;
-  opacity: 0;
-  cursor: pointer;
-  height: 0;
-  padding-right: ${space(1)};
-
-  & + svg {
-    position: absolute;
-    left: 50%;
-    top: 50%;
-    transform: translate(-50%, -50%);
-    width: 1em;
-    height: 1.4em;
-    display: none;
-  }
-
-  &:checked + svg {
-    display: block;
-  }
-`;
-
-const MenuItemCheckbox = styled(
-  forwardRef(
-    (props: MenuItemCheckboxProps, ref: React.Ref<HTMLDivElement> | undefined) => {
-      const {children, checked, className, style, ...rest} = props;
-
-      return (
-        // @ts-ignore this ref is forwarded
-        <MenuItem ref={ref} {...rest}>
-          <label className={className} style={style}>
-            <MenuLeadingItem>
-              <Input type="checkbox" checked={checked} onChange={() => void 0} />
-              <IconCheckmark />
-            </MenuLeadingItem>
-            <MenuContent>{children}</MenuContent>
-          </label>
-        </MenuItem>
-      );
-    }
-  )
-)`
-  cursor: pointer;
-  display: flex;
-  align-items: center;
-  font-weight: normal;
-  padding: 0 ${space(1)};
-  border-radius: ${p => p.theme.borderRadius};
-  box-sizing: border-box;
-  background: ${p => (p.tabIndex === 0 ? p.theme.hover : undefined)};
-
-  &:focus {
-    color: ${p => p.theme.textColor};
-    background: ${p => p.theme.hover};
-  }
-`;
-
-interface MenuItemProps
-  extends React.DetailedHTMLProps<React.HTMLAttributes<HTMLDivElement>, HTMLDivElement> {
-  children: React.ReactNode;
-}
-
-const MenuItem = styled(
-  forwardRef((props: MenuItemProps, ref: React.Ref<HTMLDivElement> | undefined) => {
-    const {children, ...rest} = props;
-    return (
-      <div ref={ref} role="menuitem" {...rest}>
-        {children}
-      </div>
-    );
-  })
-)`
-  cursor: pointer;
-  color: ${p => p.theme.textColor};
-  background: transparent;
-  padding: 0 ${space(0.5)};
-
-  &:focus {
-    outline: none;
-  }
-
-  &:active: {
-    background: transparent;
-  }
-`;
-
-const MenuGroup = styled('div')`
-  padding-top: 0;
-  padding-bottom: ${space(1)};
-
-  &:last-of-type {
-    padding-bottom: 0;
-  }
-`;
-
-interface MenuHeadingProps
-  extends React.DetailedHTMLProps<React.HTMLAttributes<HTMLDivElement>, HTMLDivElement> {
-  children: React.ReactNode;
-}
-
-const MenuHeading = styled((props: MenuHeadingProps) => {
-  const {children, ...rest} = props;
-  return <div {...rest}>{children}</div>;
-})`
-  text-transform: uppercase;
-  line-height: 1.5;
-  font-weight: 600;
-  color: ${p => p.theme.subText};
-  margin-bottom: 0;
-  cursor: default;
-  font-size: 75%;
-  padding: ${space(0.5)} ${space(1.5)};
-`;
-
-const Layer = styled('div')`
-  width: 100%;
-  height: 100%;
-  position: absolute;
-  left: 0;
-  top: 0;
-  z-index: ${p => p.theme.zIndex.dropdown - 1};
-`;
+import {
+  ProfilingContextMenu,
+  ProfilingContextMenuGroup,
+  ProfilingContextMenuHeading,
+  ProfilingContextMenuItemCheckbox,
+  ProfilingContextMenuLayer,
+} from './ProfilingContextMenu/profilingContextMenu';
 
 const FLAMEGRAPH_COLOR_CODINGS: FlamegraphColorCodings = [
   'by symbol name',
@@ -205,173 +28,78 @@ const FLAMEGRAPH_VIEW_OPTIONS: FlamegraphViewOptions = ['top down', 'bottom up']
 const FLAMEGRAPH_SORTING_OPTIONS: FlamegraphSorting = ['left heavy', 'call order'];
 const FLAMEGRAPH_AXIS_OPTIONS: FlamegraphAxisOptions = ['standalone', 'transaction'];
 
-function computeBestContextMenuPosition(mouse: Rect, container: Rect, target: Rect) {
-  const maxY = Math.floor(container.height - target.height);
-  const minY = container.top;
-
-  const minX = container.left;
-  const maxX = Math.floor(container.right - target.width);
-
-  // We add a tiny offset so that the menu is not directly where the user places their cursor.
-  const OFFSET = 6;
-
-  return {
-    left: clamp(mouse.x + OFFSET, minX, maxX),
-    top: clamp(mouse.y + OFFSET, minY, maxY),
-  };
-}
-
 interface FlameGraphOptionsContextMenuProps {
-  container: HTMLElement | null;
-  contextMenuCoordinates: Rect | null;
-  contextMenuProps: ReturnType<typeof useContextMenu>;
+  contextMenu: ReturnType<typeof useContextMenu>;
 }
 
 export function FlamegraphOptionsContextMenu(props: FlameGraphOptionsContextMenuProps) {
-  const [containerCoordinates, setContainerCoordinates] = useState<Rect | null>(null);
-  const [menuCoordinates, setMenuCoordinates] = useState<Rect | null>(null);
-  const {open, setOpen, menuRef, getMenuProps, getMenuItemProps} = props.contextMenuProps;
-
   const [preferences, dispatch] = useFlamegraphPreferences();
 
-  useEffect(() => {
-    const listener = (event: MouseEvent | TouchEvent) => {
-      // Do nothing if clicking ref's element or descendent elements
-      if (!menuRef || menuRef.contains(event.target as Node)) {
-        return;
-      }
-
-      setOpen(false);
-    };
-
-    document.addEventListener('mousedown', listener);
-    document.addEventListener('touchstart', listener);
-    return () => {
-      document.removeEventListener('mousedown', listener);
-      document.removeEventListener('touchstart', listener);
-    };
-  }, [menuRef, setOpen]);
-
-  // Observe the menu
-  useEffect(() => {
-    if (!menuRef) {
-      return undefined;
-    }
-
-    const resizeObserver = new window.ResizeObserver(entries => {
-      const contentRect = entries[0].contentRect;
-      setMenuCoordinates(new Rect(0, 0, contentRect.width, contentRect.height));
-    });
-
-    resizeObserver.observe(menuRef);
-
-    return () => {
-      resizeObserver.disconnect();
-    };
-  }, [menuRef]);
-
-  // Observe the container
-  useEffect(() => {
-    if (!props.container) {
-      return undefined;
-    }
-
-    const resizeObserver = new window.ResizeObserver(entries => {
-      const contentRect = entries[0].contentRect;
-      setContainerCoordinates(new Rect(0, 0, contentRect.width, contentRect.height));
-    });
-
-    resizeObserver.observe(props.container);
-
-    return () => {
-      resizeObserver.disconnect();
-    };
-  }, [props.container]);
-
-  const position =
-    props.contextMenuCoordinates && containerCoordinates && menuCoordinates
-      ? computeBestContextMenuPosition(
-          props.contextMenuCoordinates,
-          containerCoordinates,
-          menuCoordinates
-        )
-      : null;
-
-  return (
+  return props.contextMenu.open ? (
     <Fragment>
-      {open ? (
-        <Layer
-          onClick={() => {
-            setOpen(false);
-          }}
-        />
-      ) : null}
-      {props.contextMenuCoordinates ? (
-        <Menu
-          {...getMenuProps()}
-          style={{
-            position: 'absolute',
-            visibility: open ? 'initial' : 'hidden',
-            left: position?.left ?? -9999,
-            top: position?.top ?? -9999,
-            pointerEvents: open ? 'initial' : 'none',
-            maxHeight: containerCoordinates?.height ?? 'auto',
-          }}
-        >
-          <MenuGroup>
-            <MenuHeading>{t('Color Coding')}</MenuHeading>
-            {FLAMEGRAPH_COLOR_CODINGS.map((coding, idx) => (
-              <MenuItemCheckbox
-                key={idx}
-                {...getMenuItemProps()}
-                onClick={() => dispatch({type: 'set color coding', payload: coding})}
-                checked={preferences.colorCoding === coding}
-              >
-                {coding}
-              </MenuItemCheckbox>
-            ))}
-          </MenuGroup>
-          <MenuGroup>
-            <MenuHeading>{t('View')}</MenuHeading>
-            {FLAMEGRAPH_VIEW_OPTIONS.map((view, idx) => (
-              <MenuItemCheckbox
-                key={idx}
-                {...getMenuItemProps()}
-                onClick={() => dispatch({type: 'set view', payload: view})}
-                checked={preferences.view === view}
-              >
-                {view}
-              </MenuItemCheckbox>
-            ))}
-          </MenuGroup>
-          <MenuGroup>
-            <MenuHeading>{t('Sorting')}</MenuHeading>
-            {FLAMEGRAPH_SORTING_OPTIONS.map((sorting, idx) => (
-              <MenuItemCheckbox
-                key={idx}
-                {...getMenuItemProps()}
-                onClick={() => dispatch({type: 'set sorting', payload: sorting})}
-                checked={preferences.sorting === sorting}
-              >
-                {sorting}
-              </MenuItemCheckbox>
-            ))}
-          </MenuGroup>
-          <MenuGroup>
-            <MenuHeading>{t('X Axis')}</MenuHeading>
-            {FLAMEGRAPH_AXIS_OPTIONS.map((axis, idx) => (
-              <MenuItemCheckbox
-                key={idx}
-                {...getMenuItemProps()}
-                onClick={() => dispatch({type: 'set xAxis', payload: axis})}
-                checked={preferences.xAxis === axis}
-              >
-                {axis}
-              </MenuItemCheckbox>
-            ))}
-          </MenuGroup>
-        </Menu>
-      ) : null}
+      <ProfilingContextMenuLayer onClick={() => props.contextMenu.setOpen(false)} />
+      <ProfilingContextMenu
+        {...props.contextMenu.getMenuProps()}
+        style={{
+          position: 'absolute',
+          left: props.contextMenu.position?.left ?? -9999,
+          top: props.contextMenu.position?.top ?? -9999,
+          maxHeight: props.contextMenu.containerCoordinates?.height ?? 'auto',
+        }}
+      >
+        <ProfilingContextMenuGroup>
+          <ProfilingContextMenuHeading>{t('Color Coding')}</ProfilingContextMenuHeading>
+          {FLAMEGRAPH_COLOR_CODINGS.map((coding, idx) => (
+            <ProfilingContextMenuItemCheckbox
+              key={idx}
+              {...props.contextMenu.getMenuItemProps()}
+              onClick={() => dispatch({type: 'set color coding', payload: coding})}
+              checked={preferences.colorCoding === coding}
+            >
+              {coding}
+            </ProfilingContextMenuItemCheckbox>
+          ))}
+        </ProfilingContextMenuGroup>
+        <ProfilingContextMenuGroup>
+          <ProfilingContextMenuHeading>{t('View')}</ProfilingContextMenuHeading>
+          {FLAMEGRAPH_VIEW_OPTIONS.map((view, idx) => (
+            <ProfilingContextMenuItemCheckbox
+              key={idx}
+              {...props.contextMenu.getMenuItemProps()}
+              onClick={() => dispatch({type: 'set view', payload: view})}
+              checked={preferences.view === view}
+            >
+              {view}
+            </ProfilingContextMenuItemCheckbox>
+          ))}
+        </ProfilingContextMenuGroup>
+        <ProfilingContextMenuGroup>
+          <ProfilingContextMenuHeading>{t('Sorting')}</ProfilingContextMenuHeading>
+          {FLAMEGRAPH_SORTING_OPTIONS.map((sorting, idx) => (
+            <ProfilingContextMenuItemCheckbox
+              key={idx}
+              {...props.contextMenu.getMenuItemProps()}
+              onClick={() => dispatch({type: 'set sorting', payload: sorting})}
+              checked={preferences.sorting === sorting}
+            >
+              {sorting}
+            </ProfilingContextMenuItemCheckbox>
+          ))}
+        </ProfilingContextMenuGroup>
+        <ProfilingContextMenuGroup>
+          <ProfilingContextMenuHeading>{t('X Axis')}</ProfilingContextMenuHeading>
+          {FLAMEGRAPH_AXIS_OPTIONS.map((axis, idx) => (
+            <ProfilingContextMenuItemCheckbox
+              key={idx}
+              {...props.contextMenu.getMenuItemProps()}
+              onClick={() => dispatch({type: 'set xAxis', payload: axis})}
+              checked={preferences.xAxis === axis}
+            >
+              {axis}
+            </ProfilingContextMenuItemCheckbox>
+          ))}
+        </ProfilingContextMenuGroup>
+      </ProfilingContextMenu>
     </Fragment>
-  );
+  ) : null;
 }

static/app/components/profiling/flamegraphZoomView.tsx

@@ -14,17 +14,15 @@ import {
 } from 'sentry/utils/profiling/flamegraph/useFlamegraphState';
 import {useFlamegraphTheme} from 'sentry/utils/profiling/flamegraph/useFlamegraphTheme';
 import {FlamegraphCanvas} from 'sentry/utils/profiling/flamegraphCanvas';
-import {FlamegraphFrame} from 'sentry/utils/profiling/flamegraphFrame';
 import {FlamegraphView} from 'sentry/utils/profiling/flamegraphView';
 import {formatColorForFrame, Rect} from 'sentry/utils/profiling/gl/utils';
+import {useContextMenu} from 'sentry/utils/profiling/hooks/useContextMenu';
 import {FlamegraphRenderer} from 'sentry/utils/profiling/renderers/flamegraphRenderer';
 import {GridRenderer} from 'sentry/utils/profiling/renderers/gridRenderer';
 import {SelectedFrameRenderer} from 'sentry/utils/profiling/renderers/selectedFrameRenderer';
 import {TextRenderer} from 'sentry/utils/profiling/renderers/textRenderer';
 import usePrevious from 'sentry/utils/usePrevious';
 
-import {useContextMenu} from '../../utils/profiling/hooks/useContextMenu';
-
 import {BoundTooltip} from './boundTooltip';
 import {FlamegraphOptionsContextMenu} from './flamegraphOptionsContextMenu';
 
@@ -327,9 +325,8 @@ function FlamegraphZoomView({
       setConfigSpaceCursor(null);
     };
 
-    const onZoomIntoFrame = (frame: FlamegraphFrame) => {
+    const onZoomIntoFrame = () => {
       setConfigSpaceCursor(null);
-      dispatchFlamegraphState({type: 'set selected node', payload: frame});
     };
 
     scheduler.on('resetZoom', onResetZoom);
@@ -589,32 +586,7 @@ function FlamegraphZoomView({
     };
   }, [flamegraphCanvasRef, zoom, scroll]);
 
-  // Context menu coordinates
-  const contextMenuProps = useContextMenu();
-  const [contextMenuCoordinates, setContextMenuCoordinates] = useState<Rect | null>(null);
-  const onContextMenu = useCallback(
-    (evt: React.MouseEvent) => {
-      evt.preventDefault();
-      evt.stopPropagation();
-
-      if (!flamegraphCanvasRef) {
-        return;
-      }
-
-      const parentPosition = flamegraphCanvasRef.getBoundingClientRect();
-
-      setContextMenuCoordinates(
-        new Rect(
-          evt.clientX - parentPosition.left,
-          evt.clientY - parentPosition.top,
-          0,
-          0
-        )
-      );
-      contextMenuProps.setOpen(true);
-    },
-    [flamegraphCanvasRef, contextMenuProps]
-  );
+  const contextMenu = useContextMenu({container: flamegraphCanvasRef});
 
   return (
     <Fragment>
@@ -625,7 +597,7 @@ function FlamegraphZoomView({
           onMouseUp={onCanvasMouseUp}
           onMouseMove={onCanvasMouseMove}
           onMouseLeave={onCanvasMouseLeave}
-          onContextMenu={onContextMenu}
+          onContextMenu={contextMenu.handleContextMenu}
           style={{cursor: lastInteraction === 'pan' ? 'grab' : 'default'}}
         />
         <Canvas
@@ -634,13 +606,7 @@ function FlamegraphZoomView({
             pointerEvents: 'none',
           }}
         />
-        {contextMenuProps.open ? (
-          <FlamegraphOptionsContextMenu
-            container={flamegraphCanvasRef}
-            contextMenuCoordinates={contextMenuCoordinates}
-            contextMenuProps={contextMenuProps}
-          />
-        ) : null}
+        <FlamegraphOptionsContextMenu contextMenu={contextMenu} />
         {flamegraphCanvas &&
         flamegraphRenderer &&
         flamegraphView &&

static/app/utils/profiling/hooks/useContextMenu.tsx

@@ -1,50 +1,176 @@
-import {useState} from 'react';
+import {useCallback, useEffect, useState} from 'react';
+
+import {clamp} from 'sentry/utils/profiling/colors/utils';
+import {Rect} from 'sentry/utils/profiling/gl/utils';
 
 import {useKeyboardNavigation} from './useKeyboardNavigation';
 
-export function useContextMenu() {
+export function computeBestContextMenuPosition(
+  mouse: Rect,
+  container: Rect,
+  target: Rect
+) {
+  const maxY = Math.floor(container.height - target.height);
+  const minY = container.top;
+
+  const minX = container.left;
+  const maxX = Math.floor(container.right - target.width);
+
+  // We add a tiny offset so that the menu is not directly where the user places their cursor.
+  const OFFSET = 6;
+
+  return {
+    left: clamp(mouse.x + OFFSET, minX, maxX),
+    top: clamp(mouse.y + OFFSET, minY, maxY),
+  };
+}
+
+interface UseContextMenuOptions {
+  container: HTMLElement | null;
+}
+
+export function useContextMenu({container}: UseContextMenuOptions) {
   const [open, setOpen] = useState<boolean>(false);
+  const [menuCoordinates, setMenuCoordinates] = useState<Rect | null>(null);
+  const [contextMenuCoordinates, setContextMenuCoordinates] = useState<Rect | null>(null);
+  const [containerCoordinates, setContainerCoordinates] = useState<Rect | null>(null);
+
   const itemProps = useKeyboardNavigation();
 
-  function wrapSetOpen(newOpen: boolean) {
-    if (!newOpen) {
-      itemProps.setTabIndex(null);
-    }
-    setOpen(newOpen);
-  }
+  // We wrap the setOpen function in a useEffect so that we also clear the keyboard
+  // tabIndex when a menu is closed. This prevents tabIndex from being persisted between render
+  const wrapSetOpen = useCallback(
+    (newOpen: boolean) => {
+      if (!newOpen) {
+        itemProps.setTabIndex(null);
+      }
+      setOpen(newOpen);
+    },
+    [itemProps]
+  );
 
-  function getMenuProps() {
-    const menuProps = itemProps.getMenuKeyboardEventHandlers();
+  const getMenuProps = useCallback(() => {
+    const menuProps = itemProps.getMenuProps();
 
     return {
       ...menuProps,
       onKeyDown: (evt: React.KeyboardEvent) => {
         if (evt.key === 'Escape') {
-          setOpen(false);
+          wrapSetOpen(false);
         }
         menuProps.onKeyDown(evt);
       },
     };
-  }
+  }, [itemProps, wrapSetOpen]);
 
-  function getMenuItemProps() {
-    const menuItemProps = itemProps.getMenuItemKeyboardEventHandlers();
+  const getMenuItemProps = useCallback(() => {
+    const menuItemProps = itemProps.getItemProps();
 
     return {
       ...menuItemProps,
       onKeyDown: (evt: React.KeyboardEvent) => {
         if (evt.key === 'Escape') {
-          setOpen(false);
+          wrapSetOpen(false);
         }
         menuItemProps.onKeyDown(evt);
       },
     };
-  }
+  }, [itemProps, wrapSetOpen]);
+
+  const handleContextMenu = useCallback(
+    (evt: React.MouseEvent) => {
+      if (!container) {
+        return;
+      }
+      evt.preventDefault();
+      evt.stopPropagation();
+
+      const parentPosition = container.getBoundingClientRect();
+
+      setContextMenuCoordinates(
+        new Rect(
+          evt.clientX - parentPosition.left,
+          evt.clientY - parentPosition.top,
+          0,
+          0
+        )
+      );
+      wrapSetOpen(true);
+    },
+    [wrapSetOpen, container]
+  );
+
+  useEffect(() => {
+    const listener = (event: MouseEvent | TouchEvent) => {
+      // Do nothing if clicking ref's element or descendent elements
+      if (!itemProps.menuRef || itemProps.menuRef.contains(event.target as Node)) {
+        return;
+      }
+
+      setOpen(false);
+    };
+
+    document.addEventListener('mousedown', listener);
+    document.addEventListener('touchstart', listener);
+    return () => {
+      document.removeEventListener('mousedown', listener);
+      document.removeEventListener('touchstart', listener);
+    };
+  }, [itemProps.menuRef]);
+
+  // Observe the menu
+  useEffect(() => {
+    if (!itemProps.menuRef) {
+      return undefined;
+    }
+
+    const resizeObserver = new window.ResizeObserver(entries => {
+      const contentRect = entries[0].contentRect;
+      setMenuCoordinates(new Rect(0, 0, contentRect.width, contentRect.height));
+    });
+
+    resizeObserver.observe(itemProps.menuRef);
+
+    return () => {
+      resizeObserver.disconnect();
+    };
+  }, [itemProps.menuRef]);
+
+  // Observe the container
+  useEffect(() => {
+    if (!container) {
+      return undefined;
+    }
+
+    const resizeObserver = new window.ResizeObserver(entries => {
+      const contentRect = entries[0].contentRect;
+      setContainerCoordinates(new Rect(0, 0, contentRect.width, contentRect.height));
+    });
+
+    resizeObserver.observe(container);
+
+    return () => {
+      resizeObserver.disconnect();
+    };
+  }, [container]);
+
+  const position =
+    contextMenuCoordinates && containerCoordinates && menuCoordinates
+      ? computeBestContextMenuPosition(
+          contextMenuCoordinates,
+          containerCoordinates,
+          menuCoordinates
+        )
+      : null;
 
   return {
     open,
     setOpen: wrapSetOpen,
+    position,
+    containerCoordinates,
+    contextMenuCoordinates: position,
     menuRef: itemProps.menuRef,
+    handleContextMenu,
     getMenuProps,
     getMenuItemProps,
   };

static/app/utils/profiling/hooks/useKeyboardNavigation.tsx

@@ -1,11 +1,54 @@
-import {useEffect, useState} from 'react';
+import {useCallback, useEffect, useState} from 'react';
 
-export function useKeyboardNavigation() {
-  const [menuRef, setMenuRef] = useState<HTMLDivElement | null>(null);
+export function useRovingTabIndex(items: any[]) {
   const [tabIndex, setTabIndex] = useState<number | null>(null);
 
+  const onKeyDown = useCallback(
+    (evt: React.KeyboardEvent) => {
+      if (items.length === 0) {
+        return;
+      }
+
+      if (evt.key === 'Escape') {
+        setTabIndex(null);
+      }
+
+      if (evt.key === 'ArrowDown' || evt.key === 'Tab') {
+        evt.preventDefault();
+
+        if (tabIndex === items.length - 1 || tabIndex === null) {
+          setTabIndex(0);
+        } else {
+          setTabIndex((tabIndex ?? 0) + 1);
+        }
+      }
+
+      if (evt.key === 'ArrowUp' || (evt.key === 'Tab' && evt.shiftKey)) {
+        evt.preventDefault();
+
+        if (tabIndex === 0 || tabIndex === null) {
+          setTabIndex(items.length - 1);
+        } else {
+          setTabIndex((tabIndex ?? 0) - 1);
+        }
+      }
+    },
+    [tabIndex, items]
+  );
+
+  return {
+    tabIndex,
+    setTabIndex,
+    onKeyDown,
+  };
+}
+
+export function useKeyboardNavigation() {
+  const [menuRef, setMenuRef] = useState<HTMLDivElement | null>(null);
   const items: {id: number; node: HTMLElement | null}[] = [];
 
+  const {tabIndex, setTabIndex, onKeyDown} = useRovingTabIndex(items);
+
   useEffect(() => {
     if (menuRef) {
       if (tabIndex === null) {
@@ -14,54 +57,15 @@ export function useKeyboardNavigation() {
     }
   }, [menuRef, tabIndex]);
 
-  useEffect(() => {
-    if (typeof tabIndex !== 'number') {
-      return;
-    }
-    if (items[tabIndex]?.node) {
-      items[tabIndex]?.node?.focus();
-    }
-    // We only want to focus the element if the tabIndex changes
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [tabIndex]);
-
-  function getMenuKeyboardEventHandlers() {
+  function getMenuProps() {
     return {
       tabIndex: -1,
       ref: setMenuRef,
-      onKeyDown: (evt: React.KeyboardEvent) => {
-        if (items.length === 0) {
-          return;
-        }
-
-        if (evt.key === 'Escape') {
-          setTabIndex(null);
-        }
-
-        if (evt.key === 'ArrowDown' || evt.key === 'Tab') {
-          evt.preventDefault();
-
-          if (tabIndex === items.length - 1 || tabIndex === null) {
-            setTabIndex(0);
-          } else {
-            setTabIndex((tabIndex ?? 0) + 1);
-          }
-        }
-
-        if (evt.key === 'ArrowUp' || (evt.key === 'Tab' && evt.shiftKey)) {
-          evt.preventDefault();
-
-          if (tabIndex === 0 || tabIndex === null) {
-            setTabIndex(items.length - 1);
-          } else {
-            setTabIndex((tabIndex ?? 0) - 1);
-          }
-        }
-      },
+      onKeyDown,
     };
   }
 
-  function getMenuItemKeyboardEventHandlers() {
+  function getItemProps() {
     const idx = items.length;
     items.push({id: idx, node: null});
 
@@ -69,52 +73,23 @@ export function useKeyboardNavigation() {
       tabIndex: tabIndex === idx ? 0 : -1,
       ref: (node: HTMLElement | null) => {
         if (items[idx]) {
+          if (tabIndex === idx) {
+            node?.focus();
+          }
           items[idx].node = node;
         }
       },
       onMouseEnter: () => {
         setTabIndex(idx);
       },
-      onKeyDown: (evt: React.KeyboardEvent) => {
-        if (items.length === 0) {
-          return;
-        }
-
-        if (evt.key === 'Escape') {
-          setTabIndex(null);
-        }
-
-        if (evt.key === 'Enter' || evt.key === ' ') {
-          items?.[idx]?.node?.click?.();
-        }
-
-        if (evt.key === 'ArrowDown' || evt.key === 'Tab') {
-          evt.preventDefault();
-
-          if (tabIndex === items.length || tabIndex === null) {
-            setTabIndex(0);
-          } else {
-            setTabIndex((tabIndex ?? 0) + 1);
-          }
-        }
-
-        if (evt.key === 'ArrowUp' || (evt.key === 'Tab' && evt.shiftKey)) {
-          evt.preventDefault();
-
-          if (tabIndex === 0 || tabIndex === null) {
-            setTabIndex(items.length);
-          } else {
-            setTabIndex((tabIndex ?? 0) - 1);
-          }
-        }
-      },
+      onKeyDown,
     };
   }
 
   return {
     menuRef,
-    getMenuItemKeyboardEventHandlers,
-    getMenuKeyboardEventHandlers,
+    getItemProps,
+    getMenuProps,
     tabIndex,
     setTabIndex,
   };

static/app/utils/profiling/hooks/useVirtualizedTree/useVirtualizedTree.tsx

@@ -93,6 +93,7 @@ export function useVirtualizedTree<T extends TreeLike>(
     const visibleItems: {
       item: VirtualizedTreeNode<T>;
       key: number;
+      ref: HTMLElement | null;
       styles: React.CSSProperties;
     }[] = [];
 
@@ -126,6 +127,7 @@ export function useVirtualizedTree<T extends TreeLike>(
       if (elementTop >= viewport.top && elementBottom <= viewport.bottom) {
         visibleItems[visibleItemIndex] = {
           key: indexPointer,
+          ref: null,
           styles: {position: 'absolute', top: elementTop},
           item: tree.flattened[indexPointer],
         };
@@ -225,10 +227,63 @@ export function useVirtualizedTree<T extends TreeLike>(
     return {height: tree.flattened.length * props.rowHeight};
   }, [tree.flattened.length, props.rowHeight]);
 
+  const [tabIndexKey, setTabIndexKey] = useState<number | null>(null);
+  const handleRowClick = useCallback((key: number) => {
+    setTabIndexKey(key);
+  }, []);
+
+  const handleRowKeyDown = useCallback(
+    (key: number, event: React.KeyboardEvent) => {
+      if (event.key === 'Enter') {
+        handleExpandTreeNode(tree.flattened[key], {expandChildren: true});
+      }
+
+      if (event.key === 'ArrowDown') {
+        event.preventDefault();
+        const indexInVisibleItems = items.findIndex(i => i.key === key);
+
+        if (indexInVisibleItems !== -1) {
+          const nextIndex = indexInVisibleItems + 1;
+
+          // Bound check if we are at end of list
+          if (nextIndex > tree.flattened.length - 1) {
+            return;
+          }
+
+          setTabIndexKey(items[nextIndex].key);
+          items[nextIndex].ref?.focus({preventScroll: true});
+          items[nextIndex].ref?.scrollIntoView({block: 'nearest'});
+        }
+      }
+
+      if (event.key === 'ArrowUp') {
+        event.preventDefault();
+        const indexInVisibleItems = items.findIndex(i => i.key === key);
+
+        if (indexInVisibleItems !== -1) {
+          const nextIndex = indexInVisibleItems - 1;
+
+          // Bound check if we are at start of list
+          if (nextIndex < 0) {
+            return;
+          }
+
+          setTabIndexKey(items[nextIndex].key);
+          items[nextIndex].ref?.focus({preventScroll: true});
+          items[nextIndex].ref?.scrollIntoView({block: 'nearest'});
+        }
+      }
+    },
+    [handleExpandTreeNode, items, tree.flattened]
+  );
+
   return {
     tree,
     items,
+    tabIndexKey,
     handleScroll,
+    handleRowClick,
+    handleRowKeyDown,
     handleExpandTreeNode,
     handleSortingChange,
     scrollContainerStyles,