Move subscription tracking logic over to batchActions file

Author: markerikson

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

@@ -1,7 +1,12 @@
 import type { QueryThunk, RejectedAction } from '../buildThunks'
 import type { InternalHandlerBuilder } from './types'
-import type { SubscriptionState } from '../apiState'
+import type {
+  SubscriptionState,
+  QuerySubstateIdentifier,
+  Subscribers,
+} from '../apiState'
 import { produceWithPatches } from 'immer'
+import { createSlice, PayloadAction, AnyAction } from '@reduxjs/toolkit'
 
 // Copied from https://github.com/feross/queue-microtask
 let promise: Promise<any>
@@ -19,14 +24,69 @@ const queueMicrotaskShim =
 export const buildBatchedActionsHandler: InternalHandlerBuilder<
   [actionShouldContinue: boolean, subscriptionExists: boolean]
 > = ({ api, queryThunk, internalState }) => {
-  const { actuallyMutateSubscriptions } = api.internalActions
   const subscriptionsPrefix = `${api.reducerPath}/subscriptions`
 
   let previousSubscriptions: SubscriptionState =
     null as unknown as SubscriptionState
 
   let dispatchQueued = false
 
+  const { updateSubscriptionOptions, unsubscribeQueryResult } =
+    api.internalActions
+
+  // Actually intentionally mutate the subscriptions state used in the middleware
+  // This is done to speed up perf when loading many components
+  const actuallyMutateSubscriptions = (
+    mutableState: SubscriptionState,
+    action: AnyAction
+  ) => {
+    if (updateSubscriptionOptions.match(action)) {
+      const { queryCacheKey, requestId, options } = action.payload
+
+      if (mutableState?.[queryCacheKey]?.[requestId]) {
+        mutableState[queryCacheKey]![requestId] = options
+      }
+      return true
+    }
+    if (unsubscribeQueryResult.match(action)) {
+      const { queryCacheKey, requestId } = action.payload
+      if (mutableState[queryCacheKey]) {
+        delete mutableState[queryCacheKey]![requestId]
+      }
+      return true
+    }
+    if (api.internalActions.removeQueryResult.match(action)) {
+      delete mutableState[action.payload.queryCacheKey]
+      return true
+    }
+    if (queryThunk.pending.match(action)) {
+      const {
+        meta: { arg, requestId },
+      } = action
+      if (arg.subscribe) {
+        const substate = (mutableState[arg.queryCacheKey] ??= {})
+        substate[requestId] =
+          arg.subscriptionOptions ?? substate[requestId] ?? {}
+
+        return true
+      }
+    }
+    if (queryThunk.rejected.match(action)) {
+      const {
+        meta: { condition, arg, requestId },
+      } = action
+      if (condition && arg.subscribe) {
+        const substate = (mutableState[arg.queryCacheKey] ??= {})
+        substate[requestId] =
+          arg.subscriptionOptions ?? substate[requestId] ?? {}
+
+        return true
+      }
+    }
+
+    return false
+  }
+
   return (action, mwApi) => {
     if (!previousSubscriptions) {
       // Initialize it the first time this handler runs

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

@@ -412,62 +412,6 @@ export function buildSlice({
     },
   })
 
-  const { updateSubscriptionOptions, unsubscribeQueryResult } =
-    subscriptionSlice.actions
-
-  // Actually intentionally mutate the subscriptions state used in the middleware
-  // This is done to speed up perf when loading many components
-  const actuallyMutateSubscriptions = (
-    draft: SubscriptionState,
-    action: AnyAction
-  ) => {
-    if (updateSubscriptionOptions.match(action)) {
-      const { queryCacheKey, requestId, options } = action.payload
-
-      if (draft?.[queryCacheKey]?.[requestId]) {
-        draft[queryCacheKey]![requestId] = options
-      }
-      return true
-    }
-    if (unsubscribeQueryResult.match(action)) {
-      const { queryCacheKey, requestId } = action.payload
-      if (draft[queryCacheKey]) {
-        delete draft[queryCacheKey]![requestId]
-      }
-      return true
-    }
-    if (querySlice.actions.removeQueryResult.match(action)) {
-      delete draft[action.payload.queryCacheKey]
-      return true
-    }
-    if (queryThunk.pending.match(action)) {
-      const {
-        meta: { arg, requestId },
-      } = action
-      if (arg.subscribe) {
-        const substate = (draft[arg.queryCacheKey] ??= {})
-        substate[requestId] =
-          arg.subscriptionOptions ?? substate[requestId] ?? {}
-
-        return true
-      }
-    }
-    if (queryThunk.rejected.match(action)) {
-      const {
-        meta: { condition, arg, requestId },
-      } = action
-      if (condition && arg.subscribe) {
-        const substate = (draft[arg.queryCacheKey] ??= {})
-        substate[requestId] =
-          arg.subscriptionOptions ?? substate[requestId] ?? {}
-
-        return true
-      }
-    }
-
-    return false
-  }
-
   const internalSubscriptionsSlice = createSlice({
     name: `${reducerPath}/internalSubscriptions`,
     initialState: initialState as SubscriptionState,
@@ -536,7 +480,6 @@ export function buildSlice({
     /** @deprecated has been renamed to `removeMutationResult` */
     unsubscribeMutationResult: mutationSlice.actions.removeMutationResult,
     resetApiState,
-    actuallyMutateSubscriptions,
   }
 
   return { reducer, actions }