Restructure provided tags to improve deletion perf

markerikson · markerikson · commit 25b4042a7d53 · 2025-03-27T15:52:31.000-04:00
diff --git a/packages/toolkit/src/query/core/apiState.ts b/packages/toolkit/src/query/core/apiState.ts
@@ -3,6 +3,7 @@ import type { BaseQueryError } from '../baseQueryTypes'
 import type {
   BaseEndpointDefinition,
   EndpointDefinitions,
+  FullTagDescription,
   InfiniteQueryDefinition,
   MutationDefinition,
   PageParamFrom,
@@ -15,14 +16,8 @@ import type { Id, WithRequiredProp } from '../tsHelpers'
 export type QueryCacheKey = string & { _type: 'queryCacheKey' }
 export type QuerySubstateIdentifier = { queryCacheKey: QueryCacheKey }
 export type MutationSubstateIdentifier =
-  | {
-      requestId: string
-      fixedCacheKey?: string
-    }
-  | {
-      requestId?: string
-      fixedCacheKey: string
-    }
+  | { requestId: string; fixedCacheKey?: string }
+  | { requestId?: string; fixedCacheKey: string }
 
 export type RefetchConfigOptions = {
   refetchOnMountOrArgChange: boolean | number
@@ -225,18 +220,15 @@ export type QuerySubState<
   D extends BaseEndpointDefinition<any, any, any>,
   DataType = ResultTypeFrom<D>,
 > = Id<
-  | ({
-      status: QueryStatus.fulfilled
-    } & WithRequiredProp<
+  | ({ status: QueryStatus.fulfilled } & WithRequiredProp<
       BaseQuerySubState<D, DataType>,
       'data' | 'fulfilledTimeStamp'
     > & { error: undefined })
-  | ({
-      status: QueryStatus.pending
-    } & BaseQuerySubState<D, DataType>)
-  | ({
-      status: QueryStatus.rejected
-    } & WithRequiredProp<BaseQuerySubState<D, DataType>, 'error'>)
+  | ({ status: QueryStatus.pending } & BaseQuerySubState<D, DataType>)
+  | ({ status: QueryStatus.rejected } & WithRequiredProp<
+      BaseQuerySubState<D, DataType>,
+      'error'
+    >)
   | {
       status: QueryStatus.uninitialized
       originalArgs?: undefined
@@ -274,18 +266,17 @@ type BaseMutationSubState<D extends BaseEndpointDefinition<any, any, any>> = {
 }
 
 export type MutationSubState<D extends BaseEndpointDefinition<any, any, any>> =
-  | (({
-      status: QueryStatus.fulfilled
-    } & WithRequiredProp<
+  | (({ status: QueryStatus.fulfilled } & WithRequiredProp<
       BaseMutationSubState<D>,
       'data' | 'fulfilledTimeStamp'
     >) & { error: undefined })
-  | (({
-      status: QueryStatus.pending
-    } & BaseMutationSubState<D>) & { data?: undefined })
-  | ({
-      status: QueryStatus.rejected
-    } & WithRequiredProp<BaseMutationSubState<D>, 'error'>)
+  | (({ status: QueryStatus.pending } & BaseMutationSubState<D>) & {
+      data?: undefined
+    })
+  | ({ status: QueryStatus.rejected } & WithRequiredProp<
+      BaseMutationSubState<D>,
+      'error'
+    >)
   | {
       requestId?: undefined
       status: QueryStatus.uninitialized
@@ -309,10 +300,13 @@ export type CombinedState<
 }
 
 export type InvalidationState<TagTypes extends string> = {
-  [_ in TagTypes]: {
-    [id: string]: Array<QueryCacheKey>
-    [id: number]: Array<QueryCacheKey>
+  tags: {
+    [_ in TagTypes]: {
+      [id: string]: Array<QueryCacheKey>
+      [id: number]: Array<QueryCacheKey>
+    }
   }
+  keys: Record<QueryCacheKey, Array<FullTagDescription<any>>>
 }
 
 export type QueryState<D extends EndpointDefinitions> = {
@@ -346,6 +340,4 @@ export type RootState<
   Definitions extends EndpointDefinitions,
   TagTypes extends string,
   ReducerPath extends string,
-> = {
-  [P in ReducerPath]: CombinedState<Definitions, TagTypes, P>
-}
+> = { [P in ReducerPath]: CombinedState<Definitions, TagTypes, P> }
diff --git a/packages/toolkit/src/query/core/buildSelectors.ts b/packages/toolkit/src/query/core/buildSelectors.ts
@@ -199,10 +199,7 @@ export function buildSelectors<
   function withRequestFlags<T extends { status: QueryStatus }>(
     substate: T,
   ): T & RequestStatusFlags {
-    return {
-      ...substate,
-      ...getRequestStatusFlags(substate.status),
-    }
+    return { ...substate, ...getRequestStatusFlags(substate.status) }
   }
 
   function selectApiState(rootState: RootState) {
@@ -344,7 +341,7 @@ export function buildSelectors<
     const apiState = state[reducerPath]
     const toInvalidate = new Set<QueryCacheKey>()
     for (const tag of tags.filter(isNotNullish).map(expandTagDescription)) {
-      const provided = apiState.provided[tag.type]
+      const provided = apiState.provided.tags[tag.type]
       if (!provided) {
         continue
       }
diff --git a/packages/toolkit/src/query/core/buildSlice.ts b/packages/toolkit/src/query/core/buildSlice.ts
@@ -217,10 +217,7 @@ export function buildSlice({
 
   function writeFulfilledCacheEntry(
     draft: QueryState<any>,
-    meta: {
-      arg: QueryThunkArg
-      requestId: string
-    } & {
+    meta: { arg: QueryThunkArg; requestId: string } & {
       fulfilledTimeStamp: number
       baseQueryMeta: unknown
     },
@@ -297,11 +294,7 @@ export function buildSlice({
           action: PayloadAction<
             ProcessedQueryUpsertEntry[],
             string,
-            {
-              RTK_autoBatch: boolean
-              requestId: string
-              timestamp: number
-            }
+            { RTK_autoBatch: boolean; requestId: string; timestamp: number }
           >,
         ) {
           for (const entry of action.payload) {
@@ -488,9 +481,14 @@ export function buildSlice({
     | ReturnType<ReturnType<InfiniteQueryThunk<any>>>
   >
 
+  const initialInvalidationState: InvalidationState<string> = {
+    tags: {},
+    keys: {},
+  }
+
   const invalidationSlice = createSlice({
     name: `${reducerPath}/invalidation`,
-    initialState: initialState as InvalidationState<string>,
+    initialState: initialInvalidationState,
     reducers: {
       updateProvidedBy: {
         reducer(
@@ -502,24 +500,21 @@ export function buildSlice({
         ) {
           const { queryCacheKey, providedTags } = action.payload
 
-          for (const tagTypeSubscriptions of Object.values(draft)) {
-            for (const idSubscriptions of Object.values(tagTypeSubscriptions)) {
-              const foundAt = idSubscriptions.indexOf(queryCacheKey)
-              if (foundAt !== -1) {
-                idSubscriptions.splice(foundAt, 1)
-              }
-            }
-          }
+          removeCacheKeyFromTags(draft, queryCacheKey)
 
           for (const { type, id } of providedTags) {
-            const subscribedQueries = ((draft[type] ??= {})[
+            const subscribedQueries = ((draft.tags[type] ??= {})[
               id || '__internal_without_id'
             ] ??= [])
             const alreadySubscribed = subscribedQueries.includes(queryCacheKey)
             if (!alreadySubscribed) {
               subscribedQueries.push(queryCacheKey)
             }
           }
+
+          // Remove readonly from the providedTags array
+          draft.keys[queryCacheKey] =
+            providedTags as FullTagDescription<string>[]
         },
         prepare: prepareAutoBatched<{
           queryCacheKey: QueryCacheKey
@@ -532,23 +527,14 @@ export function buildSlice({
         .addCase(
           querySlice.actions.removeQueryResult,
           (draft, { payload: { queryCacheKey } }) => {
-            for (const tagTypeSubscriptions of Object.values(draft)) {
-              for (const idSubscriptions of Object.values(
-                tagTypeSubscriptions,
-              )) {
-                const foundAt = idSubscriptions.indexOf(queryCacheKey)
-                if (foundAt !== -1) {
-                  idSubscriptions.splice(foundAt, 1)
-                }
-              }
-            }
+            removeCacheKeyFromTags(draft, queryCacheKey)
           },
         )
         .addMatcher(hasRehydrationInfo, (draft, action) => {
           const { provided } = extractRehydrationInfo(action)!
           for (const [type, incomingTags] of Object.entries(provided)) {
             for (const [id, cacheKeys] of Object.entries(incomingTags)) {
-              const subscribedQueries = ((draft[type] ??= {})[
+              const subscribedQueries = ((draft.tags[type] ??= {})[
                 id || '__internal_without_id'
               ] ??= [])
               for (const queryCacheKey of cacheKeys) {
@@ -574,11 +560,7 @@ export function buildSlice({
               const action: CalculateProvidedByAction = {
                 type: 'UNKNOWN',
                 payload: value,
-                meta: {
-                  requestStatus: 'fulfilled',
-                  requestId: 'UNKNOWN',
-                  arg,
-                },
+                meta: { requestStatus: 'fulfilled', requestId: 'UNKNOWN', arg },
               }
 
               writeProvidedTagsForQuery(draft, action)
@@ -588,6 +570,28 @@ export function buildSlice({
     },
   })
 
+  function removeCacheKeyFromTags(
+    draft: InvalidationState<any>,
+    queryCacheKey: QueryCacheKey,
+  ) {
+    const existingTags = draft.keys[queryCacheKey] ?? []
+
+    // Delete this cache key from any existing tags that may have provided it
+    for (const tag of existingTags) {
+      const tagType = tag.type
+      const tagId = tag.id ?? '__internal_without_id'
+      const tagSubscriptions = draft.tags[tagType]?.[tagId]
+
+      if (tagSubscriptions) {
+        draft.tags[tagType][tagId] = tagSubscriptions.filter(
+          (qc) => qc !== queryCacheKey,
+        )
+      }
+    }
+
+    delete draft.keys[queryCacheKey]
+  }
+
   function writeProvidedTagsForQuery(
     draft: InvalidationState<string>,
     action: CalculateProvidedByAction,
diff --git a/packages/toolkit/src/query/tests/optimisticUpdates.test.tsx b/packages/toolkit/src/query/tests/optimisticUpdates.test.tsx
@@ -60,9 +60,7 @@ const api = createApi({
   }),
 })
 
-const storeRef = setupApiStore(api, {
-  ...actionsReducer,
-})
+const storeRef = setupApiStore(api, { ...actionsReducer })
 
 describe('basic lifecycle', () => {
   let onStart = vi.fn(),
@@ -96,9 +94,7 @@ describe('basic lifecycle', () => {
   test('success', async () => {
     const { result } = renderHook(
       () => extendedApi.endpoints.test.useMutation(),
-      {
-        wrapper: storeRef.wrapper,
-      },
+      { wrapper: storeRef.wrapper },
     )
 
     baseQuery.mockResolvedValue('success')
@@ -119,9 +115,7 @@ describe('basic lifecycle', () => {
   test('error', async () => {
     const { result } = renderHook(
       () => extendedApi.endpoints.test.useMutation(),
-      {
-        wrapper: storeRef.wrapper,
-      },
+      { wrapper: storeRef.wrapper },
     )
 
     baseQuery.mockRejectedValueOnce('error')
@@ -201,11 +195,7 @@ describe('updateQueryData', () => {
   test('updates (list) cache values including provided tags, undos that', async () => {
     baseQuery
       .mockResolvedValueOnce([
-        {
-          id: '3',
-          title: 'All about cheese.',
-          contents: 'TODO',
-        },
+        { id: '3', title: 'All about cheese.', contents: 'TODO' },
       ])
       .mockResolvedValueOnce(42)
     const { result } = renderHook(() => api.endpoints.listPosts.useQuery(), {
@@ -218,7 +208,7 @@ describe('updateQueryData', () => {
       provided = storeRef.store.getState().api.provided
     })
 
-    const provided3 = provided.Post['3']
+    const provided3 = provided.tags.Post['3']
 
     let returnValue!: ReturnType<ReturnType<typeof api.util.updateQueryData>>
     act(() => {
@@ -242,7 +232,7 @@ describe('updateQueryData', () => {
       provided = storeRef.store.getState().api.provided
     })
 
-    const provided4 = provided.Post['4']
+    const provided4 = provided.tags.Post['4']
 
     expect(provided4).toEqual(provided3)
 
@@ -254,19 +244,15 @@ describe('updateQueryData', () => {
       provided = storeRef.store.getState().api.provided
     })
 
-    const provided4Next = provided.Post['4']
+    const provided4Next = provided.tags.Post['4']
 
     expect(provided4Next).toEqual([])
   })
 
   test('updates (list) cache values excluding provided tags, undoes that', async () => {
     baseQuery
       .mockResolvedValueOnce([
-        {
-          id: '3',
-          title: 'All about cheese.',
-          contents: 'TODO',
-        },
+        { id: '3', title: 'All about cheese.', contents: 'TODO' },
       ])
       .mockResolvedValueOnce(42)
     const { result } = renderHook(() => api.endpoints.listPosts.useQuery(), {
@@ -301,7 +287,7 @@ describe('updateQueryData', () => {
       provided = storeRef.store.getState().api.provided
     })
 
-    const provided4 = provided.Post['4']
+    const provided4 = provided.tags.Post['4']
 
     expect(provided4).toEqual(undefined)
 
@@ -313,7 +299,7 @@ describe('updateQueryData', () => {
       provided = storeRef.store.getState().api.provided
     })
 
-    const provided4Next = provided.Post['4']
+    const provided4Next = provided.tags.Post['4']
 
     expect(provided4Next).toEqual(undefined)
   })
@@ -382,9 +368,7 @@ describe('full integration', () => {
         query: api.endpoints.post.useQuery('3'),
         mutation: api.endpoints.updatePost.useMutation(),
       }),
-      {
-        wrapper: storeRef.wrapper,
-      },
+      { wrapper: storeRef.wrapper },
     )
     await hookWaitFor(() => expect(result.current.query.isSuccess).toBeTruthy())
 
@@ -433,9 +417,7 @@ describe('full integration', () => {
         query: api.endpoints.post.useQuery('3'),
         mutation: api.endpoints.updatePost.useMutation(),
       }),
-      {
-        wrapper: storeRef.wrapper,
-      },
+      { wrapper: storeRef.wrapper },
     )
     await hookWaitFor(() => expect(result.current.query.isSuccess).toBeTruthy())
 
diff --git a/packages/toolkit/src/query/tests/optimisticUpserts.test.tsx b/packages/toolkit/src/query/tests/optimisticUpserts.test.tsx
