From 0c02efaf0f546eb9a75419d51c10db4b59f7b094 Mon Sep 17 00:00:00 2001 From: Marco polo Date: Mon, 29 Apr 2024 16:33:05 -0400 Subject: [PATCH] Add TraceContext class for tracking asynchronous dependencies for page load. (#21492) ## Summary & Motivation Similar in spirit to https://github.com/dagster-io/dagster/pull/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 --- .../ui-core/src/app/TrackedSuspense.tsx | 100 ---------------- .../app/__tests__/TrackedSuspense.test.tsx | 110 ------------------ .../src/asset-graph/useAssetGraphData.tsx | 2 + .../ui-core/src/performance/TraceContext.tsx | 95 +++++++++++++++ .../__tests__/TraceContext.test.tsx | 106 +++++++++++++++++ 5 files changed, 203 insertions(+), 210 deletions(-) delete mode 100644 js_modules/dagster-ui/packages/ui-core/src/app/TrackedSuspense.tsx delete mode 100644 js_modules/dagster-ui/packages/ui-core/src/app/__tests__/TrackedSuspense.test.tsx create mode 100644 js_modules/dagster-ui/packages/ui-core/src/performance/TraceContext.tsx create mode 100644 js_modules/dagster-ui/packages/ui-core/src/performance/__tests__/TraceContext.test.tsx diff --git a/js_modules/dagster-ui/packages/ui-core/src/app/TrackedSuspense.tsx b/js_modules/dagster-ui/packages/ui-core/src/app/TrackedSuspense.tsx deleted file mode 100644 index b54df600f51ad..0000000000000 --- a/js_modules/dagster-ui/packages/ui-core/src/app/TrackedSuspense.tsx +++ /dev/null @@ -1,100 +0,0 @@ -import { - ComponentProps, - ReactNode, - Suspense, - createContext, - useCallback, - useContext, - useLayoutEffect, - useMemo, -} from 'react'; - -// Context to store the active suspense count and a method to manipulate it -const TrackedSuspenseContext = createContext({ - useBoundary: (_id: string) => { - return { - onFallbackRendered: () => { - return () => {}; - }, - onContentRendered: () => { - return () => {}; - }, - }; - }, -}); - -export type Boundary = { - id: string; - name: string; -}; - -export const TrackedSuspenseProvider = ({ - children, - onContentRendered, - onContentRemoved, - onFallbackRendered, - onFallbackRemoved, -}: { - children: ReactNode; - onContentRendered: (boundary: Boundary) => void; - onContentRemoved: (boundary: Boundary) => void; - onFallbackRendered: (boundary: Boundary) => void; - onFallbackRemoved: (boundary: Boundary) => void; -}) => { - const useBoundary = useCallback( - (name: string) => { - // eslint-disable-next-line react-hooks/rules-of-hooks - const boundary: Boundary = useMemo(() => ({name, id: uniqueId()}), [name]); - - return { - onFallbackRendered: () => { - onFallbackRendered(boundary); - return () => { - onFallbackRemoved(boundary); - }; - }, - onContentRendered: () => { - onContentRendered(boundary); - return () => { - onContentRemoved(boundary); - }; - }, - }; - }, - [onContentRemoved, onContentRendered, onFallbackRemoved, onFallbackRendered], - ); - - return ( - ({useBoundary}), [useBoundary])}> - {children} - - ); -}; - -export const TrackedSuspense = (props: ComponentProps & {id: string}) => { - const {useBoundary} = useContext(TrackedSuspenseContext); - const boundary = useBoundary(props.id); - return ( - {props.fallback}} - > - {props.children} - - ); -}; -const OnRendered = ({ - onRendered, - children, -}: { - children: React.ReactNode; - onRendered: () => () => void; -}) => { - useLayoutEffect(onRendered, [onRendered]); - return children; -}; - -let id = 0; -function uniqueId() { - return `boundaryId${id++}`; -} diff --git a/js_modules/dagster-ui/packages/ui-core/src/app/__tests__/TrackedSuspense.test.tsx b/js_modules/dagster-ui/packages/ui-core/src/app/__tests__/TrackedSuspense.test.tsx deleted file mode 100644 index ab41c13ace1b8..0000000000000 --- a/js_modules/dagster-ui/packages/ui-core/src/app/__tests__/TrackedSuspense.test.tsx +++ /dev/null @@ -1,110 +0,0 @@ -import {act, render, waitFor} from '@testing-library/react'; - -import {TrackedSuspense, TrackedSuspenseProvider} from '../TrackedSuspense'; - -describe('TrackedSuspenseProvider', () => { - it('manages boundaries correctly', async () => { - const mockContentRendered = jest.fn(); - const mockContentRemoved = jest.fn(); - const mockFallbackRendered = jest.fn(); - const mockFallbackRemoved = jest.fn(); - - const state = { - promises: {} as Record>, - resolvers: {} as Record void>, - resolved: {} as Record, - }; - - function Suspender({id}: {id: string}) { - if (!state.promises[id]) { - state.promises[id] = new Promise((res) => { - state.resolvers[id] = (value: any) => { - state.resolved[id] = true; - res(value); - }; - }); - } - if (!state.resolved[id]) { - throw state.promises[id]; - } - return id; - } - - render( - - Loading 1...}> - Loading 2...}> - - - Loading 3...}> - - - - Loading 4...}> - - - , - ); - - // test1 is rendered since it doesn't suspend - expect(mockContentRendered).toHaveBeenCalledTimes(1); - expect(mockContentRendered).toHaveBeenNthCalledWith(1, { - name: 'test1', - id: 'boundaryId0', - }); - - expect(mockFallbackRendered).toHaveBeenCalledTimes(3); - - expect(mockFallbackRendered).toHaveBeenNthCalledWith(1, { - name: 'test2', - id: 'boundaryId1', - }); - expect(mockFallbackRendered).toHaveBeenNthCalledWith(2, { - name: 'test3', - id: 'boundaryId2', - }); - expect(mockFallbackRendered).toHaveBeenNthCalledWith(3, { - name: 'test4', - id: 'boundaryId3', - }); - expect(mockFallbackRemoved).toHaveBeenCalledTimes(0); - act(() => { - state.resolvers['test2']!(1); - }); - await waitFor(() => { - expect(mockFallbackRemoved).toHaveBeenCalledTimes(1); - expect(mockContentRendered).toHaveBeenCalledTimes(2); - expect(mockContentRendered).toHaveBeenNthCalledWith(2, { - name: 'test2', - id: 'boundaryId1', - }); - }); - act(() => { - state.resolvers['test3']!(1); - }); - await waitFor(() => { - expect(mockFallbackRemoved).toHaveBeenCalledTimes(2); - expect(mockContentRendered).toHaveBeenCalledTimes(3); - expect(mockContentRendered).toHaveBeenNthCalledWith(3, { - name: 'test3', - id: 'boundaryId2', - }); - }); - act(() => { - state.resolvers['test4']!(1); - }); - await waitFor(() => { - expect(mockFallbackRemoved).toHaveBeenCalledTimes(3); - expect(mockContentRendered).toHaveBeenCalledTimes(4); - expect(mockContentRendered).toHaveBeenNthCalledWith(4, { - name: 'test4', - id: 'boundaryId3', - }); - }); - }); -}); diff --git a/js_modules/dagster-ui/packages/ui-core/src/asset-graph/useAssetGraphData.tsx b/js_modules/dagster-ui/packages/ui-core/src/asset-graph/useAssetGraphData.tsx index b30d651985fd2..904725a1db487 100644 --- a/js_modules/dagster-ui/packages/ui-core/src/asset-graph/useAssetGraphData.tsx +++ b/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; diff --git a/js_modules/dagster-ui/packages/ui-core/src/performance/TraceContext.tsx b/js_modules/dagster-ui/packages/ui-core/src/performance/TraceContext.tsx new file mode 100644 index 0000000000000..0e4143990fcaf --- /dev/null +++ b/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({ + 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 ( + ({ + createDependency: () => null, + addDependency: () => {}, + completeDependency: () => {}, + }), + [], + )} + > + {children} + + ); +}; + +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, 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], + ); +} diff --git a/js_modules/dagster-ui/packages/ui-core/src/performance/__tests__/TraceContext.test.tsx b/js_modules/dagster-ui/packages/ui-core/src/performance/__tests__/TraceContext.test.tsx new file mode 100644 index 0000000000000..bf3f036932444 --- /dev/null +++ b/js_modules/dagster-ui/packages/ui-core/src/performance/__tests__/TraceContext.test.tsx @@ -0,0 +1,106 @@ +import {QueryResult} from '@apollo/client'; +import {renderHook} from '@testing-library/react-hooks'; +import {ReactNode, useContext} from 'react'; + +import { + CompletionType, + Dependency, + OrphanDependenciesTraceContext, + TraceContext, + useBlockTraceOnQueryResult, + useDependency, +} from '../TraceContext'; + +describe('TraceContext', () => { + it('provides default functions that do not throw', () => { + const {result} = renderHook(() => useContext(TraceContext)); + expect(() => result.current.createDependency('test')).not.toThrow(); + expect(() => result.current.addDependency(null)).not.toThrow(); + expect(() => result.current.completeDependency(null, CompletionType.SUCCESS)).not.toThrow(); + }); +}); + +describe('useBlockTraceOnQueryResult', () => { + it('handles complete and error actions based on query data presence', () => { + const mockAddDependency = jest.fn(); + const mockCompleteDependency = jest.fn(); + const context = { + createDependency: () => new Dependency('test'), + addDependency: mockAddDependency, + completeDependency: mockCompleteDependency, + }; + const wrapper = ({children}: {children: ReactNode}) => ( + {children} + ); + + const {rerender} = renderHook( + ({queryResult}: {queryResult: QueryResult}) => + useBlockTraceOnQueryResult(queryResult, 'testDep'), + {initialProps: {queryResult: {data: null, error: null}}, wrapper} as any, + ); + + rerender({queryResult: {data: {}}} as any); + expect(mockCompleteDependency).toHaveBeenCalledTimes(1); + expect(mockCompleteDependency).toHaveBeenCalledWith( + expect.any(Dependency), + CompletionType.SUCCESS, + ); + + rerender({queryResult: {data: null, error: {}} as any}); + expect(mockCompleteDependency).toHaveBeenCalledTimes(2); + expect(mockCompleteDependency).toHaveBeenCalledWith( + expect.any(Dependency), + CompletionType.ERROR, + ); + }); +}); + +describe('useDependency', () => { + it('adds, completes, and cancels dependencies correctly', () => { + const mockAddDependency = jest.fn(); + const mockCompleteDependency = jest.fn(); + const wrapper = ({children}: {children: ReactNode}) => ( + new Dependency('test'), + addDependency: mockAddDependency, + completeDependency: mockCompleteDependency, + }} + > + {children} + + ); + + const {unmount} = renderHook(() => useDependency('testDep'), {wrapper}); + + expect(mockAddDependency).toHaveBeenCalledTimes(1); + unmount(); + expect(mockCompleteDependency).toHaveBeenCalledTimes(1); + expect(mockCompleteDependency).toHaveBeenCalledWith( + expect.any(Dependency), + CompletionType.CANCELLED, + ); + }); +}); + +describe('OrphanDependenciesTraceContext', () => { + it('ignores dependencies added within an orphaned tree', () => { + const mockAddDependency = jest.fn(); + const mockCompleteDependency = jest.fn(); + const context = { + createDependency: () => new Dependency('test'), + addDependency: mockAddDependency, + completeDependency: mockCompleteDependency, + }; + const wrapper = ({children}: {children: ReactNode}) => ( + + {children} + + ); + + const {unmount} = renderHook(() => useDependency('testDep'), {wrapper}); + unmount(); + expect(mockAddDependency).not.toHaveBeenCalled(); + expect(mockCompleteDependency).not.toHaveBeenCalled(); + }); +});