diff --git a/packages/runtime/src/Suspense.tsx b/packages/runtime/src/Suspense.tsx index e9ca9e727..004b37fb8 100644 --- a/packages/runtime/src/Suspense.tsx +++ b/packages/runtime/src/Suspense.tsx @@ -47,7 +47,29 @@ const getHydrateData = (id: string) => { }; }; -const querySet = new Set(); +class QueryClient { + queryCache: Map>; + + constructor() { + this.queryCache = new Map(); + } + + setQueryData(key: string, promise: Promise) { + if (key) { + this.queryCache.set(key, promise); + } + } + + getQueryData(key: string) { + return this.queryCache.get(key); + } + + removeQueryData(key: string) { + this.queryCache.delete(key); + } +} + +const queryClient = new QueryClient(); interface SuspenseDataProps { queryKey?: string; @@ -61,28 +83,26 @@ export function useSuspenseData(request?: Request, options?: SuspenseDataProps) const { data, done, promise, update, error, id } = suspenseState; const { hasHydrateData, data: hydrateData } = getHydrateData(id); - const queryInProcess = options?.queryKey && querySet.has(options.queryKey); + const queryInProcess = options?.queryKey && queryClient.getQueryData(options.queryKey); let thenable: Promise = null; if (!hasHydrateData && !error && !done && !promise && request && !queryInProcess) { thenable = request(requestContext); thenable.then((response) => { - querySet.delete(options?.queryKey); + queryClient.removeQueryData(options?.queryKey); update({ done: true, data: response, promise: null, }); }).catch(e => { - querySet.delete(options?.queryKey); + queryClient.removeQueryData(options?.queryKey); update({ done: true, error: e, promise: null, }); }); - if (options?.queryKey) { - querySet.add(options.queryKey); - } + queryClient.setQueryData(options?.queryKey, thenable); } React.useEffect(() => { @@ -133,7 +153,7 @@ export function useSuspenseData(request?: Request, options?: SuspenseDataProps) throw thenable; } else { // Throw a pending promise to Suspense, otherwise the component will be error because of undefined data. - throw new Promise(() => {}); + throw (queryClient.getQueryData(options?.queryKey) || new Promise(() => {})); } }