Merge pull request #3116 from GeorchW/delayed-tag-invalidations

Author: markerikson

packages/toolkit/src/query/apiTypes.ts

@@ -45,6 +45,7 @@ export type Module<Name extends ModuleName> = {
       | 'refetchOnMountOrArgChange'
       | 'refetchOnFocus'
       | 'refetchOnReconnect'
+      | 'invalidationBehavior'
       | 'tagTypes'
     >,
     context: ApiContext<Definitions>

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

@@ -254,6 +254,7 @@ export type ConfigState<ReducerPath> = RefetchConfigOptions & {
 
 export type ModifiableConfigState = {
   keepUnusedDataFor: number
+  invalidationBehavior: 'delayed' | 'immediately'
 } & RefetchConfigOptions
 
 export type MutationState<D extends EndpointDefinitions> = {

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

@@ -1,8 +1,16 @@
-import { isAnyOf, isFulfilled, isRejectedWithValue } from '../rtkImports'
+import {
+  isAnyOf,
+  isFulfilled,
+  isRejected,
+  isRejectedWithValue,
+} from '../rtkImports'
 
-import type { FullTagDescription } from '../../endpointDefinitions'
+import type {
+  EndpointDefinitions,
+  FullTagDescription,
+} from '../../endpointDefinitions'
 import { calculateProvidedBy } from '../../endpointDefinitions'
-import type { QueryCacheKey } from '../apiState'
+import type { CombinedState, QueryCacheKey } from '../apiState'
 import { QueryStatus } from '../apiState'
 import { calculateProvidedByThunk } from '../buildThunks'
 import type {
@@ -18,6 +26,7 @@ export const buildInvalidationByTagsHandler: InternalHandlerBuilder = ({
   context,
   context: { endpointDefinitions },
   mutationThunk,
+  queryThunk,
   api,
   assertTagType,
   refetchQuery,
@@ -29,6 +38,13 @@ export const buildInvalidationByTagsHandler: InternalHandlerBuilder = ({
     isRejectedWithValue(mutationThunk)
   )
 
+  const isQueryEnd = isAnyOf(
+    isFulfilled(mutationThunk, queryThunk),
+    isRejected(mutationThunk, queryThunk)
+  )
+
+  let pendingTagInvalidations: FullTagDescription<string>[] = []
+
   const handler: ApiMiddlewareInternalHandler = (action, mwApi) => {
     if (isThunkActionWithTags(action)) {
       invalidateTags(
@@ -38,12 +54,11 @@ export const buildInvalidationByTagsHandler: InternalHandlerBuilder = ({
           endpointDefinitions,
           assertTagType
         ),
-        mwApi,
-        internalState
+        mwApi
       )
-    }
-
-    if (api.util.invalidateTags.match(action)) {
+    } else if (isQueryEnd(action)) {
+      invalidateTags([], mwApi)
+    } else if (api.util.invalidateTags.match(action)) {
       invalidateTags(
         calculateProvidedBy(
           action.payload,
@@ -53,21 +68,38 @@ export const buildInvalidationByTagsHandler: InternalHandlerBuilder = ({
           undefined,
           assertTagType
         ),
-        mwApi,
-        internalState
+        mwApi
       )
     }
   }
 
+  function hasPendingRequests(state: CombinedState<EndpointDefinitions, string, string>) {
+    for (const key in state.queries) {
+      if (state.queries[key]?.status === QueryStatus.pending) return true;
+    }
+    for (const key in state.mutations) {
+      if (state.mutations[key]?.status === QueryStatus.pending) return true;
+    }
+    return false;
+  }
+
   function invalidateTags(
-    tags: readonly FullTagDescription<string>[],
-    mwApi: SubMiddlewareApi,
-    internalState: InternalMiddlewareState
+    newTags: readonly FullTagDescription<string>[],
+    mwApi: SubMiddlewareApi
   ) {
     const rootState = mwApi.getState()
-
     const state = rootState[reducerPath]
 
+    pendingTagInvalidations.push(...newTags)
+
+    if (state.config.invalidationBehavior === 'delayed' && hasPendingRequests(state)) {
+      return
+    }
+
+    const tags = pendingTagInvalidations
+    pendingTagInvalidations = []
+    if (tags.length === 0) return
+
     const toInvalidate = api.util.selectInvalidatedBy(rootState, tags)
 
     context.batch(() => {

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

@@ -452,6 +452,7 @@ export const coreModule = (): Module<CoreModule> => ({
       refetchOnMountOrArgChange,
       refetchOnFocus,
       refetchOnReconnect,
+      invalidationBehavior,
     },
     context
   ) {
@@ -514,6 +515,7 @@ export const coreModule = (): Module<CoreModule> => ({
         refetchOnMountOrArgChange,
         keepUnusedDataFor,
         reducerPath,
+        invalidationBehavior,
       },
     })
 

packages/toolkit/src/query/createApi.ts

@@ -151,6 +151,16 @@ export interface CreateApiOptions<
    * Note: requires [`setupListeners`](./setupListeners) to have been called.
    */
   refetchOnReconnect?: boolean
+  /**
+   * Defaults to `'immediately'`. This setting allows you to control when tags are invalidated after a mutation.
+   *
+   * - `'immediately'`: Queries are invalidated instantly after the mutation finished, even if they are running.
+   *   If the query provides tags that were invalidated while it ran, it won't be re-fetched.
+   * - `'delayed'`: Invalidation only happens after all queries and mutations are settled.
+   *   This ensures that queries are always invalidated correctly and automatically "batches" invalidations of concurrent mutations.
+   *   Note that if you constantly have some queries (or mutations) running, this can delay tag invalidations indefinitely.
+   */
+  invalidationBehavior?: 'delayed' | 'immediately'
   /**
    * A function that is passed every dispatched action. If this returns something other than `undefined`,
    * that return value will be used to rehydrate fulfilled & errored queries.
@@ -255,6 +265,7 @@ export function buildCreateApi<Modules extends [Module<any>, ...Module<any>[]]>(
       refetchOnMountOrArgChange: false,
       refetchOnFocus: false,
       refetchOnReconnect: false,
+      invalidationBehavior: 'delayed',
       ...options,
       extractRehydrationInfo,
       serializeQueryArgs(queryArgsApi) {

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

@@ -51,6 +51,7 @@ describe('buildSlice', () => {
       api: {
         config: {
           focused: true,
+          invalidationBehavior: 'delayed',
           keepUnusedDataFor: 60,
           middlewareRegistered: true,
           online: true,

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

@@ -0,0 +1,109 @@
+import { createApi, QueryStatus } from '@reduxjs/toolkit/query'
+import { getLog } from 'console-testing-library'
+import { actionsReducer, setupApiStore, waitMs } from './helpers'
+
+// We need to be able to control when which query resolves to simulate race
+// conditions properly, that's the purpose of this factory.
+const createPromiseFactory = () => {
+  const resolveQueue: (() => void)[] = []
+  const createPromise = () =>
+    new Promise<void>((resolve) => {
+      resolveQueue.push(resolve)
+    })
+  const resolveOldest = () => {
+    resolveQueue.shift()?.()
+  }
+  return { createPromise, resolveOldest }
+}
+
+const getEatenBananaPromises = createPromiseFactory()
+const eatBananaPromises = createPromiseFactory()
+
+let eatenBananas = 0
+const api = createApi({
+  invalidationBehavior: 'delayed',
+  baseQuery: () => undefined as any,
+  tagTypes: ['Banana'],
+  endpoints: (build) => ({
+    // Eat a banana.
+    eatBanana: build.mutation<unknown, void>({
+      queryFn: async () => {
+        await eatBananaPromises.createPromise()
+        eatenBananas += 1
+        return { data: null, meta: {} }
+      },
+      invalidatesTags: ['Banana'],
+    }),
+
+    // Get the number of eaten bananas.
+    getEatenBananas: build.query<number, void>({
+      queryFn: async (arg, arg1, arg2, arg3) => {
+        const result = eatenBananas
+        await getEatenBananaPromises.createPromise()
+        return { data: result }
+      },
+      providesTags: ['Banana'],
+    }),
+  }),
+})
+const { getEatenBananas, eatBanana } = api.endpoints
+
+const storeRef = setupApiStore(api, {
+  ...actionsReducer,
+})
+
+it('invalidates a query after a corresponding mutation', async () => {
+  eatenBananas = 0
+
+  const query = storeRef.store.dispatch(getEatenBananas.initiate())
+  const getQueryState = () =>
+    storeRef.store.getState().api.queries[query.queryCacheKey]
+  getEatenBananaPromises.resolveOldest()
+  await waitMs(2)
+
+  expect(getQueryState()?.data).toBe(0)
+  expect(getQueryState()?.status).toBe(QueryStatus.fulfilled)
+
+  const mutation = storeRef.store.dispatch(eatBanana.initiate())
+  const getMutationState = () =>
+    storeRef.store.getState().api.mutations[mutation.requestId]
+  eatBananaPromises.resolveOldest()
+  await waitMs(2)
+
+  expect(getMutationState()?.status).toBe(QueryStatus.fulfilled)
+  expect(getQueryState()?.data).toBe(0)
+  expect(getQueryState()?.status).toBe(QueryStatus.pending)
+
+  getEatenBananaPromises.resolveOldest()
+  await waitMs(2)
+
+  expect(getQueryState()?.data).toBe(1)
+  expect(getQueryState()?.status).toBe(QueryStatus.fulfilled)
+})
+
+it('invalidates a query whose corresponding mutation finished while the query was in flight', async () => {
+  eatenBananas = 0
+
+  const query = storeRef.store.dispatch(getEatenBananas.initiate())
+  const getQueryState = () =>
+    storeRef.store.getState().api.queries[query.queryCacheKey]
+
+  const mutation = storeRef.store.dispatch(eatBanana.initiate())
+  const getMutationState = () =>
+    storeRef.store.getState().api.mutations[mutation.requestId]
+  eatBananaPromises.resolveOldest()
+  await waitMs(2)
+  expect(getMutationState()?.status).toBe(QueryStatus.fulfilled)
+
+  getEatenBananaPromises.resolveOldest()
+  await waitMs(2)
+  expect(getQueryState()?.data).toBe(0)
+  expect(getQueryState()?.status).toBe(QueryStatus.pending)
+
+  // should already be refetching
+  getEatenBananaPromises.resolveOldest()
+  await waitMs(2)
+
+  expect(getQueryState()?.status).toBe(QueryStatus.fulfilled)
+  expect(getQueryState()?.data).toBe(1)
+})

packages/toolkit/src/tests/combineSlices.test.ts

@@ -33,6 +33,7 @@ const api = {
       subscriptions: {},
       config: {
         reducerPath: 'api',
+        invalidationBehavior: 'delayed',
         online: false,
         focused: false,
         keepUnusedDataFor: 60,