Move useFullAssetGraphData to use webworker too (#26385)

Similar to #26315, moves
`buildGraphData` for the `useFullAssetGraphData` hook into a webworker.


## Test plan

Tested using app-proxy and loading a big asset graph

js_modules/dagster-ui/packages/ui-core/jest/mocks/ComputeGraphData.worker.ts

@@ -1,6 +1,8 @@
 import {setFeatureFlagsInternal} from '../../src/app/Flags';
+import {assertUnreachable} from '../../src/app/Util';
 import {computeGraphData} from '../../src/asset-graph/ComputeGraphData';
-import {ComputeGraphDataMessageType} from '../../src/asset-graph/ComputeGraphData.types';
+import {ComputeGraphDataWorkerMessageType} from '../../src/asset-graph/ComputeGraphData.types';
+import {buildGraphData} from '../../src/asset-graph/Utils';
 
 // eslint-disable-next-line import/no-default-export
 export default class MockWorker {
@@ -18,13 +20,19 @@ export default class MockWorker {
   }
 
   // mock expects data: { } instead of e: { data: { } }
-  async postMessage(data: ComputeGraphDataMessageType) {
+  async postMessage(data: ComputeGraphDataWorkerMessageType) {
     if (data.type === 'computeGraphData') {
       if (data.flagAssetSelectionSyntax) {
         setFeatureFlagsInternal({flagAssetSelectionSyntax: true});
       }
       const state = await computeGraphData(data);
       this.onmessage.forEach((onmessage) => onmessage({data: {...state, id: data.id}}));
+    } else if (data.type === 'buildGraphData') {
+      this.onmessage.forEach((onmessage) =>
+        onmessage({data: {...buildGraphData(data.nodes), id: data.id}}),
+      );
+    } else {
+      assertUnreachable(data);
     }
   }
 }

js_modules/dagster-ui/packages/ui-core/src/asset-graph/ComputeGraphData.types.ts

@@ -1,13 +1,25 @@
 import {AssetNodeForGraphQueryFragment} from './types/useAssetGraphData.types';
 import {AssetGraphFetchScope, AssetGraphQueryItem} from './useAssetGraphData';
 
-export type ComputeGraphDataMessageType = {
+type BaseType = {
   id: number;
+  flagAssetSelectionSyntax?: boolean;
+};
+
+export type ComputeGraphDataMessageType = BaseType & {
   type: 'computeGraphData';
   repoFilteredNodes?: AssetNodeForGraphQueryFragment[];
   graphQueryItems?: AssetGraphQueryItem[];
   opsQuery: string;
   kinds: AssetGraphFetchScope['kinds'];
   hideEdgesToNodesOutsideQuery?: boolean;
-  flagAssetSelectionSyntax?: boolean;
 };
+
+export type BuildGraphDataMessageType = BaseType & {
+  nodes: AssetNodeForGraphQueryFragment[];
+  type: 'buildGraphData';
+};
+
+export type ComputeGraphDataWorkerMessageType =
+  | ComputeGraphDataMessageType
+  | BuildGraphDataMessageType;

js_modules/dagster-ui/packages/ui-core/src/asset-graph/ComputeGraphData.worker.ts

@@ -1,20 +1,27 @@
 import {FeatureFlag} from 'shared/app/FeatureFlags.oss';
 
 import {computeGraphData} from './ComputeGraphData';
-import {ComputeGraphDataMessageType} from './ComputeGraphData.types';
+import {BuildGraphDataMessageType, ComputeGraphDataMessageType} from './ComputeGraphData.types';
+import {buildGraphData} from './Utils';
 import {setFeatureFlags} from '../app/Flags';
+import {assertUnreachable} from '../app/Util';
 
-type WorkerMessageData = ComputeGraphDataMessageType;
+type WorkerMessageData = ComputeGraphDataMessageType | BuildGraphDataMessageType;
 
 self.addEventListener('message', async (event: MessageEvent & {data: WorkerMessageData}) => {
-  const {data} = event;
+  const data: WorkerMessageData = event.data;
+
+  if (data.flagAssetSelectionSyntax) {
+    setFeatureFlags({[FeatureFlag.flagAssetSelectionSyntax]: true});
+  }
 
   if (data.type === 'computeGraphData') {
-    if (data.flagAssetSelectionSyntax) {
-      setFeatureFlags({[FeatureFlag.flagAssetSelectionSyntax]: true});
-    }
     const state = await computeGraphData(data);
     self.postMessage({...state, id: data.id});
+  } else if (data.type === 'buildGraphData') {
+    self.postMessage({...buildGraphData(data.nodes), id: data.id});
+  } else {
+    assertUnreachable(data);
   }
 });
 

js_modules/dagster-ui/packages/ui-core/src/asset-graph/useAssetGraphData.tsx

@@ -5,10 +5,10 @@ import {useEffect, useMemo, useRef, useState} from 'react';
 import {FeatureFlag} from 'shared/app/FeatureFlags.oss';
 
 import {ASSET_NODE_FRAGMENT} from './AssetNode';
-import {GraphData, buildGraphData, tokenForAssetKey} from './Utils';
+import {GraphData, buildGraphData as buildGraphDataImpl, tokenForAssetKey} from './Utils';
 import {gql} from '../apollo-client';
 import {computeGraphData as computeGraphDataImpl} from './ComputeGraphData';
-import {ComputeGraphDataMessageType} from './ComputeGraphData.types';
+import {BuildGraphDataMessageType, ComputeGraphDataMessageType} from './ComputeGraphData.types';
 import {throttleLatest} from './throttleLatest';
 import {featureEnabled} from '../app/Flags';
 import {
@@ -61,12 +61,33 @@ export function useFullAssetGraphData(options: AssetGraphFetchScope) {
   });
 
   const nodes = fetchResult.data?.assetNodes;
-  const queryItems = useMemo(() => (nodes ? buildGraphQueryItems(nodes) : []), [nodes]);
-
-  const fullAssetGraphData = useMemo(
-    () => (queryItems ? buildGraphData(queryItems.map((n) => n.node)) : null),
-    [queryItems],
+  const queryItems = useMemo(
+    () => (nodes ? buildGraphQueryItems(nodes) : []).map(({node}) => node),
+    [nodes],
   );
+
+  const [fullAssetGraphData, setFullAssetGraphData] = useState<GraphData | null>(null);
+  useBlockTraceUntilTrue('FullAssetGraphData', !!fullAssetGraphData);
+
+  const lastProcessedRequestRef = useRef(0);
+  const currentRequestRef = useRef(0);
+
+  useEffect(() => {
+    if (options.loading) {
+      return;
+    }
+    const requestId = ++currentRequestRef.current;
+    buildGraphData({
+      nodes: queryItems,
+      flagAssetSelectionSyntax: featureEnabled(FeatureFlag.flagAssetSelectionSyntax),
+    })?.then((data) => {
+      if (lastProcessedRequestRef.current < requestId) {
+        lastProcessedRequestRef.current = requestId;
+        setFullAssetGraphData(data);
+      }
+    });
+  }, [options.loading, queryItems]);
+
   return fullAssetGraphData;
 }
 
@@ -317,14 +338,16 @@ const computeGraphData = throttleLatest(
   2000,
 );
 
-const getWorker = memoize(() => new Worker(new URL('./ComputeGraphData.worker', import.meta.url)));
+const getWorker = memoize(
+  (_key: string = '') => new Worker(new URL('./ComputeGraphData.worker', import.meta.url)),
+);
 
 let _id = 0;
 async function computeGraphDataWrapper(
   props: Omit<ComputeGraphDataMessageType, 'id' | 'type'>,
 ): Promise<GraphDataState> {
   if (featureEnabled(FeatureFlag.flagAssetSelectionWorker)) {
-    const worker = getWorker();
+    const worker = getWorker('computeGraphWorker');
     return new Promise<GraphDataState>((resolve) => {
       const id = ++_id;
       const callback = (event: MessageEvent) => {
@@ -345,3 +368,39 @@ async function computeGraphDataWrapper(
   }
   return computeGraphDataImpl(props);
 }
+
+const buildGraphData = throttleLatest(
+  indexedDBAsyncMemoize<BuildGraphDataMessageType, GraphData, typeof buildGraphDataWrapper>(
+    buildGraphDataWrapper,
+    (props) => {
+      return JSON.stringify(props);
+    },
+  ),
+  2000,
+);
+
+async function buildGraphDataWrapper(
+  props: Omit<BuildGraphDataMessageType, 'id' | 'type'>,
+): Promise<GraphData> {
+  if (featureEnabled(FeatureFlag.flagAssetSelectionWorker)) {
+    const worker = getWorker('buildGraphWorker');
+    return new Promise<GraphData>((resolve) => {
+      const id = ++_id;
+      const callback = (event: MessageEvent) => {
+        const data = event.data as GraphData & {id: number};
+        if (data.id === id) {
+          resolve(data);
+          worker.removeEventListener('message', callback);
+        }
+      };
+      worker.addEventListener('message', callback);
+      const message: BuildGraphDataMessageType = {
+        type: 'buildGraphData',
+        id,
+        ...props,
+      };
+      worker.postMessage(message);
+    });
+  }
+  return buildGraphDataImpl(props.nodes);
+}