Handle batched upserts in cache lifecycles

Author: markerikson

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

@@ -44,12 +44,14 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
   context,
   internalState,
 }) => {
-  const { removeQueryResult, unsubscribeQueryResult } = api.internalActions
+  const { removeQueryResult, unsubscribeQueryResult, cacheEntriesUpserted } =
+    api.internalActions
 
   const canTriggerUnsubscribe = isAnyOf(
     unsubscribeQueryResult.match,
     queryThunk.fulfilled,
     queryThunk.rejected,
+    cacheEntriesUpserted.match,
   )
 
   function anySubscriptionsRemainingForKey(queryCacheKey: string) {
@@ -66,16 +68,27 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
   ) => {
     if (canTriggerUnsubscribe(action)) {
       const state = mwApi.getState()[reducerPath]
-      const { queryCacheKey } = unsubscribeQueryResult.match(action)
-        ? action.payload
-        : action.meta.arg
-
-      handleUnsubscribe(
-        queryCacheKey,
-        state.queries[queryCacheKey]?.endpointName,
-        mwApi,
-        state.config,
-      )
+      let queryCacheKeys: QueryCacheKey[]
+
+      if (cacheEntriesUpserted.match(action)) {
+        queryCacheKeys = action.payload.map(
+          (entry) => entry.queryDescription.queryCacheKey,
+        )
+      } else {
+        const { queryCacheKey } = unsubscribeQueryResult.match(action)
+          ? action.payload
+          : action.meta.arg
+        queryCacheKeys = [queryCacheKey]
+      }
+
+      for (const queryCacheKey of queryCacheKeys) {
+        handleUnsubscribe(
+          queryCacheKey,
+          state.queries[queryCacheKey]?.endpointName,
+          mwApi,
+          state.config,
+        )
+      }
     }
 
     if (api.util.resetApiState.match(action)) {
@@ -132,6 +145,7 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
       if (currentTimeout) {
         clearTimeout(currentTimeout)
       }
+
       currentRemovalTimeouts[queryCacheKey] = setTimeout(() => {
         if (!anySubscriptionsRemainingForKey(queryCacheKey)) {
           api.dispatch(removeQueryResult({ queryCacheKey }))

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

@@ -183,24 +183,68 @@ export const buildCacheLifecycleHandler: InternalHandlerBuilder = ({
   }
   const lifecycleMap: Record<string, CacheLifecycle> = {}
 
+  function resolveLifecycleEntry(
+    cacheKey: string,
+    data: unknown,
+    meta: unknown,
+  ) {
+    const lifecycle = lifecycleMap[cacheKey]
+
+    if (lifecycle?.valueResolved) {
+      lifecycle.valueResolved({
+        data,
+        meta,
+      })
+      delete lifecycle.valueResolved
+    }
+  }
+
+  function removeLifecycleEntry(cacheKey: string) {
+    const lifecycle = lifecycleMap[cacheKey]
+    if (lifecycle) {
+      delete lifecycleMap[cacheKey]
+      lifecycle.cacheEntryRemoved()
+    }
+  }
+
   const handler: ApiMiddlewareInternalHandler = (
     action,
     mwApi,
     stateBefore,
   ) => {
     const cacheKey = getCacheKey(action)
 
-    if (queryThunk.pending.match(action)) {
+    function checkForNewCacheKey(
+      endpointName: string,
+      cacheKey: string,
+      requestId: string,
+      originalArgs: unknown,
+    ) {
       const oldState = stateBefore[reducerPath].queries[cacheKey]
       const state = mwApi.getState()[reducerPath].queries[cacheKey]
       if (!oldState && state) {
-        handleNewKey(
-          action.meta.arg.endpointName,
-          action.meta.arg.originalArgs,
-          cacheKey,
-          mwApi,
+        handleNewKey(endpointName, originalArgs, cacheKey, mwApi, requestId)
+      }
+    }
+
+    if (queryThunk.pending.match(action)) {
+      checkForNewCacheKey(
+        action.meta.arg.endpointName,
+        cacheKey,
+        action.meta.requestId,
+        action.meta.arg.originalArgs,
+      )
+    } else if (api.internalActions.cacheEntriesUpserted.match(action)) {
+      for (const { queryDescription, value } of action.payload) {
+        const { endpointName, originalArgs, queryCacheKey } = queryDescription
+        checkForNewCacheKey(
+          endpointName,
+          queryCacheKey,
           action.meta.requestId,
+          originalArgs,
         )
+
+        resolveLifecycleEntry(queryCacheKey, value, {})
       }
     } else if (mutationThunk.pending.match(action)) {
       const state = mwApi.getState()[reducerPath].mutations[cacheKey]
@@ -214,27 +258,15 @@ export const buildCacheLifecycleHandler: InternalHandlerBuilder = ({
         )
       }
     } else if (isFulfilledThunk(action)) {
-      const lifecycle = lifecycleMap[cacheKey]
-      if (lifecycle?.valueResolved) {
-        lifecycle.valueResolved({
-          data: action.payload,
-          meta: action.meta.baseQueryMeta,
-        })
-        delete lifecycle.valueResolved
-      }
+      resolveLifecycleEntry(cacheKey, action.payload, action.meta.baseQueryMeta)
     } else if (
       api.internalActions.removeQueryResult.match(action) ||
       api.internalActions.removeMutationResult.match(action)
     ) {
-      const lifecycle = lifecycleMap[cacheKey]
-      if (lifecycle) {
-        delete lifecycleMap[cacheKey]
-        lifecycle.cacheEntryRemoved()
-      }
+      removeLifecycleEntry(cacheKey)
     } else if (api.util.resetApiState.match(action)) {
-      for (const [cacheKey, lifecycle] of Object.entries(lifecycleMap)) {
-        delete lifecycleMap[cacheKey]
-        lifecycle.cacheEntryRemoved()
+      for (const cacheKey of Object.keys(lifecycleMap)) {
+        removeLifecycleEntry(cacheKey)
       }
     }
   }

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

@@ -76,6 +76,11 @@ type NormalizedQueryUpsertEntryPayload = {
   value: any
 }
 
+export type ProcessedQueryUpsertEntry = {
+  queryDescription: QueryThunkArg
+  value: unknown
+}
+
 /**
  * A typesafe representation of a util action creator that accepts cache entry descriptions to upsert
  */
@@ -90,7 +95,7 @@ export type UpsertEntries<Definitions extends EndpointDefinitions> = <
       >
     },
   ],
-) => PayloadAction<NormalizedQueryUpsertEntryPayload>
+) => PayloadAction<NormalizedQueryUpsertEntryPayload[]>
 
 function updateQuerySubstateIfExists(
   state: QueryState<any>,
@@ -276,7 +281,7 @@ export function buildSlice({
         reducer(
           draft,
           action: PayloadAction<
-            NormalizedQueryUpsertEntryPayload[],
+            ProcessedQueryUpsertEntry[],
             string,
             {
               RTK_autoBatch: boolean
@@ -286,19 +291,7 @@ export function buildSlice({
           >,
         ) {
           for (const entry of action.payload) {
-            const { endpointName, args, value } = entry
-            const endpointDefinition = definitions[endpointName]
-
-            const arg: QueryThunkArg = {
-              type: 'query',
-              endpointName: endpointName,
-              originalArgs: entry.args,
-              queryCacheKey: serializeQueryArgs({
-                queryArgs: args,
-                endpointDefinition,
-                endpointName,
-              }),
-            }
+            const { queryDescription: arg, value } = entry
             writePendingCacheEntry(draft, arg, true, {
               arg,
               requestId: action.meta.requestId,
@@ -313,13 +306,31 @@ export function buildSlice({
                 fulfilledTimeStamp: action.meta.timestamp,
                 baseQueryMeta: {},
               },
-              entry.value,
+              value,
             )
           }
         },
         prepare: (payload: NormalizedQueryUpsertEntryPayload[]) => {
+          const queryDescriptions: ProcessedQueryUpsertEntry[] = payload.map(
+            (entry) => {
+              const { endpointName, args, value } = entry
+              const endpointDefinition = definitions[endpointName]
+              const queryDescription: QueryThunkArg = {
+                type: 'query',
+                endpointName: endpointName,
+                originalArgs: entry.args,
+                queryCacheKey: serializeQueryArgs({
+                  queryArgs: args,
+                  endpointDefinition,
+                  endpointName,
+                }),
+              }
+              return { queryDescription, value }
+            },
+          )
+
           const result = {
-            payload,
+            payload: queryDescriptions,
             meta: {
               [SHOULD_AUTOBATCH]: true,
               requestId: nanoid(),