Fixed memory leak in rapid hook arg changing

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

@@ -49,6 +49,7 @@ export const THIRTY_TWO_BIT_MAX_TIMER_SECONDS = 2_147_483_647 / 1_000 - 1
 export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
   reducerPath,
   api,
+  queryThunk,
   context,
   internalState,
 }) => {
@@ -66,9 +67,18 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
     mwApi,
     internalState
   ) => {
-    if (unsubscribeQueryResult.match(action)) {
+    if (unsubscribeQueryResult.match(action) ||
+      queryThunk.fulfilled.match(action) ||
+      queryThunk.rejected.match(action)
+    ) {
       const state = mwApi.getState()[reducerPath]
-      const { queryCacheKey } = action.payload
+      let queryCacheKey: QueryCacheKey
+
+      if (unsubscribeQueryResult.match(action)) {
+        queryCacheKey = action.payload.queryCacheKey
+      } else {
+        queryCacheKey = action.meta.arg.queryCacheKey
+      }
 
       handleUnsubscribe(
         queryCacheKey,

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

@@ -730,6 +730,67 @@ 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'))
+
+      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({