Skip to content

[perf] Share cache manager and add global cache for useAllAssets #29179

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed

[perf] Share cache manager and add global cache for useAllAssets #29179

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
153 changes: 84 additions & 69 deletions js_modules/dagster-ui/packages/ui-core/src/assets/AssetsCatalogTable.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ import {AssetGroupSelector} from '../graphql/types';
import {useUpdatingRef} from '../hooks/useUpdatingRef';
import {useBlockTraceUntilTrue} from '../performance/TraceContext';
import {fetchPaginatedData} from '../runs/fetchPaginatedBucketData';
import {CacheManager} from '../search/useIndexedDBCachedQuery';
import {getCacheManager} from '../search/useIndexedDBCachedQuery';
import {SyntaxError} from '../selection/CustomErrorListener';
import {LoadingSpinner} from '../ui/Loading';

Expand All @@ -54,7 +54,7 @@ export function useCachedAssets({
}) {
const {localCacheIdPrefix} = useContext(AppContext);
const cacheManager = useMemo(
() => new CacheManager<AssetTableFragment[]>(`${localCacheIdPrefix}/allAssetNodes`),
() => getCacheManager<AssetTableFragment[]>(`${localCacheIdPrefix}/allAssetNodes`),
[localCacheIdPrefix],
);

Expand All @@ -69,6 +69,11 @@ export function useCachedAssets({
return {cacheManager};
}

// Module-level cache variables
let globalAssetsPromise: Promise<Asset[]> | null = null;
let cachedAllAssets: Asset[] | null = null;
let cachedAssetsFetchTime: number = 0;

export function useAllAssets({
batchLimit = DEFAULT_BATCH_LIMIT,
groupSelector,
Expand All @@ -77,14 +82,16 @@ export function useAllAssets({
const [{error, assets}, setErrorAndAssets] = useState<{
error: PythonErrorFragment | undefined;
assets: Asset[] | undefined;
}>({error: undefined, assets: undefined});
}>({error: undefined, assets: cachedAllAssets || undefined});

const assetsRef = useUpdatingRef(assets);

const {cacheManager} = useCachedAssets({
onAssetsLoaded: useCallback(
(data) => {
console.log('onAssetsLoaded', data);
if (!assetsRef.current) {
console.log('onAssetsLoaded', data);
setErrorAndAssets({
error: undefined,
assets: data,
Expand All @@ -95,69 +102,57 @@ export function useAllAssets({
),
});

const allAssetsQuery = useCallback(async () => {
if (groupSelector) {
return;
// Query function for all assets
const fetchAllAssets = useCallback(async () => {
if (cachedAllAssets && Date.now() - cachedAssetsFetchTime < 6000) {
return cachedAllAssets;
}
try {
const data = await fetchPaginatedData({
async fetchData(cursor: string | null | undefined) {
const {data} = await client.query<
AssetCatalogTableQuery,
AssetCatalogTableQueryVariables
>({
query: ASSET_CATALOG_TABLE_QUERY,
fetchPolicy: 'no-cache',
variables: {
cursor,
limit: batchLimit,
},
});

if (data.assetsOrError.__typename === 'PythonError') {
return {
data: [],
cursor: undefined,
hasMore: false,
error: data.assetsOrError,
};
}
const assets = data.assetsOrError.nodes;
const hasMoreData = assets.length === batchLimit;
const nextCursor = data.assetsOrError.cursor;
return {
data: assets,
cursor: nextCursor,
hasMore: hasMoreData,
error: undefined,
};
},
});
cacheManager.set(data, AssetCatalogTableQueryVersion);
setErrorAndAssets({error: undefined, assets: data});
} catch (e: any) {
if (e.__typename === 'PythonError') {
setErrorAndAssets(({assets}) => ({
error: e,
assets,
}));
}
if (!globalAssetsPromise) {
globalAssetsPromise = (async () => {
const allAssets = await fetchPaginatedData({
async fetchData(cursor: string | null | undefined) {
const {data} = await client.query<
AssetCatalogTableQuery,
AssetCatalogTableQueryVariables
>({
query: ASSET_CATALOG_TABLE_QUERY,
fetchPolicy: 'no-cache',
variables: {cursor, limit: batchLimit},
});
if (data.assetsOrError.__typename === 'PythonError') {
return {
data: [],
cursor: undefined,
hasMore: false,
error: data.assetsOrError,
};
}
const assets = data.assetsOrError.nodes;
const hasMoreData = assets.length === batchLimit;
const nextCursor = data.assetsOrError.cursor;
return {data: assets, cursor: nextCursor, hasMore: hasMoreData, error: undefined};
},
});
cachedAssetsFetchTime = Date.now();
cachedAllAssets = allAssets;
cacheManager.set(allAssets, AssetCatalogTableQueryVersion);
globalAssetsPromise = null;
return allAssets;
})();
}
}, [batchLimit, cacheManager, client, groupSelector]);
return globalAssetsPromise;
}, [batchLimit, cacheManager, client]);

// Query function for group assets
const groupQuery = useCallback(async () => {
if (!groupSelector) {
return;
}
function onData(queryData: typeof data) {
const cacheKey = JSON.stringify(groupSelector);
if (groupTableCache.has(cacheKey)) {
const cachedData = groupTableCache.get(cacheKey);
setErrorAndAssets({
error: undefined,
assets: queryData.assetNodes?.map(definitionToAssetTableFragment),
assets: cachedData.assetNodes?.map(definitionToAssetTableFragment),
});
}
const cacheKey = JSON.stringify(groupSelector);
if (groupTableCache.has(cacheKey)) {
onData(groupTableCache.get(cacheKey));
return;
}
const {data} = await client.query<
AssetCatalogGroupTableQuery,
Expand All @@ -168,23 +163,43 @@ export function useAllAssets({
fetchPolicy: 'no-cache',
});
groupTableCache.set(cacheKey, data);
onData(data);
}, [groupSelector, client]);
setErrorAndAssets({
error: undefined,
assets: data.assetNodes?.map(definitionToAssetTableFragment),
});
}, [client, groupSelector]);

const query = groupSelector ? groupQuery : allAssetsQuery;
const query = groupSelector ? groupQuery : fetchAllAssets;

useEffect(() => {
query();
}, [query]);
if (groupSelector) {
groupQuery();
} else {
fetchAllAssets()
.then((allAssets) => setErrorAndAssets({error: undefined, assets: allAssets}))
.catch((e: any) => {
if (e.__typename === 'PythonError') {
setErrorAndAssets((prev) => ({error: e, assets: prev.assets}));
}
});
}
}, [fetchAllAssets, groupQuery, groupSelector]);

return useMemo(() => {
return {
const assetsByAssetKey = useMemo(
() => Object.fromEntries(assets?.map((asset) => [tokenForAssetKey(asset.key), asset]) ?? []),
[assets],
);

return useMemo(
() => ({
assets,
assetsByAssetKey,
error,
loading: !assets && !error,
query,
};
}, [assets, error, query]);
}),
[assets, assetsByAssetKey, error, query],
);
}

interface AssetCatalogTableProps {
Expand Down Expand Up @@ -253,7 +268,7 @@ export const AssetsCatalogTable = ({
[filtered, prefixPath, view],
);

const refreshState = useRefreshAtInterval({
const refreshState = useRefreshAtInterval<any>({
refresh: query,
intervalMs: 4 * FIFTEEN_SECONDS,
leading: true,
Expand Down
3 changes: 0 additions & 3 deletions js_modules/dagster-ui/packages/ui-core/src/assets/__tests__/useAllAssets.test.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -112,9 +112,6 @@ describe('useAllAssets', () => {
},
});

await waitFor(() => {
expect(result.current.assets?.length).toBe(1);
});
await waitFor(() => {
expect(result.current.assets?.length).toBe(5);
});
Expand Down
23 changes: 15 additions & 8 deletions js_modules/dagster-ui/packages/ui-core/src/search/useIndexedDBCachedQuery.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
import memoize from 'lodash/memoize';
import React, {useCallback, useEffect, useMemo} from 'react';

import {
Expand All @@ -12,6 +11,7 @@ import {usePreviousDistinctValue} from '../hooks/usePrevious';
import {useUpdatingRef} from '../hooks/useUpdatingRef';
import {CompletionType, useBlockTraceUntilTrue} from '../performance/TraceContext';
import {cache} from '../util/idb-lru-cache';
import {weakMapMemoize} from '../util/weakMapMemoize';

type CacheData<TQuery> = {
data: TQuery;
Expand All @@ -20,7 +20,7 @@ type CacheData<TQuery> = {

export const KEY_PREFIX = 'indexdbQueryCache:';

export class CacheManager<TQuery> {
class CacheManager<TQuery> {
private cache: ReturnType<typeof cache<CacheData<TQuery>>> | undefined;
private key: string;
private current?: CacheData<TQuery>;
Expand Down Expand Up @@ -63,8 +63,7 @@ export class CacheManager<TQuery> {
async set(data: TQuery, version: number | string): Promise<void> {
if (
this.current?.data === data ||
(JSON.stringify(this.current?.data) === JSON.stringify(data) &&
this.current?.version === version)
(stringified(this.current?.data) === stringified(data) && this.current?.version === version)
) {
return;
}
Expand Down Expand Up @@ -249,7 +248,7 @@ export function useIndexedDBCachedQuery<TQuery, TVariables extends OperationVari
data &&
// Work around a weird jest issue where it returns an empty object if no mocks are found...
Object.keys(data).length &&
(!dataRef.current || JSON.stringify(dataRef.current) !== JSON.stringify(data))
(!dataRef.current || stringified(dataRef.current) !== stringified(data))
) {
setData(data);
}
Expand Down Expand Up @@ -348,7 +347,7 @@ export function useGetData() {

export function useGetCachedData() {
return useCallback(async <TQuery,>({key, version}: {key: string; version: number | string}) => {
const cacheManager = new CacheManager<TQuery>(key);
const cacheManager = getCacheManager<TQuery>(key);
return await cacheManager.get(version);
}, []);
}
Expand All @@ -359,13 +358,21 @@ export function useClearCachedData() {
}, []);
}

let getCacheManager = memoize(<TQuery,>(key: string) => {
export let getCacheManager = weakMapMemoize(<TQuery,>(key: string) => {
return new CacheManager<TQuery>(key);
});

export const __resetForJest = () => {
Object.keys(globalFetchStates).forEach((key) => delete globalFetchStates[key]);
getCacheManager = memoize(<TQuery,>(key: string) => {
getCacheManager = weakMapMemoize(<TQuery,>(key: string) => {
return new CacheManager<TQuery>(key);
});
};

const stringified = weakMapMemoize((data: any) => {
try {
return JSON.stringify(data);
} catch {
return '';
}
});