Implement aborting running queries on cache entry removal

Author: markerikson

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

@@ -1,5 +1,5 @@
 import type { QueryDefinition } from '../../endpointDefinitions'
-import type { ConfigState, QueryCacheKey } from '../apiState'
+import type { ConfigState, QueryCacheKey, QuerySubState } from '../apiState'
 import { isAnyOf } from '../rtkImports'
 import type {
   ApiMiddlewareInternalHandler,
@@ -11,16 +11,6 @@ import type {
 
 export type ReferenceCacheCollection = never
 
-function isObjectEmpty(obj: Record<any, any>) {
-  // Apparently a for..in loop is faster than `Object.keys()` here:
-  // https://stackoverflow.com/a/59787784/62937
-  for (const k in obj) {
-    // If there is at least one key, it's not empty
-    return false
-  }
-  return true
-}
-
 export type CacheCollectionQueryExtraOptions = {
   /**
    * Overrides the api-wide definition of `keepUnusedDataFor` for this endpoint only. _(This value is in seconds.)_
@@ -44,6 +34,7 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
   context,
   internalState,
   selectors: { selectQueryEntry, selectConfig },
+  getRunningQueryThunk,
 }) => {
   const { removeQueryResult, unsubscribeQueryResult, cacheEntriesUpserted } =
     api.internalActions
@@ -57,7 +48,18 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
 
   function anySubscriptionsRemainingForKey(queryCacheKey: string) {
     const subscriptions = internalState.currentSubscriptions[queryCacheKey]
-    return !!subscriptions && !isObjectEmpty(subscriptions)
+    if (!subscriptions) {
+      return false
+    }
+
+    // Check if there are any keys that are NOT _running subscriptions
+    for (const key in subscriptions) {
+      if (!key.endsWith('_running')) {
+        return true
+      }
+    }
+    // Only _running subscriptions remain (or empty)
+    return false
   }
 
   const currentRemovalTimeouts: QueryStateMeta<TimeoutId> = {}
@@ -69,6 +71,7 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
   ) => {
     const state = mwApi.getState()
     const config = selectConfig(state)
+
     if (canTriggerUnsubscribe(action)) {
       let queryCacheKeys: QueryCacheKey[]
 
@@ -114,18 +117,20 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
     const state = api.getState()
     for (const queryCacheKey of cacheKeys) {
       const entry = selectQueryEntry(state, queryCacheKey)
-      handleUnsubscribe(queryCacheKey, entry?.endpointName, api, config)
+      if (entry?.endpointName) {
+        handleUnsubscribe(queryCacheKey, entry.endpointName, api, config)
+      }
     }
   }
 
   function handleUnsubscribe(
     queryCacheKey: QueryCacheKey,
-    endpointName: string | undefined,
+    endpointName: string,
     api: SubMiddlewareApi,
     config: ConfigState<string>,
   ) {
     const endpointDefinition = context.endpointDefinitions[
-      endpointName!
+      endpointName
     ] as QueryDefinition<any, any, any, any>
     const keepUnusedDataFor =
       endpointDefinition?.keepUnusedDataFor ?? config.keepUnusedDataFor
@@ -151,6 +156,15 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
 
       currentRemovalTimeouts[queryCacheKey] = setTimeout(() => {
         if (!anySubscriptionsRemainingForKey(queryCacheKey)) {
+          // Try to abort any running query for this cache key
+          const entry = selectQueryEntry(api.getState(), queryCacheKey)
+
+          if (entry?.endpointName) {
+            const runningQuery = api.dispatch(
+              getRunningQueryThunk(entry.endpointName, entry.originalArgs),
+            )
+            runningQuery?.abort()
+          }
           api.dispatch(removeQueryResult({ queryCacheKey }))
         }
         delete currentRemovalTimeouts![queryCacheKey]

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

@@ -1190,6 +1190,89 @@ describe('hooks tests', () => {
       ).toBe(-1)
     })
 
+    test('query thunk should be aborted when component unmounts and cache entry is removed', async () => {
+      let abortSignalFromQueryFn: AbortSignal | undefined
+
+      const pokemonApi = createApi({
+        baseQuery: fetchBaseQuery({ baseUrl: 'https://pokeapi.co/api/v2/' }),
+        endpoints: (builder) => ({
+          getTest: builder.query<string, number>({
+            async queryFn(arg, { signal }) {
+              // Capture the abort signal for testing
+              abortSignalFromQueryFn = signal
+
+              // Simulate a long-running request that should be aborted
+              await new Promise((resolve, reject) => {
+                const timeout = setTimeout(resolve, 5000) // 5 second delay
+
+                signal.addEventListener('abort', () => {
+                  clearTimeout(timeout)
+                  reject(new Error('Aborted'))
+                })
+              })
+
+              return { data: 'data!' }
+            },
+            keepUnusedDataFor: 0.01, // Very short timeout (10ms)
+          }),
+        }),
+      })
+
+      const storeRef = setupApiStore(pokemonApi, undefined, {
+        withoutTestLifecycles: true,
+      })
+
+      function TestComponent() {
+        const { data, isFetching } = pokemonApi.endpoints.getTest.useQuery(1)
+
+        return (
+          <div>
+            <div data-testid="isFetching">{String(isFetching)}</div>
+            <div data-testid="data">{data || 'no data'}</div>
+          </div>
+        )
+      }
+
+      function App() {
+        const [showComponent, setShowComponent] = useState(true)
+
+        return (
+          <div>
+            {showComponent && <TestComponent />}
+            <button
+              data-testid="unmount"
+              onClick={() => setShowComponent(false)}
+            >
+              Unmount Component
+            </button>
+          </div>
+        )
+      }
+
+      render(<App />, { wrapper: storeRef.wrapper })
+
+      // Wait for the query to start
+      await waitFor(() =>
+        expect(screen.getByTestId('isFetching').textContent).toBe('true'),
+      )
+
+      // Verify we have an abort signal
+      expect(abortSignalFromQueryFn).toBeDefined()
+      expect(abortSignalFromQueryFn!.aborted).toBe(false)
+
+      // Unmount the component
+      fireEvent.click(screen.getByTestId('unmount'))
+
+      // Wait for the cache entry to be removed (keepUnusedDataFor: 0.01s = 10ms)
+      await act(async () => {
+        await delay(100) // Wait a bit longer than keepUnusedDataFor
+      })
+
+      // The abort signal should now be aborted
+      // This test is expected to fail initially until we implement the abort logic
+      expect(abortSignalFromQueryFn!.aborted).toBe(true)
+    })
+
     describe('Hook middleware requirements', () => {
       const consoleErrorSpy = vi
         .spyOn(console, 'error')
@@ -1898,7 +1981,6 @@ describe('hooks tests', () => {
       const checkNumQueries = (count: number) => {
         const cacheEntries = Object.keys(storeRef.store.getState().api.queries)
         const queries = cacheEntries.length
-        //console.log('queries', queries, storeRef.store.getState().api.queries)
 
         expect(queries).toBe(count)
       }