Rewrite RTKQ internal subscription lookups and subscription syncing

- Changed `subscriptionUpdated` from a microtask to a 500ms throttle
- Exposed the RTKQ internal middleware state to be returned via an
  internal action, so that hooks can read that state directly without
  needing to call `dispatch()` on every render.
- Reworked tests to completely ignore `subscriptionsUpdated` in any
  action sequence checks due to timing changes and irrelevance
- Fixed case where `invalidateTags` was still reading from store state

Author: markerikson

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

@@ -265,19 +265,16 @@ export function buildInitiate({
   function middlewareWarning(dispatch: Dispatch) {
     if (process.env.NODE_ENV !== 'production') {
       if ((middlewareWarning as any).triggered) return
-      const registered:
-        | ReturnType<typeof api.internalActions.internal_probeSubscription>
-        | boolean = dispatch(
-        api.internalActions.internal_probeSubscription({
-          queryCacheKey: 'DOES_NOT_EXIST',
-          requestId: 'DUMMY_REQUEST_ID',
-        })
-      )
+      const returnedValue = dispatch(api.internalActions.getRTKQInternalState())
 
       ;(middlewareWarning as any).triggered = true
 
-      // The RTKQ middleware _should_ always return a boolean for `probeSubscription`
-      if (typeof registered !== 'boolean') {
+      // The RTKQ middleware should return the internal state object,
+      // but it should _not_ be the action object.
+      if (
+        typeof returnedValue !== 'object' ||
+        typeof returnedValue?.type === 'string'
+      ) {
         // Otherwise, must not have been added
         throw new Error(
           `Warning: Middleware for RTK-Query API at reducerPath "${api.reducerPath}" has not been added to the store.

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

@@ -1,17 +1,20 @@
-import type { InternalHandlerBuilder } from './types'
+import type { InternalHandlerBuilder, InternalMiddlewareState } from './types'
 import type { SubscriptionState } from '../apiState'
 import { produceWithPatches } from 'immer'
 import type { Action } from '@reduxjs/toolkit'
 
 export const buildBatchedActionsHandler: InternalHandlerBuilder<
-  [actionShouldContinue: boolean, subscriptionExists: boolean]
+  [
+    actionShouldContinue: boolean,
+    returnValue: InternalMiddlewareState | boolean
+  ]
 > = ({ api, queryThunk, internalState }) => {
   const subscriptionsPrefix = `${api.reducerPath}/subscriptions`
 
   let previousSubscriptions: SubscriptionState =
     null as unknown as SubscriptionState
 
-  let dispatchQueued = false
+  let updateSyncTimer: ReturnType<typeof window.setTimeout> | null = null
 
   const { updateSubscriptionOptions, unsubscribeQueryResult } =
     api.internalActions
@@ -82,7 +85,10 @@ export const buildBatchedActionsHandler: InternalHandlerBuilder<
   return (
     action,
     mwApi
-  ): [actionShouldContinue: boolean, hasSubscription: boolean] => {
+  ): [
+    actionShouldContinue: boolean,
+    result: InternalMiddlewareState | boolean
+  ] => {
     if (!previousSubscriptions) {
       // Initialize it the first time this handler runs
       previousSubscriptions = JSON.parse(
@@ -92,16 +98,16 @@ export const buildBatchedActionsHandler: InternalHandlerBuilder<
 
     if (api.util.resetApiState.match(action)) {
       previousSubscriptions = internalState.currentSubscriptions = {}
+      updateSyncTimer = null
       return [true, false]
     }
 
     // Intercept requests by hooks to see if they're subscribed
-    // Necessary because we delay updating store state to the end of the tick
-    if (api.internalActions.internal_probeSubscription.match(action)) {
-      const { queryCacheKey, requestId } = action.payload
-      const hasSubscription =
-        !!internalState.currentSubscriptions[queryCacheKey]?.[requestId]
-      return [false, hasSubscription]
+    // We return the internal state reference so that hooks
+    // can do their own checks to see if they're still active.
+    // It's stupid and hacky, but it does cut down on some dispatch calls.
+    if (api.internalActions.getRTKQInternalState.match(action)) {
+      return [false, internalState]
     }
 
     // Update subscription data based on this action
@@ -110,9 +116,16 @@ export const buildBatchedActionsHandler: InternalHandlerBuilder<
       action
     )
 
+    let actionShouldContinue = true
+
     if (didMutate) {
-      if (!dispatchQueued) {
-        queueMicrotask(() => {
+      if (!updateSyncTimer) {
+        // We only use the subscription state for the Redux DevTools at this point,
+        // as the real data is kept here in the middleware.
+        // Given that, we can throttle synchronizing this state significantly to
+        // save on overall perf.
+        // 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)
@@ -127,25 +140,23 @@ export const buildBatchedActionsHandler: InternalHandlerBuilder<
           mwApi.next(api.internalActions.subscriptionsUpdated(patches))
           // Save the cloned state for later reference
           previousSubscriptions = newSubscriptions
-          dispatchQueued = false
-        })
-        dispatchQueued = true
+          updateSyncTimer = null
+        }, 500)
       }
 
       const isSubscriptionSliceAction =
         typeof action.type == 'string' &&
         !!action.type.startsWith(subscriptionsPrefix)
+
       const isAdditionalSubscriptionAction =
         queryThunk.rejected.match(action) &&
         action.meta.condition &&
         !!action.meta.arg.subscribe
 
-      const actionShouldContinue =
+      actionShouldContinue =
         !isSubscriptionSliceAction && !isAdditionalSubscriptionAction
-
-      return [actionShouldContinue, false]
     }
 
-    return [true, false]
+    return [actionShouldContinue, false]
   }
 }

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

@@ -71,6 +71,7 @@ export function buildMiddleware<
       >),
       internalState,
       refetchQuery,
+      isThisApiSliceAction,
     }
 
     const handlers = handlerBuilders.map((build) => build(builderArgs))
@@ -93,18 +94,15 @@ export function buildMiddleware<
 
         const stateBefore = mwApi.getState()
 
-        const [actionShouldContinue, hasSubscription] = batchedActionsHandler(
-          action,
-          mwApiWithNext,
-          stateBefore
-        )
+        const [actionShouldContinue, internalProbeResult] =
+          batchedActionsHandler(action, mwApiWithNext, stateBefore)
 
         let res: any
 
         if (actionShouldContinue) {
           res = next(action)
         } else {
-          res = hasSubscription
+          res = internalProbeResult
         }
 
         if (!!mwApi.getState()[reducerPath]) {

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

@@ -9,6 +9,7 @@ import type {
   SubMiddlewareApi,
   InternalHandlerBuilder,
   ApiMiddlewareInternalHandler,
+  InternalMiddlewareState,
 } from './types'
 
 export const buildInvalidationByTagsHandler: InternalHandlerBuilder = ({
@@ -19,6 +20,7 @@ export const buildInvalidationByTagsHandler: InternalHandlerBuilder = ({
   api,
   assertTagType,
   refetchQuery,
+  internalState,
 }) => {
   const { removeQueryResult } = api.internalActions
   const isThunkActionWithTags = isAnyOf(
@@ -35,7 +37,8 @@ export const buildInvalidationByTagsHandler: InternalHandlerBuilder = ({
           endpointDefinitions,
           assertTagType
         ),
-        mwApi
+        mwApi,
+        internalState
       )
     }
 
@@ -49,16 +52,19 @@ export const buildInvalidationByTagsHandler: InternalHandlerBuilder = ({
           undefined,
           assertTagType
         ),
-        mwApi
+        mwApi,
+        internalState
       )
     }
   }
 
   function invalidateTags(
     tags: readonly FullTagDescription<string>[],
-    mwApi: SubMiddlewareApi
+    mwApi: SubMiddlewareApi,
+    internalState: InternalMiddlewareState
   ) {
     const rootState = mwApi.getState()
+
     const state = rootState[reducerPath]
 
     const toInvalidate = api.util.selectInvalidatedBy(rootState, tags)
@@ -67,7 +73,8 @@ export const buildInvalidationByTagsHandler: InternalHandlerBuilder = ({
       const valuesArray = Array.from(toInvalidate.values())
       for (const { queryCacheKey } of valuesArray) {
         const querySubState = state.queries[queryCacheKey]
-        const subscriptionSubState = state.subscriptions[queryCacheKey] ?? {}
+        const subscriptionSubState =
+          internalState.currentSubscriptions[queryCacheKey] ?? {}
 
         if (querySubState) {
           if (Object.keys(subscriptionSubState).length === 0) {

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

@@ -61,6 +61,7 @@ export interface BuildSubMiddlewareInput
     queryCacheKey: string,
     override?: Partial<QueryThunkArg>
   ): AsyncThunkAction<ThunkResult, QueryThunkArg, {}>
+  isThisApiSliceAction: (action: Action) => boolean
 }
 
 export type SubMiddlewareBuilder = (

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

@@ -1,4 +1,4 @@
-import type { PayloadAction, UnknownAction } from '@reduxjs/toolkit'
+import type { Action, PayloadAction, UnknownAction } from '@reduxjs/toolkit'
 import {
   combineReducers,
   createAction,
@@ -443,12 +443,7 @@ export function buildSlice({
       ) {
         // Dummy
       },
-      internal_probeSubscription(
-        d,
-        a: PayloadAction<{ queryCacheKey: string; requestId: string }>
-      ) {
-        // dummy
-      },
+      getRTKQInternalState() {},
     },
   })
 

packages/toolkit/src/query/react/buildHooks.ts

@@ -53,6 +53,7 @@ import { UNINITIALIZED_VALUE } from './constants'
 import { useShallowStableValue } from './useShallowStableValue'
 import type { BaseQueryFn } from '../baseQueryTypes'
 import { defaultSerializeQueryArgs } from '../defaultSerializeQueryArgs'
+import { InternalMiddlewareState } from '../core/buildMiddleware/types'
 
 // Copy-pasted from React-Redux
 export const useIsomorphicLayoutEffect =
@@ -681,6 +682,27 @@ export function buildHooks<Definitions extends EndpointDefinitions>({
         Definitions
       >
       const dispatch = useDispatch<ThunkDispatch<any, any, UnknownAction>>()
+      const internalStateRef = useRef<InternalMiddlewareState | null>(null)
+      if (!internalStateRef.current) {
+        const returnedValue = dispatch(
+          api.internalActions.getRTKQInternalState()
+        )
+
+        if (process.env.NODE_ENV !== 'production') {
+          if (
+            typeof returnedValue !== 'object' ||
+            typeof returnedValue?.type === 'string'
+          ) {
+            throw new Error(
+              `Warning: Middleware for RTK-Query API at reducerPath "${api.reducerPath}" has not been added to the store.
+    You must add the middleware for RTK-Query to function correctly!`
+            )
+          }
+        }
+
+        internalStateRef.current =
+          returnedValue as unknown as InternalMiddlewareState
+      }
       const stableArg = useStableQueryArgs(
         skip ? skipToken : arg,
         // Even if the user provided a per-endpoint `serializeQueryArgs` with
@@ -704,28 +726,15 @@ export function buildHooks<Definitions extends EndpointDefinitions>({
 
       let { queryCacheKey, requestId } = promiseRef.current || {}
 
-      // HACK Because the latest state is in the middleware, we actually
-      // dispatch an action that will be intercepted and returned.
+      // HACK We've saved the middleware internal state into a ref,
+      // and that state object gets directly mutated. But, we've _got_ a reference
+      // to it locally, so we can just read the data directly here in the hook.
       let currentRenderHasSubscription = false
       if (queryCacheKey && requestId) {
-        // This _should_ return a boolean, even if the types don't line up
-        const returnedValue = dispatch(
-          api.internalActions.internal_probeSubscription({
-            queryCacheKey,
-            requestId,
-          })
-        )
-
-        if (process.env.NODE_ENV !== 'production') {
-          if (typeof returnedValue !== 'boolean') {
-            throw new Error(
-              `Warning: Middleware for RTK-Query API at reducerPath "${api.reducerPath}" has not been added to the store.
-    You must add the middleware for RTK-Query to function correctly!`
-            )
-          }
-        }
-
-        currentRenderHasSubscription = !!returnedValue
+        currentRenderHasSubscription =
+          !!internalStateRef.current?.currentSubscriptions?.[queryCacheKey]?.[
+            requestId
+          ]
       }
 
       const subscriptionRemoved =