feat(profiling): lift activeProfileIndex so we stop reading from multiple sources (#34060)

* feat(profiling): add encode/decode state utility fns

* feat(profiling): make guards exhaustive

* feat(profiling): use sync functionality

* feat(profiling): lift index to view

* feat(profiling): move selectedNode to flamegraph state

* fix(profiling): wrong comment

static/app/components/profiling/flamegraph.tsx

@@ -1,4 +1,4 @@
-import {Fragment, ReactElement, useCallback, useMemo, useState} from 'react';
+import {Fragment, ReactElement, useMemo} from 'react';
 import styled from '@emotion/styled';
 
 import {FlamegraphOptionsMenu} from 'sentry/components/profiling/flamegraphOptionsMenu';
@@ -7,12 +7,16 @@ import {FlamegraphToolbar} from 'sentry/components/profiling/flamegraphToolbar';
 import {FlamegraphViewSelectMenu} from 'sentry/components/profiling/flamegraphViewSelectMenu';
 import {FlamegraphZoomView} from 'sentry/components/profiling/flamegraphZoomView';
 import {FlamegraphZoomViewMinimap} from 'sentry/components/profiling/flamegraphZoomViewMinimap';
-import {ProfileDragDropImport} from 'sentry/components/profiling/profileDragDropImport';
+import {
+  ProfileDragDropImport,
+  ProfileDragDropImportProps,
+} from 'sentry/components/profiling/profileDragDropImport';
 import {ThreadMenuSelector} from 'sentry/components/profiling/threadSelector';
 import {CanvasPoolManager} from 'sentry/utils/profiling/canvasScheduler';
 import {Flamegraph as FlamegraphModel} from 'sentry/utils/profiling/flamegraph';
 import {FlamegraphTheme} from 'sentry/utils/profiling/flamegraph/flamegraphTheme';
 import {useFlamegraphPreferences} from 'sentry/utils/profiling/flamegraph/useFlamegraphPreferences';
+import {useFlamegraphProfiles} from 'sentry/utils/profiling/flamegraph/useFlamegraphProfiles';
 import {useFlamegraphTheme} from 'sentry/utils/profiling/flamegraph/useFlamegraphTheme';
 import {Rect} from 'sentry/utils/profiling/gl/utils';
 import {ProfileGroup} from 'sentry/utils/profiling/profile/importProfile';
@@ -24,44 +28,39 @@ function getTransactionConfigSpace(profiles: Profile[]): Rect {
   return new Rect(startedAt, 0, endedAt - startedAt, 0);
 }
 interface FlamegraphProps {
+  onImport: ProfileDragDropImportProps['onImport'];
   profiles: ProfileGroup;
 }
 
 function Flamegraph(props: FlamegraphProps): ReactElement {
   const flamegraphTheme = useFlamegraphTheme();
   const [{sorting, view, xAxis}, dispatch] = useFlamegraphPreferences();
-  const canvasPoolManager = useMemo(() => new CanvasPoolManager(), []);
-
-  const [activeProfileIndex, setActiveProfileIndex] = useState<number | null>(null);
-  const [importedProfiles, setImportedProfiles] = useState<ProfileGroup | null>(null);
+  const [{activeProfileIndex}, dispatchActiveProfileIndex] = useFlamegraphProfiles();
 
-  // once an import occurs, it will always take precedence over the profile in the props
-  const profiles = importedProfiles ?? props.profiles;
+  const canvasPoolManager = useMemo(() => new CanvasPoolManager(), []);
 
   const flamegraph = useMemo(() => {
+    if (
+      !props.profiles.profiles[activeProfileIndex ?? props.profiles.activeProfileIndex]
+    ) {
+      // This could happen if activeProfileIndex was initialized from query string, but for some
+      // reason the profile was removed from the list of profiles.
+      return FlamegraphModel.Empty();
+    }
     // if the activeProfileIndex is null, use the activeProfileIndex from the profile group
-    const profileIndex = activeProfileIndex ?? profiles.activeProfileIndex;
-
-    const flamegraphModel = new FlamegraphModel(
-      profiles.profiles[profileIndex],
-      profileIndex,
+    return new FlamegraphModel(
+      props.profiles.profiles[activeProfileIndex ?? props.profiles.activeProfileIndex],
+      activeProfileIndex ?? props.profiles.activeProfileIndex,
       {
         inverted: view === 'bottom up',
         leftHeavy: sorting === 'left heavy',
         configSpace:
           xAxis === 'transaction'
-            ? getTransactionConfigSpace(profiles.profiles)
+            ? getTransactionConfigSpace(props.profiles.profiles)
             : undefined,
       }
     );
-
-    return flamegraphModel;
-  }, [profiles, activeProfileIndex, sorting, xAxis, view]);
-
-  const onImport = useCallback((profile: ProfileGroup) => {
-    setActiveProfileIndex(null);
-    setImportedProfiles(profile);
-  }, []);
+  }, [props.profiles, activeProfileIndex, sorting, xAxis, view]);
 
   return (
     <Fragment>
@@ -69,7 +68,9 @@ function Flamegraph(props: FlamegraphProps): ReactElement {
         <ThreadMenuSelector
           profileGroup={props.profiles}
           activeProfileIndex={flamegraph.profileIndex}
-          onProfileIndexChange={setActiveProfileIndex}
+          onProfileIndexChange={index =>
+            dispatchActiveProfileIndex({type: 'set active profile index', payload: index})
+          }
         />
         <FlamegraphViewSelectMenu
           view={view}
@@ -95,9 +96,8 @@ function Flamegraph(props: FlamegraphProps): ReactElement {
         />
       </FlamegraphZoomViewMinimapContainer>
       <FlamegraphZoomViewContainer>
-        <ProfileDragDropImport onImport={onImport}>
+        <ProfileDragDropImport onImport={props.onImport}>
           <FlamegraphZoomView
-            key={`${profiles.traceID}-${flamegraph.profileIndex}`}
             flamegraph={flamegraph}
             canvasPoolManager={canvasPoolManager}
           />

static/app/components/profiling/flamegraphZoomView.tsx

@@ -51,7 +51,6 @@ function FlamegraphZoomView({
   const [flamegraphState, dispatchFlamegraphState] = useFlamegraphState();
   const [canvasBounds, setCanvasBounds] = useState<Rect>(Rect.Empty());
   const [startPanVector, setStartPanVector] = useState<vec2 | null>(null);
-  const [selectedNode, setSelectedNode] = useState<FlamegraphFrame | null>(null);
   const [configSpaceCursor, setConfigSpaceCursor] = useState<vec2 | null>(null);
 
   const flamegraphRenderer = useMemoWithPrevious<FlamegraphRenderer | null>(
@@ -256,14 +255,17 @@ function FlamegraphZoomView({
     };
 
     const drawSelectedFrameBorder = () => {
-      if (selectedNode) {
+      if (flamegraphState.profiles.selectedNode) {
         selectedFrameRenderer.draw(
           new Rect(
-            selectedNode.start,
+            flamegraphState.profiles.selectedNode.start,
             flamegraph.inverted
-              ? flamegraphRenderer.configSpace.height - selectedNode.depth - 1
-              : selectedNode.depth,
-            selectedNode.end - selectedNode.start,
+              ? flamegraphRenderer.configSpace.height -
+                flamegraphState.profiles.selectedNode.depth -
+                1
+              : flamegraphState.profiles.selectedNode.depth,
+            flamegraphState.profiles.selectedNode.end -
+              flamegraphState.profiles.selectedNode.start,
             1
           ),
           {
@@ -275,7 +277,7 @@ function FlamegraphZoomView({
         );
       }
 
-      if (hoveredNode && selectedNode !== hoveredNode) {
+      if (hoveredNode && flamegraphState.profiles.selectedNode !== hoveredNode) {
         selectedFrameRenderer.draw(
           new Rect(
             hoveredNode.start,
@@ -330,7 +332,7 @@ function FlamegraphZoomView({
     flamegraphRenderer,
     textRenderer,
     gridRenderer,
-    selectedNode,
+    flamegraphState.profiles.selectedNode,
     hoveredNode,
     selectedFrameRenderer,
   ]);
@@ -372,7 +374,7 @@ function FlamegraphZoomView({
       );
 
       setConfigSpaceCursor(null);
-      setSelectedNode(frame);
+      dispatchFlamegraphState({type: 'set selected node', payload: frame});
 
       scheduler.draw();
     };
@@ -388,7 +390,7 @@ function FlamegraphZoomView({
       scheduler.off('resetZoom', onResetZoom);
       scheduler.off('zoomIntoFrame', onZoomIntoFrame);
     };
-  }, [scheduler, flamegraphRenderer]);
+  }, [scheduler, flamegraphRenderer, dispatchFlamegraphState]);
 
   useEffect(() => {
     if (!flamegraphCanvasRef || !flamegraphOverlayCanvasRef || !flamegraphRenderer) {
@@ -457,11 +459,15 @@ function FlamegraphZoomView({
       // Only dispatch the zoom action if the new clicked node is not the same as the old selected node.
       // This essentialy tracks double click action on a rectangle
       if (lastInteraction === 'click') {
-        if (hoveredNode && selectedNode && hoveredNode === selectedNode) {
+        if (
+          hoveredNode &&
+          flamegraphState.profiles.selectedNode &&
+          hoveredNode === flamegraphState.profiles.selectedNode
+        ) {
           canvasPoolManager.dispatch('zoomIntoFrame', [hoveredNode]);
         }
         canvasPoolManager.dispatch('selectedNode', [hoveredNode]);
-        setSelectedNode(hoveredNode);
+        dispatchFlamegraphState({type: 'set selected node', payload: hoveredNode});
       }
 
       setLastInteraction(null);
@@ -470,7 +476,8 @@ function FlamegraphZoomView({
     [
       flamegraphRenderer,
       configSpaceCursor,
-      selectedNode,
+      flamegraphState.profiles.selectedNode,
+      dispatchFlamegraphState,
       hoveredNode,
       canvasPoolManager,
       lastInteraction,

static/app/components/profiling/profileDragDropImport.tsx

@@ -7,7 +7,7 @@ import {
   ProfileGroup,
 } from 'sentry/utils/profiling/profile/importProfile';
 
-interface ProfileImportProps {
+export interface ProfileDragDropImportProps {
   children: React.ReactNode;
   onImport: (profile: ProfileGroup) => void;
 }
@@ -15,7 +15,7 @@ interface ProfileImportProps {
 function ProfileDragDropImport({
   onImport,
   children,
-}: ProfileImportProps): React.ReactElement {
+}: ProfileDragDropImportProps): React.ReactElement {
   const [dropState, setDropState] = React.useState<'idle' | 'dragover' | 'processing'>(
     'idle'
   );

static/app/utils/profiling/flamegraph/flamegraphStateProvider/flamegraphProfiles.tsx

@@ -1,19 +1,37 @@
-type FlamegraphProfilesgAction = {
+import {FlamegraphFrame} from 'sentry/utils/profiling/flamegraphFrame';
+
+type SetProfilesActiveIndex = {
   payload: number;
   type: 'set active profile index';
 };
 
-type FlamegraphProfiles = {
+type SetSelectedNode = {
+  payload: FlamegraphFrame | null;
+  type: 'set selected node';
+};
+
+type FlamegraphProfilesAction = SetProfilesActiveIndex | SetSelectedNode;
+
+type FlamegraphProfilesState = {
   activeProfileIndex: number | null;
+  selectedNode: FlamegraphFrame | null;
 };
 
 export function flamegraphProfilesReducer(
-  state: FlamegraphProfiles,
-  action: FlamegraphProfilesgAction
-): FlamegraphProfiles {
+  state: FlamegraphProfilesState,
+  action: FlamegraphProfilesAction
+): FlamegraphProfilesState {
   switch (action.type) {
+    case 'set selected node': {
+      return {...state, selectedNode: action.payload};
+    }
     case 'set active profile index': {
-      return {activeProfileIndex: action.payload};
+      // When the profile index changes, we want to drop the selected and hovered nodes
+      return {
+        ...state,
+        selectedNode: null,
+        activeProfileIndex: action.payload,
+      };
     }
     default: {
       return state;

static/app/utils/profiling/flamegraph/flamegraphStateProvider/index.tsx

@@ -163,6 +163,7 @@ interface FlamegraphStateProviderProps {
 const DEFAULT_FLAMEGRAPH_STATE: FlamegraphState = {
   profiles: {
     activeProfileIndex: null,
+    selectedNode: null,
   },
   position: {
     view: Rect.Empty(),
@@ -185,6 +186,7 @@ export function FlamegraphStateProvider(
 ): React.ReactElement {
   const reducer = useUndoableReducer(combinedReducers, {
     profiles: {
+      selectedNode: null,
       activeProfileIndex:
         props.initialState?.profiles?.activeProfileIndex ??
         DEFAULT_FLAMEGRAPH_STATE.profiles.activeProfileIndex,

static/app/utils/profiling/flamegraph/useFlamegraphProfiles.tsx

@@ -0,0 +1,29 @@
+import {useContext} from 'react';
+
+import {
+  FlamegraphStateContext,
+  FlamegraphStateContextValue,
+} from './flamegraphStateProvider/index';
+
+export function useFlamegraphProfiles(): [
+  FlamegraphStateContextValue[0]['profiles'],
+  FlamegraphStateContextValue[1]
+] {
+  const context = useContext(FlamegraphStateContext);
+
+  if (context === null) {
+    throw new Error('useFlamegraphProfiles called outside of FlamegraphStateProvider');
+  }
+
+  return [context[0].profiles, context[1]];
+}
+
+export function useFlamegraphProfilesValue(): FlamegraphStateContextValue[0]['profiles'] {
+  const context = useContext(FlamegraphStateContext);
+
+  if (context === null) {
+    throw new Error('useFlamegraphProfiles called outside of FlamegraphStateProvider');
+  }
+
+  return context[0].profiles;
+}

static/app/views/profiling/flamegraph.tsx

@@ -9,6 +9,7 @@ import * as Layout from 'sentry/components/layouts/thirds';
 import LoadingIndicator from 'sentry/components/loadingIndicator';
 import {Breadcrumb} from 'sentry/components/profiling/breadcrumb';
 import {Flamegraph} from 'sentry/components/profiling/flamegraph';
+import {ProfileDragDropImportProps} from 'sentry/components/profiling/profileDragDropImport';
 import SentryDocumentTitle from 'sentry/components/sentryDocumentTitle';
 import {t} from 'sentry/locale';
 import {Organization, Project} from 'sentry/types';
@@ -84,23 +85,10 @@ function FlamegraphView(props: FlamegraphViewProps): React.ReactElement {
   });
 
   const initialFlamegraphPreferencesState = useMemo((): DeepPartial<FlamegraphState> => {
-    if (requestState.type !== 'resolved') {
-      return decodeFlamegraphStateFromQueryParams(location.query);
-    }
-
-    const decodedState = decodeFlamegraphStateFromQueryParams(location.query);
-
-    return {
-      ...decodedState,
-      profiles: {
-        // We either initialize from query param or from the response
-        activeProfileIndex:
-          decodedState?.profiles?.activeProfileIndex ??
-          requestState.data.activeProfileIndex ??
-          0,
-      },
-    };
-  }, [requestState, location.query]);
+    return decodeFlamegraphStateFromQueryParams(location.query);
+    // We only want to decode this when our component mounts
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, []);
 
   useEffect(() => {
     if (!props.params.eventId || !props.params.projectId) {
@@ -124,6 +112,10 @@ function FlamegraphView(props: FlamegraphViewProps): React.ReactElement {
     };
   }, [props.params.eventId, props.params.projectId, api, organization]);
 
+  const onImport: ProfileDragDropImportProps['onImport'] = profiles => {
+    setRequestState({type: 'resolved', data: profiles});
+  };
+
   return (
     <SentryDocumentTitle title={t('Profiling')} orgSlug={organization.slug}>
       <FlamegraphStateProvider initialState={initialFlamegraphPreferencesState}>
@@ -157,13 +149,13 @@ function FlamegraphView(props: FlamegraphViewProps): React.ReactElement {
                 </Alert>
               ) : requestState.type === 'loading' ? (
                 <Fragment>
-                  <Flamegraph profiles={LoadingGroup} />
+                  <Flamegraph onImport={onImport} profiles={LoadingGroup} />
                   <LoadingIndicatorContainer>
                     <LoadingIndicator />
                   </LoadingIndicatorContainer>
                 </Fragment>
               ) : requestState.type === 'resolved' ? (
-                <Flamegraph profiles={requestState.data} />
+                <Flamegraph onImport={onImport} profiles={requestState.data} />
               ) : null}
             </FlamegraphContainer>
           </FlamegraphThemeProvider>

tests/js/spec/utils/useUndoableReducer.spec.tsx

@@ -2,8 +2,6 @@ import {useReducer} from 'react';
 
 import {reactHooks} from 'sentry-test/reactTestingLibrary';
 
-import {combinedReducers} from 'sentry/utils/profiling/flamegraph/flamegraphStateProvider';
-import {Rect} from 'sentry/utils/profiling/gl/utils';
 import {makeCombinedReducers} from 'sentry/utils/useCombinedReducer';
 import {
   makeUndoableReducer,
@@ -166,39 +164,27 @@ describe('makeUndoableReducer', () => {
   });
 
   it('can work with objects', () => {
+    const combinedReducers = makeCombinedReducers({
+      math: (state: number, action: {type: 'add'} | {type: 'subtract'}) =>
+        action.type === 'add' ? state + 1 : state - 1,
+    });
     const {result} = reactHooks.renderHook(() =>
       useReducer(makeUndoableReducer(combinedReducers), {
         previous: undefined,
         current: {
-          profiles: {
-            activeProfileIndex: 0,
-          },
-          position: {view: Rect.Empty()},
-          preferences: {
-            colorCoding: 'by symbol name',
-            sorting: 'call order',
-            view: 'top down',
-            xAxis: 'standalone',
-          },
-          search: {
-            index: null,
-            results: null,
-            query: '',
-          },
+          math: 0,
         },
         next: undefined,
       })
     );
 
-    reactHooks.act(() =>
-      result.current[1]({type: 'set color coding', payload: 'by library'})
-    );
-    expect(result.current[0].current.preferences.colorCoding).toBe('by library');
+    reactHooks.act(() => result.current[1]({type: 'add'}));
+    expect(result.current[0].current.math).toBe(1);
 
     reactHooks.act(() => result.current[1]({type: 'undo'}));
-    expect(result.current[0].current.preferences.colorCoding).toBe('by symbol name');
+    expect(result.current[0].current.math).toBe(0);
 
     reactHooks.act(() => result.current[1]({type: 'redo'}));
-    expect(result.current[0].current.preferences.colorCoding).toBe('by library');
+    expect(result.current[0].current.math).toBe(1);
   });
 });