Cleanup polls on unsubscribeQueryResult (#1933)

* Cleanup polls on unsubscribeQueryResult
* Add polling tests

Author: msutkowski

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

@@ -19,11 +19,15 @@ export const build: SubMiddlewareBuilder = ({
       timeout?: TimeoutId
       pollingInterval: number
     }> = {}
+
     return (next) =>
       (action): any => {
         const result = next(action)
 
-        if (api.internalActions.updateSubscriptionOptions.match(action)) {
+        if (
+          api.internalActions.updateSubscriptionOptions.match(action) ||
+          api.internalActions.unsubscribeQueryResult.match(action)
+        ) {
           updatePollingInterval(action.payload, mwApi)
         }
 
@@ -99,27 +103,31 @@ export const build: SubMiddlewareBuilder = ({
       }
 
       const lowestPollingInterval = findLowestPollingInterval(subscriptions)
-      const currentPoll = currentPolls[queryCacheKey]
 
       if (!Number.isFinite(lowestPollingInterval)) {
-        if (currentPoll?.timeout) {
-          clearTimeout(currentPoll.timeout)
-        }
-        delete currentPolls[queryCacheKey]
+        cleanupPollForKey(queryCacheKey)
         return
       }
 
+      const currentPoll = currentPolls[queryCacheKey]
       const nextPollTimestamp = Date.now() + lowestPollingInterval
 
       if (!currentPoll || nextPollTimestamp < currentPoll.nextPollTimestamp) {
         startNextPoll({ queryCacheKey }, api)
       }
     }
 
+    function cleanupPollForKey(key: string) {
+      const existingPoll = currentPolls[key]
+      if (existingPoll?.timeout) {
+        clearTimeout(existingPoll.timeout)
+      }
+      delete currentPolls[key]
+    }
+
     function clearPolls() {
-      for (const [key, poll] of Object.entries(currentPolls)) {
-        if (poll?.timeout) clearTimeout(poll.timeout)
-        delete currentPolls[key]
+      for (const key of Object.keys(currentPolls)) {
+        cleanupPollForKey(key)
       }
     }
   }

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

@@ -0,0 +1,110 @@
+import { createApi } from '@reduxjs/toolkit/query'
+import { setupApiStore, waitMs } from './helpers'
+
+const mockBaseQuery = jest
+  .fn()
+  .mockImplementation((args: any) => ({ data: args }))
+
+const api = createApi({
+  baseQuery: mockBaseQuery,
+  tagTypes: ['Posts'],
+  endpoints: (build) => ({
+    getPosts: build.query<unknown, number>({
+      query(pageNumber) {
+        return { url: 'posts', params: pageNumber }
+      },
+      providesTags: ['Posts'],
+    }),
+  }),
+})
+const { getPosts } = api.endpoints
+
+const storeRef = setupApiStore(api)
+
+const getSubscribersForQueryCacheKey = (queryCacheKey: string) =>
+  storeRef.store.getState()[api.reducerPath].subscriptions[queryCacheKey] || {}
+const createSubscriptionGetter = (queryCacheKey: string) => () =>
+  getSubscribersForQueryCacheKey(queryCacheKey)
+
+describe('polling tests', () => {
+  it('clears intervals when seeing a resetApiState action', async () => {
+    await storeRef.store.dispatch(
+      getPosts.initiate(1, {
+        subscriptionOptions: { pollingInterval: 10 },
+        subscribe: true,
+      })
+    )
+
+    expect(mockBaseQuery).toHaveBeenCalledTimes(1)
+
+    storeRef.store.dispatch(api.util.resetApiState())
+
+    await waitMs(30)
+
+    expect(mockBaseQuery).toHaveBeenCalledTimes(1)
+  })
+
+  it('replaces polling interval when the subscription options are updated', async () => {
+    const { requestId, queryCacheKey, ...subscription } =
+      storeRef.store.dispatch(
+        getPosts.initiate(1, {
+          subscriptionOptions: { pollingInterval: 10 },
+          subscribe: true,
+        })
+      )
+
+    const getSubs = createSubscriptionGetter(queryCacheKey)
+
+    expect(Object.keys(getSubs())).toHaveLength(1)
+    expect(getSubs()[requestId].pollingInterval).toBe(10)
+
+    subscription.updateSubscriptionOptions({ pollingInterval: 20 })
+
+    expect(Object.keys(getSubs())).toHaveLength(1)
+    expect(getSubs()[requestId].pollingInterval).toBe(20)
+  })
+
+  it(`doesn't replace the interval when removing a shared query instance with a poll `, async () => {
+    const subscriptionOne = storeRef.store.dispatch(
+      getPosts.initiate(1, {
+        subscriptionOptions: { pollingInterval: 10 },
+        subscribe: true,
+      })
+    )
+
+    storeRef.store.dispatch(
+      getPosts.initiate(1, {
+        subscriptionOptions: { pollingInterval: 10 },
+        subscribe: true,
+      })
+    )
+
+    const getSubs = createSubscriptionGetter(subscriptionOne.queryCacheKey)
+
+    expect(Object.keys(getSubs())).toHaveLength(2)
+
+    subscriptionOne.unsubscribe()
+
+    expect(Object.keys(getSubs())).toHaveLength(1)
+  })
+
+  it('uses lowest specified interval when two components are mounted', async () => {
+    storeRef.store.dispatch(
+      getPosts.initiate(1, {
+        subscriptionOptions: { pollingInterval: 30000 },
+        subscribe: true,
+      })
+    )
+
+    storeRef.store.dispatch(
+      getPosts.initiate(1, {
+        subscriptionOptions: { pollingInterval: 10 },
+        subscribe: true,
+      })
+    )
+
+    await waitMs(20)
+
+    expect(mockBaseQuery.mock.calls.length).toBeGreaterThanOrEqual(2)
+  })
+})