Rewrite polling handling for perf

# Conflicts:
#	packages/toolkit/src/query/core/buildMiddleware/polling.ts
#	packages/toolkit/src/query/core/buildMiddleware/types.ts
#	packages/toolkit/src/query/tests/polling.test.tsx

Author: markerikson

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

@@ -21,23 +21,25 @@ export const buildPollingHandler: InternalHandlerBuilder = ({
   refetchQuery,
   internalState,
 }) => {
-  const { currentPolls } = internalState
+  const { currentPolls, currentSubscriptions } = internalState
 
-  const { currentSubscriptions } = internalState
+  // Batching state for polling updates
+  const pendingPollingUpdates = new Set<string>()
+  let pollingUpdateTimer: ReturnType<typeof setTimeout> | null = null
 
   const handler: ApiMiddlewareInternalHandler = (action, mwApi) => {
     if (
       api.internalActions.updateSubscriptionOptions.match(action) ||
       api.internalActions.unsubscribeQueryResult.match(action)
     ) {
-      updatePollingInterval(action.payload, mwApi)
+      schedulePollingUpdate(action.payload.queryCacheKey, mwApi)
     }
 
     if (
       queryThunk.pending.match(action) ||
       (queryThunk.rejected.match(action) && action.meta.condition)
     ) {
-      updatePollingInterval(action.meta.arg, mwApi)
+      schedulePollingUpdate(action.meta.arg.queryCacheKey, mwApi)
     }
 
     if (
@@ -49,6 +51,27 @@ export const buildPollingHandler: InternalHandlerBuilder = ({
 
     if (api.util.resetApiState.match(action)) {
       clearPolls()
+      // Clear any pending updates
+      if (pollingUpdateTimer) {
+        clearTimeout(pollingUpdateTimer)
+        pollingUpdateTimer = null
+      }
+      pendingPollingUpdates.clear()
+    }
+  }
+
+  function schedulePollingUpdate(queryCacheKey: string, api: SubMiddlewareApi) {
+    pendingPollingUpdates.add(queryCacheKey)
+
+    if (!pollingUpdateTimer) {
+      pollingUpdateTimer = setTimeout(() => {
+        // Process all pending updates in a single batch
+        for (const key of pendingPollingUpdates) {
+          updatePollingInterval({ queryCacheKey: key as any }, api)
+        }
+        pendingPollingUpdates.clear()
+        pollingUpdateTimer = null
+      }, 0)
     }
   }
 
@@ -81,7 +104,7 @@ export const buildPollingHandler: InternalHandlerBuilder = ({
       findLowestPollingInterval(subscriptions)
     if (!Number.isFinite(lowestPollingInterval)) return
 
-    const currentPoll = currentPolls[queryCacheKey]
+    const currentPoll = currentPolls.get(queryCacheKey)
 
     if (currentPoll?.timeout) {
       clearTimeout(currentPoll.timeout)
@@ -90,7 +113,7 @@ export const buildPollingHandler: InternalHandlerBuilder = ({
 
     const nextPollTimestamp = Date.now() + lowestPollingInterval
 
-    currentPolls[queryCacheKey] = {
+    currentPolls.set(queryCacheKey, {
       nextPollTimestamp,
       pollingInterval: lowestPollingInterval,
       timeout: setTimeout(() => {
@@ -99,7 +122,7 @@ export const buildPollingHandler: InternalHandlerBuilder = ({
         }
         startNextPoll({ queryCacheKey }, api)
       }, lowestPollingInterval),
-    }
+    })
   }
 
   function updatePollingInterval(
@@ -117,6 +140,7 @@ export const buildPollingHandler: InternalHandlerBuilder = ({
     const { lowestPollingInterval } = findLowestPollingInterval(subscriptions)
 
     // HACK add extra data to track how many times this has been called in tests
+    // yes we're mutating a nonexistent field on a Map here
     if (process.env.NODE_ENV === 'test') {
       const updateCounters = ((currentPolls as any).pollUpdateCounters ??= {})
       updateCounters[queryCacheKey] ??= 0
@@ -128,7 +152,8 @@ export const buildPollingHandler: InternalHandlerBuilder = ({
       return
     }
 
-    const currentPoll = currentPolls[queryCacheKey]
+    const currentPoll = currentPolls.get(queryCacheKey)
+
     const nextPollTimestamp = Date.now() + lowestPollingInterval
 
     if (!currentPoll || nextPollTimestamp < currentPoll.nextPollTimestamp) {
@@ -137,15 +162,15 @@ export const buildPollingHandler: InternalHandlerBuilder = ({
   }
 
   function cleanupPollForKey(key: string) {
-    const existingPoll = currentPolls[key]
+    const existingPoll = currentPolls.get(key)
     if (existingPoll?.timeout) {
       clearTimeout(existingPoll.timeout)
     }
-    delete currentPolls[key]
+    currentPolls.delete(key)
   }
 
   function clearPolls() {
-    for (const key of Object.keys(currentPolls)) {
+    for (const key of currentPolls.keys()) {
       cleanupPollForKey(key)
     }
   }

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

@@ -37,13 +37,15 @@ import type { AllSelectors } from '../buildSelectors'
 export type QueryStateMeta<T> = Record<string, undefined | T>
 export type TimeoutId = ReturnType<typeof setTimeout>
 
+type QueryPollState = {
+  nextPollTimestamp: number
+  timeout?: TimeoutId
+  pollingInterval: number
+}
+
 export interface InternalMiddlewareState {
   currentSubscriptions: SubscriptionInternalState
-  currentPolls: QueryStateMeta<{
-    nextPollTimestamp: number
-    timeout?: TimeoutId
-    pollingInterval: number
-  }>
+  currentPolls: Map<string, QueryPollState>
   runningQueries: Map<
     Dispatch,
     Record<

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

@@ -621,7 +621,7 @@ export const coreModule = ({
 
     const internalState: InternalMiddlewareState = {
       currentSubscriptions: new Map(),
-      currentPolls: {},
+      currentPolls: new Map(),
       runningQueries: new Map(),
       runningMutations: new Map(),
     }

packages/toolkit/src/query/tests/infiniteQueries.test.ts

@@ -16,6 +16,7 @@ describe('Infinite queries', () => {
     name: string
   }
 
+  type HitCounter = { page: number; hitCounter: number }
   let counters: Record<string, number> = {}
   let queryCounter = 0
 
@@ -88,39 +89,41 @@ describe('Infinite queries', () => {
     }),
   })
 
-  let hitCounter = 0
-
-  type HitCounter = { page: number; hitCounter: number }
+  function createCountersApi() {
+    let hitCounter = 0
 
-  const countersApi = createApi({
-    baseQuery: fakeBaseQuery(),
-    tagTypes: ['Counter'],
-    endpoints: (build) => ({
-      counters: build.infiniteQuery<HitCounter, string, number>({
-        queryFn({ pageParam }) {
-          hitCounter++
+    const countersApi = createApi({
+      baseQuery: fakeBaseQuery(),
+      tagTypes: ['Counter'],
+      endpoints: (build) => ({
+        counters: build.infiniteQuery<HitCounter, string, number>({
+          queryFn({ pageParam }) {
+            hitCounter++
 
-          return { data: { page: pageParam, hitCounter } }
-        },
-        infiniteQueryOptions: {
-          initialPageParam: 0,
-          getNextPageParam: (
-            lastPage,
-            allPages,
-            lastPageParam,
-            allPageParams,
-          ) => lastPageParam + 1,
-        },
-        providesTags: ['Counter'],
-      }),
-      mutation: build.mutation<null, void>({
-        queryFn: async () => {
-          return { data: null }
-        },
-        invalidatesTags: ['Counter'],
+            return { data: { page: pageParam, hitCounter } }
+          },
+          infiniteQueryOptions: {
+            initialPageParam: 0,
+            getNextPageParam: (
+              lastPage,
+              allPages,
+              lastPageParam,
+              allPageParams,
+            ) => lastPageParam + 1,
+          },
+          providesTags: ['Counter'],
+        }),
+        mutation: build.mutation<null, void>({
+          queryFn: async () => {
+            return { data: null }
+          },
+          invalidatesTags: ['Counter'],
+        }),
       }),
-    }),
-  })
+    })
+
+    return countersApi
+  }
 
   let storeRef = setupApiStore(
     pokemonApi,
@@ -155,7 +158,6 @@ describe('Infinite queries', () => {
 
     counters = {}
 
-    hitCounter = 0
     queryCounter = 0
   })
 
@@ -411,6 +413,8 @@ describe('Infinite queries', () => {
       }
     }
 
+    const countersApi = createCountersApi()
+
     const storeRef = setupApiStore(
       countersApi,
       { ...actionsReducer },
@@ -465,6 +469,8 @@ describe('Infinite queries', () => {
       }
     }
 
+    const countersApi = createCountersApi()
+
     const storeRef = setupApiStore(
       countersApi,
       { ...actionsReducer },
@@ -528,6 +534,7 @@ describe('Infinite queries', () => {
   })
 
   test('Refetches on polling', async () => {
+    const countersApi = createCountersApi()
     const checkResultData = (
       result: InfiniteQueryResult,
       expectedValues: HitCounter[],

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

@@ -1,4 +1,5 @@
 import { createApi } from '@reduxjs/toolkit/query'
+import type { QueryActionCreatorResult } from '@reduxjs/toolkit/query'
 import { delay } from 'msw'
 import { setupApiStore } from '../../tests/utils/helpers'
 import type { SubscriptionSelectors } from '../core/buildMiddleware/types'
@@ -29,6 +30,11 @@ beforeEach(() => {
   ;({ getSubscriptions } = storeRef.store.dispatch(
     api.internalActions.internal_getRTKQSubscriptions(),
   ) as unknown as SubscriptionSelectors)
+
+  const currentPolls = storeRef.store.dispatch({
+    type: `${api.reducerPath}/getPolling`,
+  }) as any
+  ;(currentPolls as any).pollUpdateCounters = {}
 })
 
 const getSubscribersForQueryCacheKey = (queryCacheKey: string) =>
@@ -67,13 +73,13 @@ describe('polling tests', () => {
 
     await delay(1)
     expect(getSubs().size).toBe(1)
-    expect(getSubs().get(requestId).pollingInterval).toBe(10)
+    expect(getSubs()?.get(requestId)?.pollingInterval).toBe(10)
 
     subscription.updateSubscriptionOptions({ pollingInterval: 20 })
 
     await delay(1)
     expect(getSubs().size).toBe(1)
-    expect(getSubs().get(requestId).pollingInterval).toBe(20)
+    expect(getSubs()?.get(requestId)?.pollingInterval).toBe(20)
   })
 
   it(`doesn't replace the interval when removing a shared query instance with a poll `, async () => {
@@ -155,7 +161,7 @@ describe('polling tests', () => {
     const callsWithoutSkip = mockBaseQuery.mock.calls.length
 
     expect(callsWithSkip).toBe(1)
-    expect(callsWithoutSkip).toBeGreaterThan(2)
+    expect(callsWithoutSkip).toBeGreaterThanOrEqual(2)
 
     storeRef.store.dispatch(api.util.resetApiState())
   })
@@ -230,4 +236,51 @@ describe('polling tests', () => {
     expect(getSubs().size).toBe(1)
     expect(getSubs().get(requestId)?.skipPollingIfUnfocused).toBe(true)
   })
+
+  it('should minimize polling recalculations when adding multiple subscribers', async () => {
+    // Reset any existing state
+    const storeRef = setupApiStore(api, undefined, {
+      withoutTestLifecycles: true,
+    })
+
+    const SUBSCRIBER_COUNT = 10
+    const subscriptions: QueryActionCreatorResult<any>[] = []
+
+    // Add 10 subscribers to the same endpoint with polling enabled
+    for (let i = 0; i < SUBSCRIBER_COUNT; i++) {
+      const subscription = storeRef.store.dispatch(
+        getPosts.initiate(1, {
+          subscriptionOptions: { pollingInterval: 1000 },
+          subscribe: true,
+        }),
+      )
+      subscriptions.push(subscription)
+    }
+
+    // Wait a bit for all subscriptions to be processed
+    await Promise.all(subscriptions)
+
+    // Wait for the poll update timer
+    await delay(25)
+
+    // Get the polling state using the secret "getPolling" action
+    const currentPolls = storeRef.store.dispatch({
+      type: `${api.reducerPath}/getPolling`,
+    }) as any
+
+    // Get the query cache key for our endpoint
+    const queryCacheKey = subscriptions[0].queryCacheKey
+
+    // Check the poll update counters
+    const pollUpdateCounters = currentPolls.pollUpdateCounters || {}
+    const updateCount = pollUpdateCounters[queryCacheKey] || 0
+
+    // With batching optimization, this should be much lower than SUBSCRIBER_COUNT
+    // Ideally 1, but could be slightly higher due to timing
+    expect(updateCount).toBeGreaterThanOrEqual(1)
+    expect(updateCount).toBeLessThanOrEqual(2)
+
+    // Clean up subscriptions
+    subscriptions.forEach((sub) => sub.unsubscribe())
+  })
 })