Consolidate RTKQ middleware to simplify stack size

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

@@ -1,5 +1,5 @@
 import type { QueryThunk, RejectedAction } from '../buildThunks'
-import type { SubMiddlewareBuilder } from './types'
+import type { InternalHandlerBuilder } from './types'
 
 // Copied from https://github.com/feross/queue-microtask
 let promise: Promise<any>
@@ -14,44 +14,38 @@ const queueMicrotaskShim =
           }, 0)
         )
 
-export const build: SubMiddlewareBuilder = ({
+export const buildBatchedActionsHandler: InternalHandlerBuilder<boolean> = ({
   api,
-  context: { apiUid },
   queryThunk,
-  reducerPath,
 }) => {
-  return (mwApi) => {
-    let abortedQueryActionsQueue: RejectedAction<QueryThunk, any>[] = []
-    let dispatchQueued = false
+  let abortedQueryActionsQueue: RejectedAction<QueryThunk, any>[] = []
+  let dispatchQueued = false
 
-    return (next) => (action) => {
-      if (queryThunk.rejected.match(action)) {
-        const { condition, arg } = action.meta
+  return (action, mwApi) => {
+    if (queryThunk.rejected.match(action)) {
+      const { condition, arg } = action.meta
 
-        if (condition && arg.subscribe) {
-          // request was aborted due to condition (another query already running)
-          // _Don't_ dispatch right away - queue it for a debounced grouped dispatch
-          abortedQueryActionsQueue.push(action)
+      if (condition && arg.subscribe) {
+        // request was aborted due to condition (another query already running)
+        // _Don't_ dispatch right away - queue it for a debounced grouped dispatch
+        abortedQueryActionsQueue.push(action)
 
-          if (!dispatchQueued) {
-            queueMicrotaskShim(() => {
-              mwApi.dispatch(
-                api.internalActions.subscriptionRequestsRejected(
-                  abortedQueryActionsQueue
-                )
+        if (!dispatchQueued) {
+          queueMicrotaskShim(() => {
+            mwApi.dispatch(
+              api.internalActions.subscriptionRequestsRejected(
+                abortedQueryActionsQueue
               )
-              abortedQueryActionsQueue = []
-            })
-            dispatchQueued = true
-          }
-          // _Don't_ let the action reach the reducers now!
-          return
+            )
+            abortedQueryActionsQueue = []
+          })
+          dispatchQueued = true
         }
+        // _Don't_ let the action reach the reducers now!
+        return false
       }
-
-      const result = next(action)
-
-      return result
     }
+
+    return true
   }
 }

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

@@ -1,12 +1,12 @@
 import type { BaseQueryFn } from '../../baseQueryTypes'
 import type { QueryDefinition } from '../../endpointDefinitions'
 import type { ConfigState, QueryCacheKey } from '../apiState'
-import { QuerySubstateIdentifier } from '../apiState'
 import type {
   QueryStateMeta,
   SubMiddlewareApi,
-  SubMiddlewareBuilder,
   TimeoutId,
+  InternalHandlerBuilder,
+  ApiMiddlewareInternalHandler,
 } from './types'
 
 export type ReferenceCacheCollection = never
@@ -45,7 +45,11 @@ declare module '../../endpointDefinitions' {
 export const THIRTY_TWO_BIT_MAX_INT = 2_147_483_647
 export const THIRTY_TWO_BIT_MAX_TIMER_SECONDS = 2_147_483_647 / 1_000 - 1
 
-export const build: SubMiddlewareBuilder = ({ reducerPath, api, context }) => {
+export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
+  reducerPath,
+  api,
+  context,
+}) => {
   const { removeQueryResult, unsubscribeQueryResult } = api.internalActions
 
   function anySubscriptionsRemainingForKey(
@@ -57,88 +61,83 @@ export const build: SubMiddlewareBuilder = ({ reducerPath, api, context }) => {
     return !!subscriptions && !isObjectEmpty(subscriptions)
   }
 
-  return (mwApi) => {
-    const currentRemovalTimeouts: QueryStateMeta<TimeoutId> = {}
+  const currentRemovalTimeouts: QueryStateMeta<TimeoutId> = {}
 
-    return (next) =>
-      (action): any => {
-        const result = next(action)
+  const handler: ApiMiddlewareInternalHandler = (action, mwApi) => {
+    if (unsubscribeQueryResult.match(action)) {
+      const state = mwApi.getState()[reducerPath]
+      const { queryCacheKey } = action.payload
 
-        if (unsubscribeQueryResult.match(action)) {
-          const state = mwApi.getState()[reducerPath]
-          const { queryCacheKey } = action.payload
-
-          handleUnsubscribe(
-            queryCacheKey,
-            state.queries[queryCacheKey]?.endpointName,
-            mwApi,
-            state.config
-          )
-        }
-
-        if (api.util.resetApiState.match(action)) {
-          for (const [key, timeout] of Object.entries(currentRemovalTimeouts)) {
-            if (timeout) clearTimeout(timeout)
-            delete currentRemovalTimeouts[key]
-          }
-        }
-
-        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
-            )
-          }
-        }
+      handleUnsubscribe(
+        queryCacheKey,
+        state.queries[queryCacheKey]?.endpointName,
+        mwApi,
+        state.config
+      )
+    }
 
-        return result
+    if (api.util.resetApiState.match(action)) {
+      for (const [key, timeout] of Object.entries(currentRemovalTimeouts)) {
+        if (timeout) clearTimeout(timeout)
+        delete currentRemovalTimeouts[key]
       }
+    }
 
-    function handleUnsubscribe(
-      queryCacheKey: QueryCacheKey,
-      endpointName: string | undefined,
-      api: SubMiddlewareApi,
-      config: ConfigState<string>
-    ) {
-      const endpointDefinition = context.endpointDefinitions[
-        endpointName!
-      ] as QueryDefinition<any, any, any, any>
-      const keepUnusedDataFor =
-        endpointDefinition?.keepUnusedDataFor ?? config.keepUnusedDataFor
-
-      if (keepUnusedDataFor === Infinity) {
-        // Hey, user said keep this forever!
-        return
+    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
+        )
       }
-      // Prevent `setTimeout` timers from overflowing a 32-bit internal int, by
-      // clamping the max value to be at most 1000ms less than the 32-bit max.
-      // Look, a 24.8-day keepalive ought to be enough for anybody, right? :)
-      // Also avoid negative values too.
-      const finalKeepUnusedDataFor = Math.max(
-        0,
-        Math.min(keepUnusedDataFor, THIRTY_TWO_BIT_MAX_TIMER_SECONDS)
-      )
+    }
+  }
 
-      if (!anySubscriptionsRemainingForKey(queryCacheKey, api)) {
-        const currentTimeout = currentRemovalTimeouts[queryCacheKey]
-        if (currentTimeout) {
-          clearTimeout(currentTimeout)
-        }
-        currentRemovalTimeouts[queryCacheKey] = setTimeout(() => {
-          if (!anySubscriptionsRemainingForKey(queryCacheKey, api)) {
-            api.dispatch(removeQueryResult({ queryCacheKey }))
-          }
-          delete currentRemovalTimeouts![queryCacheKey]
-        }, finalKeepUnusedDataFor * 1000)
+  function handleUnsubscribe(
+    queryCacheKey: QueryCacheKey,
+    endpointName: string | undefined,
+    api: SubMiddlewareApi,
+    config: ConfigState<string>
+  ) {
+    const endpointDefinition = context.endpointDefinitions[
+      endpointName!
+    ] as QueryDefinition<any, any, any, any>
+    const keepUnusedDataFor =
+      endpointDefinition?.keepUnusedDataFor ?? config.keepUnusedDataFor
+
+    if (keepUnusedDataFor === Infinity) {
+      // Hey, user said keep this forever!
+      return
+    }
+    // Prevent `setTimeout` timers from overflowing a 32-bit internal int, by
+    // clamping the max value to be at most 1000ms less than the 32-bit max.
+    // Look, a 24.8-day keepalive ought to be enough for anybody, right? :)
+    // Also avoid negative values too.
+    const finalKeepUnusedDataFor = Math.max(
+      0,
+      Math.min(keepUnusedDataFor, THIRTY_TWO_BIT_MAX_TIMER_SECONDS)
+    )
+
+    if (!anySubscriptionsRemainingForKey(queryCacheKey, api)) {
+      const currentTimeout = currentRemovalTimeouts[queryCacheKey]
+      if (currentTimeout) {
+        clearTimeout(currentTimeout)
       }
+      currentRemovalTimeouts[queryCacheKey] = setTimeout(() => {
+        if (!anySubscriptionsRemainingForKey(queryCacheKey, api)) {
+          api.dispatch(removeQueryResult({ queryCacheKey }))
+        }
+        delete currentRemovalTimeouts![queryCacheKey]
+      }, finalKeepUnusedDataFor * 1000)
     }
   }
+
+  return handler
 }