Fixed memory leak in rapid hook arg changing (#4268)

* Added unsubscribeHandler for fulfilled/rejected queryThunk

* Added testcase for memory leak listed in github Issues

* Removed unnecessary delay call

* Re-defined delay function and removed Math.random() in favour of i++

* Moved near identical solutions into one conditional

* Reformat for visual clarity

* Added additional check to clearly show resolution of queries

* Refactored unsubscribe matching to make better use of type narrowing

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

@@ -1,3 +1,4 @@
+import { isAnyOf } from '@reduxjs/toolkit'
 import type { BaseQueryFn } from '../../baseQueryTypes'
 import type { QueryDefinition } from '../../endpointDefinitions'
 import type { ConfigState, QueryCacheKey } from '../apiState'
@@ -49,11 +50,18 @@ export const THIRTY_TWO_BIT_MAX_TIMER_SECONDS = 2_147_483_647 / 1_000 - 1
 export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
   reducerPath,
   api,
+  queryThunk,
   context,
   internalState,
 }) => {
   const { removeQueryResult, unsubscribeQueryResult } = api.internalActions
 
+  const canTriggerUnsubscribe = isAnyOf(
+    unsubscribeQueryResult.match,
+    queryThunk.fulfilled,
+    queryThunk.rejected
+  )
+
   function anySubscriptionsRemainingForKey(queryCacheKey: string) {
     const subscriptions = internalState.currentSubscriptions[queryCacheKey]
     return !!subscriptions && !isObjectEmpty(subscriptions)
@@ -66,9 +74,11 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
     mwApi,
     internalState,
   ) => {
-    if (unsubscribeQueryResult.match(action)) {
+    if (canTriggerUnsubscribe(action)) {
       const state = mwApi.getState()[reducerPath]
-      const { queryCacheKey } = action.payload
+      const { queryCacheKey } = unsubscribeQueryResult.match(action)
+        ? action.payload
+        : action.meta.arg
 
       handleUnsubscribe(
         queryCacheKey,

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

@@ -724,6 +724,68 @@ describe('hooks tests', () => {
       expect(res.data!.amount).toBeGreaterThan(originalAmount)
     })
 
+    // See https://github.com/reduxjs/redux-toolkit/issues/4267 - Memory leak in useQuery rapid query arg changes
+    test('Hook subscriptions are properly cleaned up when query is fulfilled/rejected', async () => {
+      // This is imported already, but it seems to be causing issues with the test on certain matrixes
+      function delay(ms: number) {
+        return new Promise((resolve) => setTimeout(resolve, ms))
+      }
+
+      const pokemonApi = createApi({
+        baseQuery: fetchBaseQuery({ baseUrl: 'https://pokeapi.co/api/v2/' }),
+        endpoints: (builder) => ({
+          getTest: builder.query<string, number>({
+            async queryFn() {
+              await new Promise((resolve) => setTimeout(resolve, 1000));
+              return { data: "data!" };
+            },
+            keepUnusedDataFor: 0,
+          }),
+        }),
+      })
+
+      const storeRef = setupApiStore(pokemonApi, undefined, {
+        withoutTestLifecycles: true,
+      })
+
+      const checkNumQueries = (count: number) => {
+        const cacheEntries = Object.keys((storeRef.store.getState()).api.queries)
+        const queries = cacheEntries.length
+
+        expect(queries).toBe(count)
+      }
+
+      let i = 0;
+
+      function User() {
+        const [fetchTest, { isFetching, isUninitialized }] =
+          pokemonApi.endpoints.getTest.useLazyQuery()
+
+        return (
+          <div>
+            <div data-testid="isUninitialized">{String(isUninitialized)}</div>
+            <div data-testid="isFetching">{String(isFetching)}</div>
+            <button data-testid="fetchButton" onClick={() => fetchTest(i++)}>
+              fetchUser
+            </button>
+          </div>
+        )
+      }
+
+      render(<User />, { wrapper: storeRef.wrapper })
+      fireEvent.click(screen.getByTestId('fetchButton'))
+      fireEvent.click(screen.getByTestId('fetchButton'))
+      fireEvent.click(screen.getByTestId('fetchButton'))
+      checkNumQueries(3)
+
+      await act(async () => {
+        await delay(1500)
+      })
+
+      // There should only be one stored query once they have had time to resolve
+      checkNumQueries( 1)
+    })
+
     // See https://github.com/reduxjs/redux-toolkit/issues/3182
     test('Hook subscriptions are properly cleaned up when changing skip back and forth', async () => {
       const pokemonApi = createApi({