Merge branch 'feature/5052-subscription-perf' of https://github.com/reduxjs/redux-toolkit into feature/5052-subscription-perf

Author: codesandbot

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

@@ -2,7 +2,7 @@ import type { InternalHandlerBuilder, SubscriptionSelectors } from './types'
 import type { SubscriptionInternalState, SubscriptionState } from '../apiState'
 import { produceWithPatches } from 'immer'
 import type { Action } from '@reduxjs/toolkit'
-import { getOrInsert } from '../../utils/getOrInsert'
+import { getOrInsertComputed, createNewMap } from '../../utils/getOrInsert'
 
 export const buildBatchedActionsHandler: InternalHandlerBuilder<
   [actionShouldContinue: boolean, returnValue: SubscriptionSelectors | boolean]
@@ -48,10 +48,10 @@ export const buildBatchedActionsHandler: InternalHandlerBuilder<
       const {
         meta: { arg, requestId },
       } = action
-      const substate = getOrInsert(
+      const substate = getOrInsertComputed(
         currentSubscriptions,
         arg.queryCacheKey,
-        new Map(),
+        createNewMap,
       )
       if (arg.subscribe) {
         substate.set(
@@ -68,10 +68,10 @@ export const buildBatchedActionsHandler: InternalHandlerBuilder<
         meta: { condition, arg, requestId },
       } = action
       if (condition && arg.subscribe) {
-        const substate = getOrInsert(
+        const substate = getOrInsertComputed(
           currentSubscriptions,
           arg.queryCacheKey,
-          new Map(),
+          createNewMap,
         )
         substate.set(
           requestId,
@@ -88,9 +88,8 @@ export const buildBatchedActionsHandler: InternalHandlerBuilder<
   const getSubscriptions = () => internalState.currentSubscriptions
   const getSubscriptionCount = (queryCacheKey: string) => {
     const subscriptions = getSubscriptions()
-    const subscriptionsForQueryArg =
-      subscriptions.get(queryCacheKey) ?? new Map()
-    return subscriptionsForQueryArg.size
+    const subscriptionsForQueryArg = subscriptions.get(queryCacheKey)
+    return subscriptionsForQueryArg?.size ?? 0
   }
   const isRequestSubscribed = (queryCacheKey: string, requestId: string) => {
     const subscriptions = getSubscriptions()

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

@@ -18,7 +18,7 @@ import type {
   InternalHandlerBuilder,
   ApiMiddlewareInternalHandler,
 } from './types'
-import { getOrInsert } from '../../utils/getOrInsert'
+import { getOrInsertComputed, createNewMap } from '../../utils/getOrInsert'
 
 export const buildInvalidationByTagsHandler: InternalHandlerBuilder = ({
   reducerPath,
@@ -110,10 +110,10 @@ export const buildInvalidationByTagsHandler: InternalHandlerBuilder = ({
       const valuesArray = Array.from(toInvalidate.values())
       for (const { queryCacheKey } of valuesArray) {
         const querySubState = state.queries[queryCacheKey]
-        const subscriptionSubState = getOrInsert(
+        const subscriptionSubState = getOrInsertComputed(
           internalState.currentSubscriptions,
           queryCacheKey,
-          new Map(),
+          createNewMap,
         )
 
         if (querySubState) {

packages/toolkit/src/query/utils/getOrInsert.ts

@@ -1,3 +1,7 @@
+// Duplicate some of the utils in `/src/utils` to ensure
+// we don't end up dragging in larger chunks of the RTK core
+// into the RTKQ bundle
+
 export function getOrInsert<K extends object, V>(
   map: WeakMap<K, V>,
   key: K,
@@ -13,3 +17,25 @@ export function getOrInsert<K extends object, V>(
 
   return map.set(key, value).get(key) as V
 }
+
+export function getOrInsertComputed<K extends object, V>(
+  map: WeakMap<K, V>,
+  key: K,
+  compute: (key: K) => V,
+): V
+export function getOrInsertComputed<K, V>(
+  map: Map<K, V>,
+  key: K,
+  compute: (key: K) => V,
+): V
+export function getOrInsertComputed<K extends object, V>(
+  map: Map<K, V> | WeakMap<K, V>,
+  key: K,
+  compute: (key: K) => V,
+): V {
+  if (map.has(key)) return map.get(key) as V
+
+  return map.set(key, compute(key)).get(key) as V
+}
+
+export const createNewMap = () => new Map()