SSR & rehydration support, suspense foundations (#1277)

Co-authored-by: Lenz Weber <mail@lenzw.de>
Co-authored-by: Josh Fraser <joshfraser91@gmail.com>
Co-authored-by: Matt Sutkowski <msutkowski@gmail.com>

packages/toolkit/src/query/apiTypes.ts

@@ -4,10 +4,16 @@ import type {
   EndpointDefinition,
   ReplaceTagTypes,
 } from './endpointDefinitions'
-import type { UnionToIntersection, NoInfer } from './tsHelpers'
+import type {
+  UnionToIntersection,
+  NoInfer,
+  WithRequiredProp,
+} from './tsHelpers'
 import type { CoreModule } from './core/module'
 import type { CreateApiOptions } from './createApi'
 import type { BaseQueryFn } from './baseQueryTypes'
+import type { CombinedState } from './core/apiState'
+import type { AnyAction } from '@reduxjs/toolkit'
 
 export interface ApiModules<
   // eslint-disable-next-line @typescript-eslint/no-unused-vars
@@ -31,8 +37,15 @@ export type Module<Name extends ModuleName> = {
     TagTypes extends string
   >(
     api: Api<BaseQuery, EndpointDefinitions, ReducerPath, TagTypes, ModuleName>,
-    options: Required<
-      CreateApiOptions<BaseQuery, Definitions, ReducerPath, TagTypes>
+    options: WithRequiredProp<
+      CreateApiOptions<BaseQuery, Definitions, ReducerPath, TagTypes>,
+      | 'reducerPath'
+      | 'serializeQueryArgs'
+      | 'keepUnusedDataFor'
+      | 'refetchOnMountOrArgChange'
+      | 'refetchOnFocus'
+      | 'refetchOnReconnect'
+      | 'tagTypes'
     >,
     context: ApiContext<Definitions>
   ): {
@@ -47,6 +60,10 @@ export interface ApiContext<Definitions extends EndpointDefinitions> {
   apiUid: string
   endpointDefinitions: Definitions
   batch(cb: () => void): void
+  extractRehydrationInfo: (
+    action: AnyAction
+  ) => CombinedState<any, any, any> | undefined
+  hasRehydrationInfo: (action: AnyAction) => boolean
 }
 
 export type Api<

packages/toolkit/src/query/core/buildInitiate.ts

@@ -5,22 +5,12 @@ import type {
   QueryArgFrom,
   ResultTypeFrom,
 } from '../endpointDefinitions'
-import type {
-  QueryThunkArg,
-  MutationThunkArg,
-  QueryThunk,
-  MutationThunk,
-} from './buildThunks'
-import type {
-  AnyAction,
-  AsyncThunk,
-  ThunkAction,
-  SerializedError,
-} from '@reduxjs/toolkit'
-import { unwrapResult } from '@reduxjs/toolkit'
+import { DefinitionType } from '../endpointDefinitions'
+import type { QueryThunk, MutationThunk } from './buildThunks'
+import type { AnyAction, ThunkAction, SerializedError } from '@reduxjs/toolkit'
 import type { QuerySubState, SubscriptionOptions, RootState } from './apiState'
 import type { InternalSerializeQueryArgs } from '../defaultSerializeQueryArgs'
-import type { Api } from '../apiTypes'
+import type { Api, ApiContext } from '../apiTypes'
 import type { ApiEndpointQuery } from './module'
 import type { BaseQueryError } from '../baseQueryTypes'
 
@@ -120,6 +110,7 @@ export type MutationActionCreatorResult<
    * A unique string generated for the request sequence
    */
   requestId: string
+
   /**
    * A method to cancel the mutation promise. Note that this is not intended to prevent the mutation
    * that was fired off from reaching the server, but only to assist in handling the response.
@@ -189,18 +180,58 @@ export function buildInitiate({
   queryThunk,
   mutationThunk,
   api,
+  context,
 }: {
   serializeQueryArgs: InternalSerializeQueryArgs
   queryThunk: QueryThunk
   mutationThunk: MutationThunk
   api: Api<any, EndpointDefinitions, any, any>
+  context: ApiContext<EndpointDefinitions>
 }) {
+  const runningQueries: Record<
+    string,
+    QueryActionCreatorResult<any> | undefined
+  > = {}
+  const runningMutations: Record<
+    string,
+    MutationActionCreatorResult<any> | undefined
+  > = {}
+
   const {
     unsubscribeQueryResult,
     removeMutationResult,
     updateSubscriptionOptions,
   } = api.internalActions
-  return { buildInitiateQuery, buildInitiateMutation }
+  return {
+    buildInitiateQuery,
+    buildInitiateMutation,
+    getRunningOperationPromises,
+    getRunningOperationPromise,
+  }
+
+  function getRunningOperationPromise(
+    endpointName: string,
+    argOrRequestId: any
+  ): any {
+    const endpointDefinition = context.endpointDefinitions[endpointName]
+    if (endpointDefinition.type === DefinitionType.query) {
+      const queryCacheKey = serializeQueryArgs({
+        queryArgs: argOrRequestId,
+        endpointDefinition,
+        endpointName,
+      })
+      return runningQueries[queryCacheKey]
+    } else {
+      return runningMutations[argOrRequestId]
+    }
+  }
+
+  function getRunningOperationPromises() {
+    return [
+      ...Object.values(runningQueries),
+      ...Object.values(runningMutations),
+    ].filter(<T>(t: T | undefined): t is T => !!t)
+  }
 
   function middlewareWarning(getState: () => RootState<{}, string, string>) {
     if (process.env.NODE_ENV !== 'production') {
@@ -242,8 +273,8 @@ Features like automatic cache collection, automatic refetching etc. will not be
         const thunkResult = dispatch(thunk)
         middlewareWarning(getState)
         const { requestId, abort } = thunkResult
-        const statePromise = Object.assign(
-          thunkResult.then(() =>
+        const statePromise: QueryActionCreatorResult<any> = Object.assign(
+          Promise.all([runningQueries[queryCacheKey], thunkResult]).then(() =>
             (api.endpoints[endpointName] as ApiEndpointQuery<any, any>).select(
               arg
             )(getState())
@@ -280,14 +311,21 @@ Features like automatic cache collection, automatic refetching etc. will not be
             },
           }
         )
+
+        if (!runningQueries[queryCacheKey]) {
+          runningQueries[queryCacheKey] = statePromise
+          statePromise.then(() => {
+            delete runningQueries[queryCacheKey]
+          })
+        }
+
         return statePromise
       }
     return queryAction
   }
 
   function buildInitiateMutation(
-    endpointName: string,
-    definition: MutationDefinition<any, any, any, any>
+    endpointName: string
   ): StartMutationActionCreator<any> {
     return (arg, { track = true, fixedCacheKey } = {}) =>
       (dispatch, getState) => {
@@ -309,14 +347,28 @@ Features like automatic cache collection, automatic refetching etc. will not be
           dispatch(removeMutationResult({ requestId, fixedCacheKey }))
         }
 
-        return Object.assign(returnValuePromise, {
+        const ret = Object.assign(returnValuePromise, {
           arg: thunkResult.arg,
           requestId,
           abort,
           unwrap: thunkResult.unwrap,
           unsubscribe: reset,
           reset,
         })
+
+        runningMutations[requestId] = ret
+        ret.then(() => {
+          delete runningMutations[requestId]
+        })
+        if (fixedCacheKey) {
+          runningMutations[fixedCacheKey] = ret
+          ret.then(() => {
+            if (runningMutations[fixedCacheKey] === ret)
+              delete runningMutations[fixedCacheKey]
+          })
+        }
+
+        return ret
       }
   }
 }

packages/toolkit/src/query/core/buildMiddleware/cacheCollection.ts

@@ -1,6 +1,6 @@
 import type { BaseQueryFn } from '../../baseQueryTypes'
 import type { QueryDefinition } from '../../endpointDefinitions'
-import type { QueryCacheKey } from '../apiState'
+import type { ConfigState, QueryCacheKey } from '../apiState'
 import { QuerySubstateIdentifier } from '../apiState'
 import type {
   QueryStateMeta,
@@ -42,15 +42,11 @@ export const build: SubMiddlewareBuilder = ({ reducerPath, api, context }) => {
           const state = mwApi.getState()[reducerPath]
           const { queryCacheKey } = action.payload
 
-          const endpointDefinition = context.endpointDefinitions[
-            state.queries[queryCacheKey]?.endpointName!
-          ] as QueryDefinition<any, any, any, any>
-
           handleUnsubscribe(
             queryCacheKey,
+            state.queries[queryCacheKey]?.endpointName,
             mwApi,
-            endpointDefinition?.keepUnusedDataFor ??
-              state.config.keepUnusedDataFor
+            state.config
           )
         }
 
@@ -61,14 +57,37 @@ export const build: SubMiddlewareBuilder = ({ reducerPath, api, context }) => {
           }
         }
 
+        if (context.hasRehydrationInfo(action)) {
+          const state = mwApi.getState()[reducerPath]
+          const { queries } = context.extractRehydrationInfo(action)!
+          for (const [queryCacheKey, queryState] of Object.entries(queries)) {
+            // Gotcha:
+            // If rehydrating before the endpoint has been injected,the global `keepUnusedDataFor`
+            // will be used instead of the endpoint-specific one.
+            handleUnsubscribe(
+              queryCacheKey as QueryCacheKey,
+              queryState?.endpointName,
+              mwApi,
+              state.config
+            )
+          }
+        }
+
         return result
       }
 
     function handleUnsubscribe(
       queryCacheKey: QueryCacheKey,
+      endpointName: string | undefined,
       api: SubMiddlewareApi,
-      keepUnusedDataFor: number
+      config: ConfigState<string>
     ) {
+      const endpointDefinition = context.endpointDefinitions[
+        endpointName!
+      ] as QueryDefinition<any, any, any, any>
+      const keepUnusedDataFor =
+        endpointDefinition?.keepUnusedDataFor ?? config.keepUnusedDataFor
+
       const currentTimeout = currentRemovalTimeouts[queryCacheKey]
       if (currentTimeout) {
         clearTimeout(currentTimeout)