Merge pull request #2789 from reduxjs/feature/v1.9-pre-beta-cleanup

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>
@@ -18,16 +23,71 @@ const queueMicrotaskShim =
 
 export const buildBatchedActionsHandler: InternalHandlerBuilder<
   [actionShouldContinue: boolean, subscriptionExists: boolean]
-> = ({ api, queryThunk }) => {
-  const { actuallyMutateSubscriptions } = api.internalActions
+> = ({ api, queryThunk, internalState }) => {
   const subscriptionsPrefix = `${api.reducerPath}/subscriptions`
 
   let previousSubscriptions: SubscriptionState =
     null as unknown as SubscriptionState
 
   let dispatchQueued = false
 
-  return (action, mwApi, internalState) => {
+  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
       previousSubscriptions = JSON.parse(

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

@@ -50,13 +50,11 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
   reducerPath,
   api,
   context,
+  internalState,
 }) => {
   const { removeQueryResult, unsubscribeQueryResult } = api.internalActions
 
-  function anySubscriptionsRemainingForKey(
-    queryCacheKey: string,
-    internalState: InternalMiddlewareState
-  ) {
+  function anySubscriptionsRemainingForKey(queryCacheKey: string) {
     const subscriptions = internalState.currentSubscriptions[queryCacheKey]
     return !!subscriptions && !isObjectEmpty(subscriptions)
   }
@@ -76,7 +74,6 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
         queryCacheKey,
         state.queries[queryCacheKey]?.endpointName,
         mwApi,
-        internalState,
         state.config
       )
     }
@@ -99,7 +96,6 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
           queryCacheKey as QueryCacheKey,
           queryState?.endpointName,
           mwApi,
-          internalState,
           state.config
         )
       }
@@ -110,7 +106,6 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
     queryCacheKey: QueryCacheKey,
     endpointName: string | undefined,
     api: SubMiddlewareApi,
-    internalState: InternalMiddlewareState,
     config: ConfigState<string>
   ) {
     const endpointDefinition = context.endpointDefinitions[
@@ -132,13 +127,13 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
       Math.min(keepUnusedDataFor, THIRTY_TWO_BIT_MAX_TIMER_SECONDS)
     )
 
-    if (!anySubscriptionsRemainingForKey(queryCacheKey, internalState)) {
+    if (!anySubscriptionsRemainingForKey(queryCacheKey)) {
       const currentTimeout = currentRemovalTimeouts[queryCacheKey]
       if (currentTimeout) {
         clearTimeout(currentTimeout)
       }
       currentRemovalTimeouts[queryCacheKey] = setTimeout(() => {
-        if (!anySubscriptionsRemainingForKey(queryCacheKey, internalState)) {
+        if (!anySubscriptionsRemainingForKey(queryCacheKey)) {
           api.dispatch(removeQueryResult({ queryCacheKey }))
         }
         delete currentRemovalTimeouts![queryCacheKey]

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

@@ -183,6 +183,7 @@ export const buildCacheLifecycleHandler: InternalHandlerBuilder = ({
   context,
   queryThunk,
   mutationThunk,
+  internalState,
 }) => {
   const isQueryThunk = isAsyncThunkAction(queryThunk)
   const isMutationThunk = isAsyncThunkAction(mutationThunk)
@@ -197,7 +198,6 @@ export const buildCacheLifecycleHandler: InternalHandlerBuilder = ({
   const handler: ApiMiddlewareInternalHandler = (
     action,
     mwApi,
-    internalState,
     stateBefore
   ) => {
     const cacheKey = getCacheKey(action)

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

@@ -59,12 +59,17 @@ export function buildMiddleware<
   > = (mwApi) => {
     let initialized = false
 
+    let internalState: InternalMiddlewareState = {
+      currentSubscriptions: {},
+    }
+
     const builderArgs = {
       ...(input as any as BuildMiddlewareInput<
         EndpointDefinitions,
         string,
         string
       >),
+      internalState,
       refetchQuery,
     }
 
@@ -73,10 +78,6 @@ export function buildMiddleware<
     const batchedActionsHandler = buildBatchedActionsHandler(builderArgs)
     const windowEventsHandler = buildWindowEventHandler(builderArgs)
 
-    let internalState: InternalMiddlewareState = {
-      currentSubscriptions: {},
-    }
-
     return (next) => {
       return (action) => {
         if (!initialized) {
@@ -92,7 +93,6 @@ export function buildMiddleware<
         const [actionShouldContinue, hasSubscription] = batchedActionsHandler(
           action,
           mwApiWithNext,
-          internalState,
           stateBefore
         )
 
@@ -108,7 +108,7 @@ export function buildMiddleware<
           // Only run these checks if the middleware is registered okay
 
           // This looks for actions that aren't specific to the API slice
-          windowEventsHandler(action, mwApiWithNext, internalState, stateBefore)
+          windowEventsHandler(action, mwApiWithNext, stateBefore)
 
           if (
             isThisApiSliceAction(action) ||
@@ -117,7 +117,7 @@ export function buildMiddleware<
             // Only run these additional checks if the actions are part of the API slice,
             // or the action has hydration-related data
             for (let handler of handlers) {
-              handler(action, mwApiWithNext, internalState, stateBefore)
+              handler(action, mwApiWithNext, stateBefore)
             }
           }
         }

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

@@ -14,37 +14,34 @@ export const buildPollingHandler: InternalHandlerBuilder = ({
   queryThunk,
   api,
   refetchQuery,
+  internalState,
 }) => {
   const currentPolls: QueryStateMeta<{
     nextPollTimestamp: number
     timeout?: TimeoutId
     pollingInterval: number
   }> = {}
 
-  const handler: ApiMiddlewareInternalHandler = (
-    action,
-    mwApi,
-    internalState
-  ) => {
+  const handler: ApiMiddlewareInternalHandler = (action, mwApi) => {
     if (
       api.internalActions.updateSubscriptionOptions.match(action) ||
       api.internalActions.unsubscribeQueryResult.match(action)
     ) {
-      updatePollingInterval(action.payload, mwApi, internalState)
+      updatePollingInterval(action.payload, mwApi)
     }
 
     if (
       queryThunk.pending.match(action) ||
       (queryThunk.rejected.match(action) && action.meta.condition)
     ) {
-      updatePollingInterval(action.meta.arg, mwApi, internalState)
+      updatePollingInterval(action.meta.arg, mwApi)
     }
 
     if (
       queryThunk.fulfilled.match(action) ||
       (queryThunk.rejected.match(action) && !action.meta.condition)
     ) {
-      startNextPoll(action.meta.arg, mwApi, internalState)
+      startNextPoll(action.meta.arg, mwApi)
     }
 
     if (api.util.resetApiState.match(action)) {
@@ -54,8 +51,7 @@ export const buildPollingHandler: InternalHandlerBuilder = ({
 
   function startNextPoll(
     { queryCacheKey }: QuerySubstateIdentifier,
-    api: SubMiddlewareApi,
-    internalState: InternalMiddlewareState
+    api: SubMiddlewareApi
   ) {
     const state = api.getState()[reducerPath]
     const querySubState = state.queries[queryCacheKey]
@@ -90,8 +86,7 @@ export const buildPollingHandler: InternalHandlerBuilder = ({
 
   function updatePollingInterval(
     { queryCacheKey }: QuerySubstateIdentifier,
-    api: SubMiddlewareApi,
-    internalState: InternalMiddlewareState
+    api: SubMiddlewareApi
   ) {
     const state = api.getState()[reducerPath]
     const querySubState = state.queries[queryCacheKey]
@@ -112,7 +107,7 @@ export const buildPollingHandler: InternalHandlerBuilder = ({
     const nextPollTimestamp = Date.now() + lowestPollingInterval
 
     if (!currentPoll || nextPollTimestamp < currentPoll.nextPollTimestamp) {
-      startNextPoll({ queryCacheKey }, api, internalState)
+      startNextPoll({ queryCacheKey }, api)
     }
   }
 

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

@@ -52,6 +52,7 @@ export type SubMiddlewareApi = MiddlewareAPI<
 
 export interface BuildSubMiddlewareInput
   extends BuildMiddlewareInput<EndpointDefinitions, string, string> {
+  internalState: InternalMiddlewareState
   refetchQuery(
     querySubState: Exclude<
       QuerySubState<any>,
@@ -73,7 +74,6 @@ export type SubMiddlewareBuilder = (
 export type ApiMiddlewareInternalHandler<ReturnType = void> = (
   action: AnyAction,
   mwApi: SubMiddlewareApi & { next: Dispatch<AnyAction> },
-  internalState: InternalMiddlewareState,
   prevState: RootState<EndpointDefinitions, string, string>
 ) => ReturnType
 

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

@@ -4,7 +4,6 @@ import { onFocus, onOnline } from '../setupListeners'
 import type {
   ApiMiddlewareInternalHandler,
   InternalHandlerBuilder,
-  InternalMiddlewareState,
   SubMiddlewareApi,
 } from './types'
 
@@ -13,26 +12,22 @@ export const buildWindowEventHandler: InternalHandlerBuilder = ({
   context,
   api,
   refetchQuery,
+  internalState,
 }) => {
   const { removeQueryResult } = api.internalActions
 
-  const handler: ApiMiddlewareInternalHandler = (
-    action,
-    mwApi,
-    internalState
-  ) => {
+  const handler: ApiMiddlewareInternalHandler = (action, mwApi) => {
     if (onFocus.match(action)) {
-      refetchValidQueries(mwApi, 'refetchOnFocus', internalState)
+      refetchValidQueries(mwApi, 'refetchOnFocus')
     }
     if (onOnline.match(action)) {
-      refetchValidQueries(mwApi, 'refetchOnReconnect', internalState)
+      refetchValidQueries(mwApi, 'refetchOnReconnect')
     }
   }
 
   function refetchValidQueries(
     api: SubMiddlewareApi,
-    type: 'refetchOnFocus' | 'refetchOnReconnect',
-    internalState: InternalMiddlewareState
+    type: 'refetchOnFocus' | 'refetchOnReconnect'
   ) {
     const state = api.getState()[reducerPath]
     const queries = state.queries