scope getRunningOperationPromise to store instance

Author: phryneas

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

@@ -14,6 +14,7 @@ import type { Api, ApiContext } from '../apiTypes'
 import type { ApiEndpointQuery } from './module'
 import type { BaseQueryError, QueryReturnValue } from '../baseQueryTypes'
 import type { QueryResultSelectorResult } from './buildSelectors'
+import { Dispatch } from 'redux'
 
 declare module './module' {
   export interface ApiEndpointQuery<
@@ -196,14 +197,14 @@ export function buildInitiate({
   api: Api<any, EndpointDefinitions, any, any>
   context: ApiContext<EndpointDefinitions>
 }) {
-  const runningQueries: Record<
-    string,
-    QueryActionCreatorResult<any> | undefined
-  > = {}
-  const runningMutations: Record<
-    string,
-    MutationActionCreatorResult<any> | undefined
-  > = {}
+  const runningQueries: Map<
+    Dispatch,
+    Record<string, QueryActionCreatorResult<any> | undefined>
+  > = new Map()
+  const runningMutations: Map<
+    Dispatch,
+    Record<string, MutationActionCreatorResult<any> | undefined>
+  > = new Map()
 
   const {
     unsubscribeQueryResult,
@@ -220,25 +221,33 @@ export function buildInitiate({
   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]
+  ): ThunkAction<any, any, any, AnyAction> {
+    return (dispatch: Dispatch) => {
+      const endpointDefinition = context.endpointDefinitions[endpointName]
+      if (endpointDefinition.type === DefinitionType.query) {
+        const queryCacheKey = serializeQueryArgs({
+          queryArgs: argOrRequestId,
+          endpointDefinition,
+          endpointName,
+        })
+        return runningQueries.get(dispatch)?.[queryCacheKey]
+      } else {
+        return runningMutations.get(dispatch)?.[argOrRequestId]
+      }
     }
   }
 
-  function getRunningOperationPromises() {
-    return [
-      ...Object.values(runningQueries),
-      ...Object.values(runningMutations),
-    ].filter(<T>(t: T | undefined): t is T => !!t)
+  function getRunningOperationPromises(): ThunkAction<
+    Promise<unknown>[],
+    any,
+    any,
+    AnyAction
+  > {
+    return (dispatch: Dispatch) =>
+      [
+        ...Object.values(runningQueries.get(dispatch) || {}),
+        ...Object.values(runningMutations.get(dispatch) || {}),
+      ].filter(<T>(t: T | undefined): t is T => !!t)
   }
 
   function middlewareWarning(getState: () => RootState<{}, string, string>) {
@@ -302,7 +311,7 @@ Features like automatic cache collection, automatic refetching etc. will not be
 
         const skippedSynchronously = stateAfter.requestId !== requestId
 
-        const runningQuery = runningQueries[queryCacheKey]
+        const runningQuery = runningQueries.get(dispatch)?.[queryCacheKey]
         const selectFromState = () => selector(getState())
 
         const statePromise: QueryActionCreatorResult<any> = Object.assign(
@@ -360,9 +369,15 @@ Features like automatic cache collection, automatic refetching etc. will not be
         )
 
         if (!runningQuery && !skippedSynchronously && !forceQueryFn) {
-          runningQueries[queryCacheKey] = statePromise
+          const running = runningQueries.get(dispatch) || {}
+          running[queryCacheKey] = statePromise
+          runningQueries.set(dispatch, running)
+
           statePromise.then(() => {
-            delete runningQueries[queryCacheKey]
+            delete running[queryCacheKey]
+            if (!Object.keys(running).length) {
+              runningQueries.delete(dispatch)
+            }
           })
         }
 
@@ -404,15 +419,24 @@ Features like automatic cache collection, automatic refetching etc. will not be
           reset,
         })
 
-        runningMutations[requestId] = ret
+        const running = runningMutations.get(dispatch) || {}
+        runningMutations.set(dispatch, running)
+        running[requestId] = ret
         ret.then(() => {
-          delete runningMutations[requestId]
+          delete running[requestId]
+          if (!Object.keys(running).length) {
+            runningMutations.delete(dispatch)
+          }
         })
         if (fixedCacheKey) {
-          runningMutations[fixedCacheKey] = ret
+          running[fixedCacheKey] = ret
           ret.then(() => {
-            if (runningMutations[fixedCacheKey] === ret)
-              delete runningMutations[fixedCacheKey]
+            if (running[fixedCacheKey] === ret) {
+              delete running[fixedCacheKey]
+              if (!Object.keys(running).length) {
+                runningMutations.delete(dispatch)
+              }
+            }
           })
         }
 

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

@@ -142,7 +142,12 @@ declare module '../apiTypes' {
          * Useful for SSR scenarios to await everything triggered in any way,
          * including via hook calls, or manually dispatching `initiate` actions.
          */
-        getRunningOperationPromises: () => Array<Promise<unknown>>
+        getRunningOperationPromises: () => ThunkAction<
+          Array<Promise<unknown>>,
+          any,
+          any,
+          AnyAction
+        >
         /**
          * If a promise is running for a given endpoint name + argument combination,
          * returns that promise. Otherwise, returns `undefined`.
@@ -152,21 +157,29 @@ declare module '../apiTypes' {
         getRunningOperationPromise<EndpointName extends QueryKeys<Definitions>>(
           endpointName: EndpointName,
           args: QueryArgFrom<Definitions[EndpointName]>
-        ):
+        ): ThunkAction<
           | QueryActionCreatorResult<
               Definitions[EndpointName] & { type: 'query' }
             >
-          | undefined
+          | undefined,
+          any,
+          any,
+          AnyAction
+        >
         getRunningOperationPromise<
           EndpointName extends MutationKeys<Definitions>
         >(
           endpointName: EndpointName,
           fixedCacheKeyOrRequestId: string
-        ):
+        ): ThunkAction<
           | MutationActionCreatorResult<
               Definitions[EndpointName] & { type: 'mutation' }
             >
-          | undefined
+          | undefined,
+          any,
+          any,
+          AnyAction
+        >
 
         /**
          * A Redux thunk that can be used to manually trigger pre-fetching of data.