Return an object of selectors from the middleware probe

Author: markerikson

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

@@ -266,7 +266,9 @@ export function buildInitiate({
   function middlewareWarning(dispatch: Dispatch) {
     if (process.env.NODE_ENV !== 'production') {
       if ((middlewareWarning as any).triggered) return
-      const returnedValue = dispatch(api.internalActions.getRTKQInternalState())
+      const returnedValue = dispatch(
+        api.internalActions.internal_getRTKQSubscriptions()
+      )
 
       ;(middlewareWarning as any).triggered = true
 

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

@@ -1,13 +1,11 @@
-import type { InternalHandlerBuilder, InternalMiddlewareState } from './types'
+import type { InternalHandlerBuilder, SubscriptionSelectors } from './types'
 import type { SubscriptionState } from '../apiState'
 import { produceWithPatches } from 'immer'
 import type { Action } from '@reduxjs/toolkit'
+import { countObjectKeys } from '../../utils/countObjectKeys'
 
 export const buildBatchedActionsHandler: InternalHandlerBuilder<
-  [
-    actionShouldContinue: boolean,
-    returnValue: InternalMiddlewareState | boolean
-  ]
+  [actionShouldContinue: boolean, returnValue: SubscriptionSelectors | boolean]
 > = ({ api, queryThunk, internalState }) => {
   const subscriptionsPrefix = `${api.reducerPath}/subscriptions`
 
@@ -82,12 +80,29 @@ export const buildBatchedActionsHandler: InternalHandlerBuilder<
     return mutated
   }
 
+  const getSubscriptions = () => internalState.currentSubscriptions
+  const getSubscriptionCount = (queryCacheKey: string) => {
+    const subscriptions = getSubscriptions()
+    const subscriptionsForQueryArg = subscriptions[queryCacheKey] ?? {}
+    return countObjectKeys(subscriptionsForQueryArg)
+  }
+  const isRequestSubscribed = (queryCacheKey: string, requestId: string) => {
+    const subscriptions = getSubscriptions()
+    return !!subscriptions?.[queryCacheKey]?.[requestId]
+  }
+
+  const subscriptionSelectors: SubscriptionSelectors = {
+    getSubscriptions,
+    getSubscriptionCount,
+    isRequestSubscribed,
+  }
+
   return (
     action,
     mwApi
   ): [
     actionShouldContinue: boolean,
-    result: InternalMiddlewareState | boolean
+    result: SubscriptionSelectors | boolean
   ] => {
     if (!previousSubscriptions) {
       // Initialize it the first time this handler runs
@@ -106,8 +121,8 @@ export const buildBatchedActionsHandler: InternalHandlerBuilder<
     // 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]
+    if (api.internalActions.internal_getRTKQSubscriptions.match(action)) {
+      return [false, subscriptionSelectors]
     }
 
     // Update subscription data based on this action

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

@@ -32,6 +32,12 @@ export interface InternalMiddlewareState {
   currentSubscriptions: SubscriptionState
 }
 
+export interface SubscriptionSelectors {
+  getSubscriptions: () => SubscriptionState
+  getSubscriptionCount: (queryCacheKey: string) => number
+  isRequestSubscribed: (queryCacheKey: string, requestId: string) => boolean
+}
+
 export interface BuildMiddlewareInput<
   Definitions extends EndpointDefinitions,
   ReducerPath extends string,

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

@@ -443,7 +443,7 @@ export function buildSlice({
       ) {
         // Dummy
       },
-      getRTKQInternalState() {},
+      internal_getRTKQSubscriptions() {},
     },
   })
 

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

@@ -53,7 +53,10 @@ import { UNINITIALIZED_VALUE } from './constants'
 import { useShallowStableValue } from './useShallowStableValue'
 import type { BaseQueryFn } from '../baseQueryTypes'
 import { defaultSerializeQueryArgs } from '../defaultSerializeQueryArgs'
-import { InternalMiddlewareState } from '../core/buildMiddleware/types'
+import {
+  InternalMiddlewareState,
+  SubscriptionSelectors,
+} from '../core/buildMiddleware/types'
 
 // Copy-pasted from React-Redux
 export const useIsomorphicLayoutEffect =
@@ -682,10 +685,10 @@ export function buildHooks<Definitions extends EndpointDefinitions>({
         Definitions
       >
       const dispatch = useDispatch<ThunkDispatch<any, any, UnknownAction>>()
-      const internalStateRef = useRef<InternalMiddlewareState | null>(null)
-      if (!internalStateRef.current) {
+      const subscriptionSelectorsRef = useRef<SubscriptionSelectors>()
+      if (!subscriptionSelectorsRef.current) {
         const returnedValue = dispatch(
-          api.internalActions.getRTKQInternalState()
+          api.internalActions.internal_getRTKQSubscriptions()
         )
 
         if (process.env.NODE_ENV !== 'production') {
@@ -700,8 +703,8 @@ export function buildHooks<Definitions extends EndpointDefinitions>({
           }
         }
 
-        internalStateRef.current =
-          returnedValue as unknown as InternalMiddlewareState
+        subscriptionSelectorsRef.current =
+          returnedValue as unknown as SubscriptionSelectors
       }
       const stableArg = useStableQueryArgs(
         skip ? skipToken : arg,
@@ -726,15 +729,15 @@ export function buildHooks<Definitions extends EndpointDefinitions>({
 
       let { queryCacheKey, requestId } = promiseRef.current || {}
 
-      // 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.
+      // HACK We've saved the middleware subscription lookup callbacks into a ref,
+      // so we can directly check here if the subscription exists for this query.
       let currentRenderHasSubscription = false
       if (queryCacheKey && requestId) {
         currentRenderHasSubscription =
-          !!internalStateRef.current?.currentSubscriptions?.[queryCacheKey]?.[
+          subscriptionSelectorsRef.current.isRequestSubscribed(
+            queryCacheKey,
             requestId
-          ]
+          )
       }
 
       const subscriptionRemoved =

packages/toolkit/src/query/tests/buildHooks.test.tsx

@@ -37,7 +37,7 @@ import type { SubscriptionOptions } from '@reduxjs/toolkit/dist/query/core/apiSt
 import type { SerializedError } from '@reduxjs/toolkit'
 import { createListenerMiddleware, configureStore } from '@reduxjs/toolkit'
 import { delay } from '../../utils'
-import type { InternalMiddlewareState } from '../core/buildMiddleware/types'
+import type { SubscriptionSelectors } from '../core/buildMiddleware/types'
 import { countObjectKeys } from '../utils/countObjectKeys'
 
 // Just setup a temporary in-memory counter for tests that `getIncrementedAmount`.
@@ -140,18 +140,8 @@ const storeRef = setupApiStore(
   }
 )
 
-function getSubscriptions() {
-  const internalState = storeRef.store.dispatch(
-    api.internalActions.getRTKQInternalState()
-  ) as unknown as InternalMiddlewareState
-  return internalState?.currentSubscriptions ?? {}
-}
-
-function getSubscriptionCount(key: string) {
-  const subscriptions = getSubscriptions()
-  const subscriptionsForQueryArg = subscriptions[key] ?? {}
-  return countObjectKeys(subscriptionsForQueryArg)
-}
+let getSubscriptions: SubscriptionSelectors['getSubscriptions']
+let getSubscriptionCount: SubscriptionSelectors['getSubscriptionCount']
 
 beforeEach(() => {
   actions = []
@@ -161,6 +151,9 @@ beforeEach(() => {
       actions.push(action)
     },
   })
+  ;({ getSubscriptions, getSubscriptionCount } = storeRef.store.dispatch(
+    api.internalActions.internal_getRTKQSubscriptions()
+  ) as unknown as SubscriptionSelectors)
 })
 
 afterEach(() => {

packages/toolkit/src/query/tests/buildInitiate.test.tsx

@@ -1,5 +1,5 @@
 import { createApi } from '../core'
-import { InternalMiddlewareState } from '../core/buildMiddleware/types'
+import type { SubscriptionSelectors } from '../core/buildMiddleware/types'
 import { fakeBaseQuery } from '../fakeBaseQuery'
 import { setupApiStore } from './helpers'
 
@@ -25,16 +25,14 @@ const api = createApi({
 
 const storeRef = setupApiStore(api)
 
-function getSubscriptions() {
-  const internalState = storeRef.store.dispatch(
-    api.internalActions.getRTKQInternalState()
-  ) as unknown as InternalMiddlewareState
-  return internalState?.currentSubscriptions ?? {}
-}
-function isRequestSubscribed(key: string, requestId: string) {
-  const subscriptions = getSubscriptions()
-  return !!subscriptions?.[key]?.[requestId]
-}
+let getSubscriptions: SubscriptionSelectors['getSubscriptions']
+let isRequestSubscribed: SubscriptionSelectors['isRequestSubscribed']
+
+beforeEach(() => {
+  ;({ getSubscriptions, isRequestSubscribed } = storeRef.store.dispatch(
+    api.internalActions.internal_getRTKQSubscriptions()
+  ) as unknown as SubscriptionSelectors)
+})
 
 test('multiple synchonrous initiate calls with pre-existing cache entry', async () => {
   const { store, api } = storeRef

packages/toolkit/src/query/tests/cleanup.test.tsx

@@ -1,13 +1,12 @@
 // tests for "cleanup-after-unsubscribe" behaviour
 import { vi } from 'vitest'
-import React, { Profiler, ProfilerOnRenderCallback } from 'react'
+import React from 'react'
 
 import { createListenerMiddleware } from '@reduxjs/toolkit'
 import { createApi, QueryStatus } from '@reduxjs/toolkit/query/react'
 import { render, waitFor, act, screen } from '@testing-library/react'
 import { setupApiStore } from './helpers'
-import { InternalMiddlewareState } from '../core/buildMiddleware/types'
-import { countObjectKeys } from '../utils/countObjectKeys'
+import { SubscriptionSelectors } from '../core/buildMiddleware/types'
 
 const tick = () => new Promise((res) => setImmediate(res))
 
@@ -161,28 +160,26 @@ test('Minimizes the number of subscription dispatches when multiple components a
     withoutTestLifecycles: true,
   })
 
-  function getSubscriptions() {
-    const internalState = storeRef.store.dispatch(
-      api.internalActions.getRTKQInternalState()
-    ) as unknown as InternalMiddlewareState
-    return internalState?.currentSubscriptions ?? {}
-  }
-
-  let getSubscriptionsA = () => {
-    return getSubscriptions()['a(undefined)']
-  }
-
   let actionTypes: unknown[] = []
 
   listenerMiddleware.startListening({
     predicate: () => true,
     effect: (action) => {
-      if (!action.type.includes('subscriptionsUpdated')) {
-        actionTypes.push(action.type)
+      if (
+        action.type.includes('subscriptionsUpdated') ||
+        action.type.includes('internal_')
+      ) {
+        return
       }
+
+      actionTypes.push(action.type)
     },
   })
 
+  const { getSubscriptionCount } = storeRef.store.dispatch(
+    api.internalActions.internal_getRTKQSubscriptions()
+  ) as unknown as SubscriptionSelectors
+
   const NUM_LIST_ITEMS = 1000
 
   function ParentComponent() {
@@ -206,9 +203,7 @@ test('Minimizes the number of subscription dispatches when multiple components a
     return screen.getAllByText(/42/).length > 0
   })
 
-  const subscriptions = getSubscriptionsA()
-
-  expect(countObjectKeys(subscriptions!)).toBe(NUM_LIST_ITEMS)
+  expect(getSubscriptionCount('a(undefined)')).toBe(NUM_LIST_ITEMS)
 
   expect(actionTypes).toEqual([
     'api/config/middlewareRegistered',

packages/toolkit/src/query/tests/polling.test.tsx

@@ -2,7 +2,7 @@ import { vi } from 'vitest'
 import { createApi } from '@reduxjs/toolkit/query'
 import { setupApiStore, waitMs } from './helpers'
 import { delay } from '../../utils'
-import type { InternalMiddlewareState } from '../core/buildMiddleware/types'
+import type { SubscriptionSelectors } from '../core/buildMiddleware/types'
 
 const mockBaseQuery = vi
   .fn()
@@ -24,12 +24,13 @@ const { getPosts } = api.endpoints
 
 const storeRef = setupApiStore(api)
 
-function getSubscriptions() {
-  const internalState = storeRef.store.dispatch(
-    api.internalActions.getRTKQInternalState()
-  ) as unknown as InternalMiddlewareState
-  return internalState?.currentSubscriptions ?? {}
-}
+let getSubscriptions: SubscriptionSelectors['getSubscriptions']
+
+beforeEach(() => {
+  ;({ getSubscriptions } = storeRef.store.dispatch(
+    api.internalActions.internal_getRTKQSubscriptions()
+  ) as unknown as SubscriptionSelectors)
+})
 
 const getSubscribersForQueryCacheKey = (queryCacheKey: string) =>
   getSubscriptions()[queryCacheKey] || {}