Rewrite providedTags handling for better perf (#4910)

* Restructure provided tags to improve deletion perf

* Accept array of provided tag entries internally for perf

Author: markerikson

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> }

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
       }

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,67 +481,66 @@ 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(
           draft,
-          action: PayloadAction<{
-            queryCacheKey: QueryCacheKey
-            providedTags: readonly FullTagDescription<string>[]
-          }>,
+          action: PayloadAction<
+            Array<{
+              queryCacheKey: QueryCacheKey
+              providedTags: readonly FullTagDescription<string>[]
+            }>
+          >,
         ) {
-          const { queryCacheKey, providedTags } = action.payload
+          for (const { queryCacheKey, providedTags } of action.payload) {
+            removeCacheKeyFromTags(draft, 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)
+            for (const { type, id } of providedTags) {
+              const subscribedQueries = ((draft.tags[type] ??= {})[
+                id || '__internal_without_id'
+              ] ??= [])
+              const alreadySubscribed =
+                subscribedQueries.includes(queryCacheKey)
+              if (!alreadySubscribed) {
+                subscribedQueries.push(queryCacheKey)
               }
             }
-          }
 
-          for (const { type, id } of providedTags) {
-            const subscribedQueries = ((draft[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
-          providedTags: readonly FullTagDescription<string>[]
-        }>(),
+        prepare:
+          prepareAutoBatched<
+            Array<{
+              queryCacheKey: QueryCacheKey
+              providedTags: readonly FullTagDescription<string>[]
+            }>
+          >(),
       },
     },
     extraReducers(builder) {
       builder
         .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) {
@@ -564,48 +556,71 @@ export function buildSlice({
         .addMatcher(
           isAnyOf(isFulfilled(queryThunk), isRejectedWithValue(queryThunk)),
           (draft, action) => {
-            writeProvidedTagsForQuery(draft, action)
+            writeProvidedTagsForQueries(draft, [action])
           },
         )
         .addMatcher(
           querySlice.actions.cacheEntriesUpserted.match,
           (draft, action) => {
-            for (const { queryDescription: arg, value } of action.payload) {
-              const action: CalculateProvidedByAction = {
-                type: 'UNKNOWN',
-                payload: value,
-                meta: {
-                  requestStatus: 'fulfilled',
-                  requestId: 'UNKNOWN',
-                  arg,
-                },
-              }
-
-              writeProvidedTagsForQuery(draft, action)
-            }
+            const mockActions: CalculateProvidedByAction[] = action.payload.map(
+              ({ queryDescription, value }) => {
+                return {
+                  type: 'UNKNOWN',
+                  payload: value,
+                  meta: {
+                    requestStatus: 'fulfilled',
+                    requestId: 'UNKNOWN',
+                    arg: queryDescription,
+                  },
+                }
+              },
+            )
+            writeProvidedTagsForQueries(draft, mockActions)
           },
         )
     },
   })
 
-  function writeProvidedTagsForQuery(
+  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 writeProvidedTagsForQueries(
     draft: InvalidationState<string>,
-    action: CalculateProvidedByAction,
+    actions: CalculateProvidedByAction[],
   ) {
-    const providedTags = calculateProvidedByThunk(
-      action,
-      'providesTags',
-      definitions,
-      assertTagType,
-    )
-    const { queryCacheKey } = action.meta.arg
+    const providedByEntries = actions.map((action) => {
+      const providedTags = calculateProvidedByThunk(
+        action,
+        'providesTags',
+        definitions,
+        assertTagType,
+      )
+      const { queryCacheKey } = action.meta.arg
+      return { queryCacheKey, providedTags }
+    })
 
     invalidationSlice.caseReducers.updateProvidedBy(
       draft,
-      invalidationSlice.actions.updateProvidedBy({
-        queryCacheKey,
-        providedTags,
-      }),
+      invalidationSlice.actions.updateProvidedBy(providedByEntries),
     )
   }