Rewrite subscription handling and polling calculations for better perf (#5064)

* Add test for subscription serialization
* Rewrite internal subscription handle for better perf
* Hack in a way to expose polling update counters in tests
* Rewrite polling handling for perf
* Use lazy getOrInsert defaults

---------

Co-authored-by: codesandbox-bot <codesandbot@codesandbox.io>

Author: markerikson

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

@@ -148,6 +148,7 @@ export type SubscriptionOptions = {
    */
   refetchOnFocus?: boolean
 }
+export type SubscribersInternal = Map<string, SubscriptionOptions>
 export type Subscribers = { [requestId: string]: SubscriptionOptions }
 export type QueryKeys<Definitions extends EndpointDefinitions> = {
   [K in keyof Definitions]: Definitions[K] extends QueryDefinition<
@@ -327,6 +328,8 @@ export type QueryState<D extends EndpointDefinitions> = {
     | undefined
 }
 
+export type SubscriptionInternalState = Map<string, SubscribersInternal>
+
 export type SubscriptionState = {
   [queryCacheKey: string]: Subscribers | undefined
 }

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

@@ -42,6 +42,7 @@ import type {
   ThunkApiMetaConfig,
 } from './buildThunks'
 import type { ApiEndpointQuery } from './module'
+import type { InternalMiddlewareState } from './buildMiddleware/types'
 
 export type BuildInitiateApiEndpointQuery<
   Definition extends QueryDefinition<any, any, any, any, any>,
@@ -270,27 +271,17 @@ export function buildInitiate({
   mutationThunk,
   api,
   context,
+  internalState,
 }: {
   serializeQueryArgs: InternalSerializeQueryArgs
   queryThunk: QueryThunk
   infiniteQueryThunk: InfiniteQueryThunk<any>
   mutationThunk: MutationThunk
   api: Api<any, EndpointDefinitions, any, any>
   context: ApiContext<EndpointDefinitions>
+  internalState: InternalMiddlewareState
 }) {
-  const runningQueries: Map<
-    Dispatch,
-    Record<
-      string,
-      | QueryActionCreatorResult<any>
-      | InfiniteQueryActionCreatorResult<any>
-      | undefined
-    >
-  > = new Map()
-  const runningMutations: Map<
-    Dispatch,
-    Record<string, MutationActionCreatorResult<any> | undefined>
-  > = new Map()
+  const { runningQueries, runningMutations } = internalState
 
   const {
     unsubscribeQueryResult,

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

@@ -1,12 +1,12 @@
 import type { InternalHandlerBuilder, SubscriptionSelectors } from './types'
-import type { SubscriptionState } from '../apiState'
+import type { SubscriptionInternalState, SubscriptionState } from '../apiState'
 import { produceWithPatches } from 'immer'
 import type { Action } from '@reduxjs/toolkit'
-import { countObjectKeys } from '../../utils/countObjectKeys'
+import { getOrInsertComputed, createNewMap } from '../../utils/getOrInsert'
 
 export const buildBatchedActionsHandler: InternalHandlerBuilder<
   [actionShouldContinue: boolean, returnValue: SubscriptionSelectors | boolean]
-> = ({ api, queryThunk, internalState }) => {
+> = ({ api, queryThunk, internalState, mwApi }) => {
   const subscriptionsPrefix = `${api.reducerPath}/subscriptions`
 
   let previousSubscriptions: SubscriptionState =
@@ -20,58 +20,63 @@ export const buildBatchedActionsHandler: InternalHandlerBuilder<
   // 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,
+    currentSubscriptions: SubscriptionInternalState,
     action: Action,
   ) => {
     if (updateSubscriptionOptions.match(action)) {
       const { queryCacheKey, requestId, options } = action.payload
 
-      if (mutableState?.[queryCacheKey]?.[requestId]) {
-        mutableState[queryCacheKey]![requestId] = options
+      const sub = currentSubscriptions.get(queryCacheKey)
+      if (sub?.has(requestId)) {
+        sub.set(requestId, options)
       }
       return true
     }
     if (unsubscribeQueryResult.match(action)) {
       const { queryCacheKey, requestId } = action.payload
-      if (mutableState[queryCacheKey]) {
-        delete mutableState[queryCacheKey]![requestId]
+      const sub = currentSubscriptions.get(queryCacheKey)
+      if (sub) {
+        sub.delete(requestId)
       }
       return true
     }
     if (api.internalActions.removeQueryResult.match(action)) {
-      delete mutableState[action.payload.queryCacheKey]
+      currentSubscriptions.delete(action.payload.queryCacheKey)
       return true
     }
     if (queryThunk.pending.match(action)) {
       const {
         meta: { arg, requestId },
       } = action
-      const substate = (mutableState[arg.queryCacheKey] ??= {})
-      substate[`${requestId}_running`] = {}
+      const substate = getOrInsertComputed(
+        currentSubscriptions,
+        arg.queryCacheKey,
+        createNewMap,
+      )
       if (arg.subscribe) {
-        substate[requestId] =
-          arg.subscriptionOptions ?? substate[requestId] ?? {}
+        substate.set(
+          requestId,
+          arg.subscriptionOptions ?? substate.get(requestId) ?? {},
+        )
       }
       return true
     }
     let mutated = false
-    if (
-      queryThunk.fulfilled.match(action) ||
-      queryThunk.rejected.match(action)
-    ) {
-      const state = mutableState[action.meta.arg.queryCacheKey] || {}
-      const key = `${action.meta.requestId}_running`
-      mutated ||= !!state[key]
-      delete state[key]
-    }
+
     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] ?? {}
+        const substate = getOrInsertComputed(
+          currentSubscriptions,
+          arg.queryCacheKey,
+          createNewMap,
+        )
+        substate.set(
+          requestId,
+          arg.subscriptionOptions ?? substate.get(requestId) ?? {},
+        )
 
         mutated = true
       }
@@ -83,12 +88,12 @@ export const buildBatchedActionsHandler: InternalHandlerBuilder<
   const getSubscriptions = () => internalState.currentSubscriptions
   const getSubscriptionCount = (queryCacheKey: string) => {
     const subscriptions = getSubscriptions()
-    const subscriptionsForQueryArg = subscriptions[queryCacheKey] ?? {}
-    return countObjectKeys(subscriptionsForQueryArg)
+    const subscriptionsForQueryArg = subscriptions.get(queryCacheKey)
+    return subscriptionsForQueryArg?.size ?? 0
   }
   const isRequestSubscribed = (queryCacheKey: string, requestId: string) => {
     const subscriptions = getSubscriptions()
-    return !!subscriptions?.[queryCacheKey]?.[requestId]
+    return !!subscriptions?.get(queryCacheKey)?.get(requestId)
   }
 
   const subscriptionSelectors: SubscriptionSelectors = {
@@ -97,6 +102,21 @@ export const buildBatchedActionsHandler: InternalHandlerBuilder<
     isRequestSubscribed,
   }
 
+  function serializeSubscriptions(
+    currentSubscriptions: SubscriptionInternalState,
+  ): SubscriptionState {
+    // We now use nested Maps for subscriptions, instead of
+    // plain Records. Stringify this accordingly so we can
+    // convert it to the shape we need for the store.
+    return JSON.parse(
+      JSON.stringify(
+        Object.fromEntries(
+          [...currentSubscriptions].map(([k, v]) => [k, Object.fromEntries(v)]),
+        ),
+      ),
+    )
+  }
+
   return (
     action,
     mwApi,
@@ -106,13 +126,14 @@ export const buildBatchedActionsHandler: InternalHandlerBuilder<
   ] => {
     if (!previousSubscriptions) {
       // Initialize it the first time this handler runs
-      previousSubscriptions = JSON.parse(
-        JSON.stringify(internalState.currentSubscriptions),
+      previousSubscriptions = serializeSubscriptions(
+        internalState.currentSubscriptions,
       )
     }
 
     if (api.util.resetApiState.match(action)) {
-      previousSubscriptions = internalState.currentSubscriptions = {}
+      previousSubscriptions = {}
+      internalState.currentSubscriptions.clear()
       updateSyncTimer = null
       return [true, false]
     }
@@ -133,6 +154,15 @@ export const buildBatchedActionsHandler: InternalHandlerBuilder<
 
     let actionShouldContinue = true
 
+    // HACK Sneak the test-only polling state back out
+    if (
+      process.env.NODE_ENV === 'test' &&
+      typeof action.type === 'string' &&
+      action.type === `${api.reducerPath}/getPolling`
+    ) {
+      return [false, internalState.currentPolls] as any
+    }
+
     if (didMutate) {
       if (!updateSyncTimer) {
         // We only use the subscription state for the Redux DevTools at this point,
@@ -142,8 +172,8 @@ export const buildBatchedActionsHandler: InternalHandlerBuilder<
         // In 1.9, it was updated in a microtask, but now we do it at most every 500ms.
         updateSyncTimer = setTimeout(() => {
           // Deep clone the current subscription data
-          const newSubscriptions: SubscriptionState = JSON.parse(
-            JSON.stringify(internalState.currentSubscriptions),
+          const newSubscriptions: SubscriptionState = serializeSubscriptions(
+            internalState.currentSubscriptions,
           )
           // Figure out a smaller diff between original and current
           const [, patches] = produceWithPatches(

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

@@ -35,10 +35,13 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
   internalState,
   selectors: { selectQueryEntry, selectConfig },
   getRunningQueryThunk,
+  mwApi,
 }) => {
   const { removeQueryResult, unsubscribeQueryResult, cacheEntriesUpserted } =
     api.internalActions
 
+  const runningQueries = internalState.runningQueries.get(mwApi.dispatch)!
+
   const canTriggerUnsubscribe = isAnyOf(
     unsubscribeQueryResult.match,
     queryThunk.fulfilled,
@@ -47,19 +50,14 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
   )
 
   function anySubscriptionsRemainingForKey(queryCacheKey: string) {
-    const subscriptions = internalState.currentSubscriptions[queryCacheKey]
+    const subscriptions = internalState.currentSubscriptions.get(queryCacheKey)
     if (!subscriptions) {
       return false
     }
 
-    // Check if there are any keys that are NOT _running subscriptions
-    for (const key in subscriptions) {
-      if (!key.endsWith('_running')) {
-        return true
-      }
-    }
-    // Only _running subscriptions remain (or empty)
-    return false
+    const hasSubscriptions = subscriptions.size > 0
+    const isRunning = runningQueries?.[queryCacheKey] !== undefined
+    return hasSubscriptions || isRunning
   }
 
   const currentRemovalTimeouts: QueryStateMeta<TimeoutId> = {}

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

@@ -45,7 +45,7 @@ export function buildMiddleware<
   ReducerPath extends string,
   TagTypes extends string,
 >(input: BuildMiddlewareInput<Definitions, ReducerPath, TagTypes>) {
-  const { reducerPath, queryThunk, api, context } = input
+  const { reducerPath, queryThunk, api, context, internalState } = input
   const { apiUid } = context
 
   const actions = {
@@ -73,10 +73,6 @@ export function buildMiddleware<
   > = (mwApi) => {
     let initialized = false
 
-    const internalState: InternalMiddlewareState = {
-      currentSubscriptions: {},
-    }
-
     const builderArgs = {
       ...(input as any as BuildMiddlewareInput<
         EndpointDefinitions,
@@ -86,6 +82,7 @@ export function buildMiddleware<
       internalState,
       refetchQuery,
       isThisApiSliceAction,
+      mwApi,
     }
 
     const handlers = handlerBuilders.map((build) => build(builderArgs))

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

@@ -17,9 +17,8 @@ import type {
   SubMiddlewareApi,
   InternalHandlerBuilder,
   ApiMiddlewareInternalHandler,
-  InternalMiddlewareState,
 } from './types'
-import { countObjectKeys } from '../../utils/countObjectKeys'
+import { getOrInsertComputed, createNewMap } from '../../utils/getOrInsert'
 
 export const buildInvalidationByTagsHandler: InternalHandlerBuilder = ({
   reducerPath,
@@ -111,11 +110,14 @@ export const buildInvalidationByTagsHandler: InternalHandlerBuilder = ({
       const valuesArray = Array.from(toInvalidate.values())
       for (const { queryCacheKey } of valuesArray) {
         const querySubState = state.queries[queryCacheKey]
-        const subscriptionSubState =
-          internalState.currentSubscriptions[queryCacheKey] ?? {}
+        const subscriptionSubState = getOrInsertComputed(
+          internalState.currentSubscriptions,
+          queryCacheKey,
+          createNewMap,
+        )
 
         if (querySubState) {
-          if (countObjectKeys(subscriptionSubState) === 0) {
+          if (subscriptionSubState.size === 0) {
             mwApi.dispatch(
               removeQueryResult({
                 queryCacheKey: queryCacheKey as QueryCacheKey,