From fb8e87e073d5b63a069661c8aa44b1c77991c5d5 Mon Sep 17 00:00:00 2001 From: Marco Salazar Date: Thu, 9 Jan 2025 10:06:52 -0500 Subject: [PATCH] fix workspace context overfetching on cascading updates --- .../WorkspaceContext/WorkspaceContext.tsx | 85 +++++++++++++------ .../__fixtures__/Workspace.fixtures.ts | 12 ++- .../__tests__/WorkspaceContext.test.tsx | 35 ++++++++ 3 files changed, 103 insertions(+), 29 deletions(-) diff --git a/js_modules/dagster-ui/packages/ui-core/src/workspace/WorkspaceContext/WorkspaceContext.tsx b/js_modules/dagster-ui/packages/ui-core/src/workspace/WorkspaceContext/WorkspaceContext.tsx index cd60b1e776c55..7883a48c23939 100644 --- a/js_modules/dagster-ui/packages/ui-core/src/workspace/WorkspaceContext/WorkspaceContext.tsx +++ b/js_modules/dagster-ui/packages/ui-core/src/workspace/WorkspaceContext/WorkspaceContext.tsx @@ -96,6 +96,7 @@ interface RefreshLocationsIfNeededParams React.SetStateAction> >; previousLocationVersionsRef: React.MutableRefObject>; + fetchingStatusesRef: React.MutableRefObject>>; } interface HandleDeletedLocationsParams @@ -113,6 +114,8 @@ interface FetchLocationDataParams setLocationEntryData: React.Dispatch< React.SetStateAction> >; + locationStatuses: Record; + fetchingStatusesRef: React.MutableRefObject>>; } const UNLOADED_CACHED_DATA = {}; @@ -129,6 +132,7 @@ export const WorkspaceProvider = ({children}: {children: React.ReactNode}) => { ); const previousLocationVersionsRef = useRef>({}); + const fetchingStatusesRef = useRef>>({}); const {locationStatuses, loading: loadingLocationStatuses} = useCodeLocationStatuses(localCacheIdPrefix); @@ -162,6 +166,7 @@ export const WorkspaceProvider = ({children}: {children: React.ReactNode}) => { getData, setLocationEntryData, previousLocationVersionsRef, + fetchingStatusesRef, }; refreshLocationsIfNeeded(params); }, [locationStatuses, locationEntryData, client, localCacheIdPrefix, getData, loadingCachedData]); @@ -197,6 +202,8 @@ export const WorkspaceProvider = ({children}: {children: React.ReactNode}) => { localCacheIdPrefix, getData, setLocationEntryData, + locationStatuses, + fetchingStatusesRef, }; return await fetchLocationData(params); }); @@ -268,6 +275,7 @@ async function refreshLocationsIfNeeded(params: RefreshLocationsIfNeededParams): getData, setLocationEntryData, previousLocationVersionsRef, + fetchingStatusesRef, } = params; const locationsToFetch = identifyStaleLocations( @@ -281,18 +289,20 @@ async function refreshLocationsIfNeeded(params: RefreshLocationsIfNeededParams): } await Promise.all( - locationsToFetch.map((location) => - fetchLocationData({ + locationsToFetch.map(async (location) => { + await fetchLocationData({ name: location.name, client, localCacheIdPrefix, getData, setLocationEntryData, - }), - ), + locationStatuses, + fetchingStatusesRef, + }); + previousLocationVersionsRef.current[location.name] = + locationStatuses[location.name]?.versionKey ?? ''; + }), ); - - previousLocationVersionsRef.current = mapLocationVersions(locationStatuses); } /** @@ -339,29 +349,49 @@ function handleDeletedLocations(params: HandleDeletedLocationsParams): void { async function fetchLocationData( params: FetchLocationDataParams, ): Promise { - const {name, client, localCacheIdPrefix, getData, setLocationEntryData} = params; - try { - const {data} = await getData({ - client, - query: LOCATION_WORKSPACE_QUERY, - key: `${localCacheIdPrefix}${locationWorkspaceKey(name)}`, - version: LocationWorkspaceQueryVersion, - variables: {name}, - bypassCache: true, - }); + const { + name, + client, + localCacheIdPrefix, + getData, + setLocationEntryData, + locationStatuses, + fetchingStatusesRef, + } = params; - const entry = data?.workspaceLocationEntryOrError; - if (entry) { - setLocationEntryData((prevData) => ({ - ...prevData, - [name]: entry, - })); - } - return data; - } catch (error) { - console.error(`Error fetching location data for ${name}:`, error); + const localKey = name + '@' + (locationStatuses[name]?.versionKey ?? ''); + if (fetchingStatusesRef.current[localKey]) { + return fetchingStatusesRef.current[localKey]; } - return undefined; + fetchingStatusesRef.current[localKey] = new Promise(async (res) => { + try { + const {data} = await getData({ + client, + query: LOCATION_WORKSPACE_QUERY, + key: `${localCacheIdPrefix}${locationWorkspaceKey(name)}`, + version: LocationWorkspaceQueryVersion, + variables: {name}, + bypassCache: true, + }); + + const entry = data?.workspaceLocationEntryOrError; + if (entry) { + setLocationEntryData((prevData) => ({ + ...prevData, + [name]: entry, + })); + } + res(data); + } catch (error) { + console.error(`Error fetching location data for ${name}:`, error); + } + res(undefined); + }); + const result = await fetchingStatusesRef.current[localKey]; + requestAnimationFrame(() => { + delete fetchingStatusesRef.current[localKey]; + }); + return result; } /** @@ -422,6 +452,7 @@ function identifyStaleLocations( const entry = locationEntryData[statusEntry.name]; const locationEntry = entry?.__typename === 'WorkspaceLocationEntry' ? entry : null; const dataVersion = locationEntry?.versionKey || ''; + return currentVersion !== prevVersion || currentVersion !== dataVersion; }); } diff --git a/js_modules/dagster-ui/packages/ui-core/src/workspace/WorkspaceContext/__fixtures__/Workspace.fixtures.ts b/js_modules/dagster-ui/packages/ui-core/src/workspace/WorkspaceContext/__fixtures__/Workspace.fixtures.ts index 8f8f12c85850a..e528c6b987f0f 100644 --- a/js_modules/dagster-ui/packages/ui-core/src/workspace/WorkspaceContext/__fixtures__/Workspace.fixtures.ts +++ b/js_modules/dagster-ui/packages/ui-core/src/workspace/WorkspaceContext/__fixtures__/Workspace.fixtures.ts @@ -33,7 +33,9 @@ export const buildCodeLocationsStatusQuery = ( export const buildWorkspaceMocks = ( entries: WorkspaceLocationEntry[], - options: Partial> = {}, + options: Partial> & { + cascadingUpdates?: boolean; + } = {}, ) => { return [ buildCodeLocationsStatusQuery( @@ -41,12 +43,13 @@ export const buildWorkspaceMocks = ( buildWorkspaceLocationStatusEntry({ ...entry, updateTimestamp: entry.updatedTimestamp, + versionKey: '' + entry.updatedTimestamp, __typename: 'WorkspaceLocationStatusEntry', }), ), options, ), - ...entries.map((entry) => + ...entries.map((entry, index) => buildQueryMock({ query: LOCATION_WORKSPACE_QUERY, variables: { @@ -56,6 +59,11 @@ export const buildWorkspaceMocks = ( workspaceLocationEntryOrError: entry, }, ...options, + ...(options.cascadingUpdates + ? { + delay: 100 * (1 + index), + } + : {}), }), ), ]; diff --git a/js_modules/dagster-ui/packages/ui-core/src/workspace/WorkspaceContext/__tests__/WorkspaceContext.test.tsx b/js_modules/dagster-ui/packages/ui-core/src/workspace/WorkspaceContext/__tests__/WorkspaceContext.test.tsx index f5714df54a5a9..dddbdeadb6d46 100644 --- a/js_modules/dagster-ui/packages/ui-core/src/workspace/WorkspaceContext/__tests__/WorkspaceContext.test.tsx +++ b/js_modules/dagster-ui/packages/ui-core/src/workspace/WorkspaceContext/__tests__/WorkspaceContext.test.tsx @@ -510,4 +510,39 @@ describe('WorkspaceContext', () => { expect(result.current.allRepos).toEqual([]); expect(result.current.data).toEqual({}); }); + + it('Doesnt overfetch the same code location version on cascading location query responses', async () => { + const {location1, location2, location3, caches} = getLocationMocks(-1); // Initialize with outdated cache + + caches.codeLocationStatusQuery.has.mockResolvedValue(false); + caches.location1.has.mockResolvedValue(false); + + const mocks = buildWorkspaceMocks([location1, location2, location3], { + cascadingUpdates: true, + maxUsageCount: 9999, + }); + const mockCbs = mocks.map(getMockResultFn); + + // Include code location status mock a second time since we call runOnlyPendingTimersAsync twice + const {result} = renderWithMocks([...mocks, mocks[0]!]); + + // Initial state + expect(result.current.allRepos).toEqual([]); + expect(result.current.data).toEqual({}); + expect(result.current.loading).toEqual(true); + + await waitFor(() => { + expect(result.current.loading).toEqual(false); + }); + + // Ensure no additional fetches were made + expect(mockCbs[1]).toHaveBeenCalledTimes(1); + expect(mockCbs[2]).toHaveBeenCalledTimes(1); + expect(mockCbs[3]).toHaveBeenCalledTimes(1); + + await act(async () => { + // Exhaust any remaining tasks so they don't affect the next test. + await jest.runAllTicks(); + }); + }); });