WIP Use middleware-internal state and fix failing RTKQ tests

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

@@ -3,7 +3,6 @@ import type { InternalHandlerBuilder } from './types'
 import type { SubscriptionState } from '../apiState'
 import { produceWithPatches } from 'immer'
 import util from 'util'
-
 // Copied from https://github.com/feross/queue-microtask
 let promise: Promise<any>
 const queueMicrotaskShim =
@@ -23,40 +22,49 @@ export const buildBatchedActionsHandler: InternalHandlerBuilder<
   const { subscriptionReducer } = api.internalActions.reducers
   const subscriptionsPrefix = `${api.reducerPath}/subscriptions`
 
-  let currentSubscriptions = subscriptionReducer(undefined, { type: 'init' })
-  let newSubscriptions = currentSubscriptions
+  let newSubscriptions: SubscriptionState = null as unknown as SubscriptionState
+  let previousSubscriptions: SubscriptionState =
+    null as unknown as SubscriptionState
 
   // let abortedQueryActionsQueue: RejectedAction<QueryThunk, any>[] = []
   let dispatchQueued = false
 
-  return (action, mwApi) => {
-    const latestState = dispatchQueued ? newSubscriptions : currentSubscriptions
+  return (action, mwApi, internalState) => {
+    if (!previousSubscriptions) {
+      previousSubscriptions = internalState.currentSubscriptions
+    }
 
     if (api.internalActions.internal_probeSubscription.match(action)) {
       const { queryCacheKey, requestId } = action.payload
-      const hasSubscription = !!latestState[queryCacheKey]?.[requestId]
+      const hasSubscription =
+        !!internalState.currentSubscriptions[queryCacheKey]?.[requestId]
       return [false, hasSubscription]
     }
 
-    newSubscriptions = subscriptionReducer(latestState, action)
+    newSubscriptions = subscriptionReducer(
+      internalState.currentSubscriptions,
+      action
+    )
 
-    if (newSubscriptions !== currentSubscriptions) {
+    if (newSubscriptions !== previousSubscriptions) {
       // console.log('Subscriptions changed: ')
 
       if (!dispatchQueued) {
         queueMicrotaskShim(() => {
           const [, patches] = produceWithPatches(
-            currentSubscriptions,
+            previousSubscriptions,
             () => newSubscriptions
           )
-          // console.log('Patches: ', util.inspect(patches, { depth: Infinity }))
+
           mwApi.next(api.internalActions.subscriptionsUpdated(patches))
-          currentSubscriptions = newSubscriptions
+          previousSubscriptions = newSubscriptions
           dispatchQueued = false
         })
         dispatchQueued = true
       }
 
+      internalState.currentSubscriptions = newSubscriptions
+
       const isSubscriptionSliceAction =
         !!action.type?.startsWith(subscriptionsPrefix)
       const isAdditionalSubscriptionAction =
@@ -66,7 +74,7 @@ export const buildBatchedActionsHandler: InternalHandlerBuilder<
 
       const actionShouldContinue =
         !isSubscriptionSliceAction && !isAdditionalSubscriptionAction
-      // console.log('Action: ', action.type, actionShouldContinue)
+
       return [actionShouldContinue, false]
     }
     /*

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

@@ -7,6 +7,7 @@ import type {
   TimeoutId,
   InternalHandlerBuilder,
   ApiMiddlewareInternalHandler,
+  InternalMiddlewareState,
 } from './types'
 
 export type ReferenceCacheCollection = never
@@ -54,16 +55,19 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
 
   function anySubscriptionsRemainingForKey(
     queryCacheKey: string,
-    api: SubMiddlewareApi
+    internalState: InternalMiddlewareState
   ) {
-    const subscriptions =
-      api.getState()[reducerPath].subscriptions[queryCacheKey]
+    const subscriptions = internalState.currentSubscriptions[queryCacheKey]
     return !!subscriptions && !isObjectEmpty(subscriptions)
   }
 
   const currentRemovalTimeouts: QueryStateMeta<TimeoutId> = {}
 
-  const handler: ApiMiddlewareInternalHandler = (action, mwApi) => {
+  const handler: ApiMiddlewareInternalHandler = (
+    action,
+    mwApi,
+    internalState
+  ) => {
     if (unsubscribeQueryResult.match(action)) {
       const state = mwApi.getState()[reducerPath]
       const { queryCacheKey } = action.payload
@@ -72,6 +76,7 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
         queryCacheKey,
         state.queries[queryCacheKey]?.endpointName,
         mwApi,
+        internalState,
         state.config
       )
     }
@@ -94,6 +99,7 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
           queryCacheKey as QueryCacheKey,
           queryState?.endpointName,
           mwApi,
+          internalState,
           state.config
         )
       }
@@ -104,6 +110,7 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
     queryCacheKey: QueryCacheKey,
     endpointName: string | undefined,
     api: SubMiddlewareApi,
+    internalState: InternalMiddlewareState,
     config: ConfigState<string>
   ) {
     const endpointDefinition = context.endpointDefinitions[
@@ -125,13 +132,13 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
       Math.min(keepUnusedDataFor, THIRTY_TWO_BIT_MAX_TIMER_SECONDS)
     )
 
-    if (!anySubscriptionsRemainingForKey(queryCacheKey, api)) {
+    if (!anySubscriptionsRemainingForKey(queryCacheKey, internalState)) {
       const currentTimeout = currentRemovalTimeouts[queryCacheKey]
       if (currentTimeout) {
         clearTimeout(currentTimeout)
       }
       currentRemovalTimeouts[queryCacheKey] = setTimeout(() => {
-        if (!anySubscriptionsRemainingForKey(queryCacheKey, api)) {
+        if (!anySubscriptionsRemainingForKey(queryCacheKey, internalState)) {
           api.dispatch(removeQueryResult({ queryCacheKey }))
         }
         delete currentRemovalTimeouts![queryCacheKey]

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

@@ -197,6 +197,7 @@ export const buildCacheLifecycleHandler: InternalHandlerBuilder = ({
   const handler: ApiMiddlewareInternalHandler = (
     action,
     mwApi,
+    internalState,
     stateBefore
   ) => {
     const cacheKey = getCacheKey(action)

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

@@ -10,7 +10,11 @@ import type { QueryThunkArg } from '../buildThunks'
 import { buildCacheCollectionHandler } from './cacheCollection'
 import { buildInvalidationByTagsHandler } from './invalidationByTags'
 import { buildPollingHandler } from './polling'
-import type { BuildMiddlewareInput, InternalHandlerBuilder } from './types'
+import type {
+  BuildMiddlewareInput,
+  InternalHandlerBuilder,
+  InternalMiddlewareState,
+} from './types'
 import { buildWindowEventHandler } from './windowEventHandling'
 import { buildCacheLifecycleHandler } from './cacheLifecycle'
 import { buildQueryLifecycleHandler } from './queryLifecycle'
@@ -69,6 +73,13 @@ export function buildMiddleware<
     const batchedActionsHandler = buildBatchedActionsHandler(builderArgs)
     const windowEventsHandler = buildWindowEventHandler(builderArgs)
 
+    let internalState: InternalMiddlewareState = {
+      currentSubscriptions: api.internalActions.reducers.subscriptionReducer(
+        undefined,
+        { type: 'init' }
+      ),
+    }
+
     return (next) => {
       return (action) => {
         if (!initialized) {
@@ -84,9 +95,11 @@ export function buildMiddleware<
         const [actionShouldContinue, hasSubscription] = batchedActionsHandler(
           action,
           mwApiWithNext,
+          internalState,
           stateBefore
         )
         let res: any
+
         if (actionShouldContinue) {
           res = next(action)
         } else {
@@ -97,7 +110,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, stateBefore)
+          windowEventsHandler(action, mwApiWithNext, internalState, stateBefore)
 
           if (
             isThisApiSliceAction(action) ||
@@ -106,7 +119,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, stateBefore)
+              handler(action, mwApiWithNext, internalState, stateBefore)
             }
           }
         }

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

@@ -6,6 +6,7 @@ import type {
   TimeoutId,
   InternalHandlerBuilder,
   ApiMiddlewareInternalHandler,
+  InternalMiddlewareState,
 } from './types'
 
 export const buildPollingHandler: InternalHandlerBuilder = ({
@@ -20,26 +21,30 @@ export const buildPollingHandler: InternalHandlerBuilder = ({
     pollingInterval: number
   }> = {}
 
-  const handler: ApiMiddlewareInternalHandler = (action, mwApi) => {
+  const handler: ApiMiddlewareInternalHandler = (
+    action,
+    mwApi,
+    internalState
+  ) => {
     if (
       api.internalActions.updateSubscriptionOptions.match(action) ||
       api.internalActions.unsubscribeQueryResult.match(action)
     ) {
-      updatePollingInterval(action.payload, mwApi)
+      updatePollingInterval(action.payload, mwApi, internalState)
     }
 
     if (
       queryThunk.pending.match(action) ||
       (queryThunk.rejected.match(action) && action.meta.condition)
     ) {
-      updatePollingInterval(action.meta.arg, mwApi)
+      updatePollingInterval(action.meta.arg, mwApi, internalState)
     }
 
     if (
       queryThunk.fulfilled.match(action) ||
       (queryThunk.rejected.match(action) && !action.meta.condition)
     ) {
-      startNextPoll(action.meta.arg, mwApi)
+      startNextPoll(action.meta.arg, mwApi, internalState)
     }
 
     if (api.util.resetApiState.match(action)) {
@@ -49,11 +54,12 @@ export const buildPollingHandler: InternalHandlerBuilder = ({
 
   function startNextPoll(
     { queryCacheKey }: QuerySubstateIdentifier,
-    api: SubMiddlewareApi
+    api: SubMiddlewareApi,
+    internalState: InternalMiddlewareState
   ) {
     const state = api.getState()[reducerPath]
     const querySubState = state.queries[queryCacheKey]
-    const subscriptions = state.subscriptions[queryCacheKey]
+    const subscriptions = internalState.currentSubscriptions[queryCacheKey]
 
     if (!querySubState || querySubState.status === QueryStatus.uninitialized)
       return
@@ -84,11 +90,12 @@ export const buildPollingHandler: InternalHandlerBuilder = ({
 
   function updatePollingInterval(
     { queryCacheKey }: QuerySubstateIdentifier,
-    api: SubMiddlewareApi
+    api: SubMiddlewareApi,
+    internalState: InternalMiddlewareState
   ) {
     const state = api.getState()[reducerPath]
     const querySubState = state.queries[queryCacheKey]
-    const subscriptions = state.subscriptions[queryCacheKey]
+    const subscriptions = internalState.currentSubscriptions[queryCacheKey]
 
     if (!querySubState || querySubState.status === QueryStatus.uninitialized) {
       return
@@ -105,7 +112,7 @@ export const buildPollingHandler: InternalHandlerBuilder = ({
     const nextPollTimestamp = Date.now() + lowestPollingInterval
 
     if (!currentPoll || nextPollTimestamp < currentPoll.nextPollTimestamp) {
-      startNextPoll({ queryCacheKey }, api)
+      startNextPoll({ queryCacheKey }, api, internalState)
     }
   }
 

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

@@ -12,7 +12,12 @@ import type {
   AssertTagTypes,
   EndpointDefinitions,
 } from '../../endpointDefinitions'
-import type { QueryStatus, QuerySubState, RootState } from '../apiState'
+import type {
+  QueryStatus,
+  QuerySubState,
+  RootState,
+  SubscriptionState,
+} from '../apiState'
 import type {
   MutationThunk,
   QueryThunk,
@@ -23,6 +28,10 @@ import type {
 export type QueryStateMeta<T> = Record<string, undefined | T>
 export type TimeoutId = ReturnType<typeof setTimeout>
 
+export interface InternalMiddlewareState {
+  currentSubscriptions: SubscriptionState
+}
+
 export interface BuildMiddlewareInput<
   Definitions extends EndpointDefinitions,
   ReducerPath extends string,
@@ -64,6 +73,7 @@ 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,6 +4,7 @@ import { onFocus, onOnline } from '../setupListeners'
 import type {
   ApiMiddlewareInternalHandler,
   InternalHandlerBuilder,
+  InternalMiddlewareState,
   SubMiddlewareApi,
 } from './types'
 
@@ -15,22 +16,27 @@ export const buildWindowEventHandler: InternalHandlerBuilder = ({
 }) => {
   const { removeQueryResult } = api.internalActions
 
-  const handler: ApiMiddlewareInternalHandler = (action, mwApi) => {
+  const handler: ApiMiddlewareInternalHandler = (
+    action,
+    mwApi,
+    internalState
+  ) => {
     if (onFocus.match(action)) {
-      refetchValidQueries(mwApi, 'refetchOnFocus')
+      refetchValidQueries(mwApi, 'refetchOnFocus', internalState)
     }
     if (onOnline.match(action)) {
-      refetchValidQueries(mwApi, 'refetchOnReconnect')
+      refetchValidQueries(mwApi, 'refetchOnReconnect', internalState)
     }
   }
 
   function refetchValidQueries(
     api: SubMiddlewareApi,
-    type: 'refetchOnFocus' | 'refetchOnReconnect'
+    type: 'refetchOnFocus' | 'refetchOnReconnect',
+    internalState: InternalMiddlewareState
   ) {
     const state = api.getState()[reducerPath]
     const queries = state.queries
-    const subscriptions = state.subscriptions
+    const subscriptions = internalState.currentSubscriptions
 
     context.batch(() => {
       for (const queryCacheKey of Object.keys(subscriptions)) {