Add TraceContext class for tracking asynchronous dependencies for pag…

…e load. (#21492)

## Summary & Motivation

Similar in spirit to #21414.
It provides a mechanism for Cloud to use to track outstanding
asynchronous dependencies for the current page.

In particular we need to know:

- When a dependency is created/added
- When a dependency is completed
- When a dependency is cancelled (eg. it didn't complete but we're
aborting it because we're navigating away)

It's important to distinguish between completed and "cancelled" loads
for understanding the metrics.

## How I Tested These Changes

jest tests

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

@@ -14,6 +14,7 @@ import {
 import {GraphQueryItem, filterByQuery} from '../app/GraphQueryImpl';
 import {AssetKey} from '../assets/types';
 import {AssetGroupSelector, PipelineSelector} from '../graphql/types';
+import {useBlockTraceOnQueryResult} from '../performance/TraceContext';
 
 export interface AssetGraphFetchScope {
   hideEdgesToNodesOutsideQuery?: boolean;
@@ -44,6 +45,7 @@ export function useAssetGraphData(opsQuery: string, options: AssetGraphFetchScop
       groupSelector: options.groupSelector,
     },
   });
+  useBlockTraceOnQueryResult(fetchResult, 'ASSET_GRAPH_QUERY');
 
   const nodes = fetchResult.data?.assetNodes;
 

js_modules/dagster-ui/packages/ui-core/src/performance/TraceContext.tsx

@@ -0,0 +1,95 @@
+import {QueryResult} from '@apollo/client';
+import {ReactNode, createContext, useContext, useLayoutEffect, useMemo} from 'react';
+
+export enum CompletionType {
+  SUCCESS = 1,
+  ERROR = 2,
+  CANCELLED = 3,
+}
+
+type TraceContextType = {
+  createDependency: (_name: string) => Dependency | null;
+  addDependency: (_dep: Dependency | null) => void;
+  completeDependency: (_dep: Dependency | null, type: CompletionType) => void;
+};
+
+export const TraceContext = createContext<TraceContextType>({
+  createDependency: (_name: string) => null,
+  addDependency: (_dep) => {},
+  completeDependency: (_dep, _type) => {},
+});
+
+/**
+ * Use this to wrap child react components who should not count
+ * toward display done. Eg. If you're re-using a component that
+ * adds dependencies but you don't want that component or its dependencies
+ * as your own dependency
+ */
+export const OrphanDependenciesTraceContext = ({children}: {children: ReactNode}) => {
+  return (
+    <TraceContext.Provider
+      value={useMemo(
+        () => ({
+          createDependency: () => null,
+          addDependency: () => {},
+          completeDependency: () => {},
+        }),
+        [],
+      )}
+    >
+      {children}
+    </TraceContext.Provider>
+  );
+};
+
+export class Dependency {
+  public readonly name: string;
+
+  constructor(name: string) {
+    this.name = name;
+  }
+}
+
+/** Use this to declare a dependency on an apollo query result */
+export function useBlockTraceOnQueryResult(queryResult: QueryResult<any>, name: string) {
+  const dep = useDependency(name);
+  const hasData = !!queryResult.data;
+  const hasError = !!queryResult.error;
+
+  useLayoutEffect(() => {
+    if (hasData) {
+      dep.completeDependency(CompletionType.SUCCESS);
+    }
+  }, [dep, hasData]);
+
+  useLayoutEffect(() => {
+    if (!hasData && hasError) {
+      dep.completeDependency(CompletionType.ERROR);
+    }
+  }, [dep, hasData, hasError]);
+}
+
+export function useDependency(name: string) {
+  const {addDependency, completeDependency, createDependency} = useContext(TraceContext);
+
+  const dependency = useMemo(() => createDependency(name), [createDependency, name]);
+
+  useLayoutEffect(() => {
+    addDependency(dependency);
+    return () => {
+      // By default cancel a dependency when the component unmounts.
+      // Rely on the user of TraceContext to track if the dependency
+      // was already completed prior to this.
+      completeDependency(dependency, CompletionType.CANCELLED);
+    };
+  }, [addDependency, completeDependency, dependency]);
+
+  return useMemo(
+    () => ({
+      completeDependency: (type: CompletionType) => {
+        completeDependency(dependency, type);
+      },
+    }),
+    [completeDependency, dependency],
+  );
+}