Byte-shave cache middleware

Author: markerikson

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

@@ -43,6 +43,7 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
   queryThunk,
   context,
   internalState,
+  selectors: { selectQueryEntry, selectConfig },
 }) => {
   const { removeQueryResult, unsubscribeQueryResult, cacheEntriesUpserted } =
     api.internalActions
@@ -66,8 +67,9 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
     mwApi,
     internalState,
   ) => {
+    const state = mwApi.getState()
+    const config = selectConfig(state)
     if (canTriggerUnsubscribe(action)) {
-      const state = mwApi.getState()[reducerPath]
       let queryCacheKeys: QueryCacheKey[]
 
       if (cacheEntriesUpserted.match(action)) {
@@ -81,14 +83,7 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
         queryCacheKeys = [queryCacheKey]
       }
 
-      for (const queryCacheKey of queryCacheKeys) {
-        handleUnsubscribe(
-          queryCacheKey,
-          state.queries[queryCacheKey]?.endpointName,
-          mwApi,
-          state.config,
-        )
-      }
+      handleUnsubscribeMany(queryCacheKeys, mwApi, config)
     }
 
     if (api.util.resetApiState.match(action)) {
@@ -99,19 +94,27 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
     }
 
     if (context.hasRehydrationInfo(action)) {
-      const state = mwApi.getState()[reducerPath]
       const { queries } = context.extractRehydrationInfo(action)!
-      for (const [queryCacheKey, queryState] of Object.entries(queries)) {
-        // Gotcha:
-        // If rehydrating before the endpoint has been injected,the global `keepUnusedDataFor`
-        // will be used instead of the endpoint-specific one.
-        handleUnsubscribe(
-          queryCacheKey as QueryCacheKey,
-          queryState?.endpointName,
-          mwApi,
-          state.config,
-        )
-      }
+      // Gotcha:
+      // If rehydrating before the endpoint has been injected,the global `keepUnusedDataFor`
+      // will be used instead of the endpoint-specific one.
+      handleUnsubscribeMany(
+        Object.keys(queries) as QueryCacheKey[],
+        mwApi,
+        config,
+      )
+    }
+  }
+
+  function handleUnsubscribeMany(
+    cacheKeys: QueryCacheKey[],
+    api: SubMiddlewareApi,
+    config: ConfigState<string>,
+  ) {
+    const state = api.getState()
+    for (const queryCacheKey of cacheKeys) {
+      const entry = selectQueryEntry(state, queryCacheKey)
+      handleUnsubscribe(queryCacheKey, entry?.endpointName, api, config)
     }
   }
 

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

@@ -2,7 +2,7 @@ import type { ThunkDispatch, UnknownAction } from '@reduxjs/toolkit'
 import type { BaseQueryFn, BaseQueryMeta } from '../../baseQueryTypes'
 import type { BaseEndpointDefinition } from '../../endpointDefinitions'
 import { DefinitionType } from '../../endpointDefinitions'
-import type { RootState } from '../apiState'
+import type { QueryCacheKey, RootState } from '../apiState'
 import type {
   MutationResultSelectorResult,
   QueryResultSelectorResult,
@@ -185,6 +185,7 @@ export const buildCacheLifecycleHandler: InternalHandlerBuilder = ({
   queryThunk,
   mutationThunk,
   internalState,
+  selectors: { selectQueryEntry, selectApiState },
 }) => {
   const isQueryThunk = isAsyncThunkAction(queryThunk)
   const isMutationThunk = isAsyncThunkAction(mutationThunk)
@@ -225,17 +226,17 @@ export const buildCacheLifecycleHandler: InternalHandlerBuilder = ({
     mwApi,
     stateBefore,
   ) => {
-    const cacheKey = getCacheKey(action)
+    const cacheKey = getCacheKey(action) as QueryCacheKey
 
     function checkForNewCacheKey(
       endpointName: string,
-      cacheKey: string,
+      cacheKey: QueryCacheKey,
       requestId: string,
       originalArgs: unknown,
     ) {
-      const oldState = stateBefore[reducerPath].queries[cacheKey]
-      const state = mwApi.getState()[reducerPath].queries[cacheKey]
-      if (!oldState && state) {
+      const oldEntry = selectQueryEntry(stateBefore, cacheKey)
+      const newEntry = selectQueryEntry(mwApi.getState(), cacheKey)
+      if (!oldEntry && newEntry) {
         handleNewKey(endpointName, originalArgs, cacheKey, mwApi, requestId)
       }
     }

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

@@ -76,11 +76,11 @@ export const buildInvalidationByTagsHandler: InternalHandlerBuilder = ({
   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
+    const { queries, mutations } = state
+    for (const cacheRecord of [queries, mutations]) {
+      for (const key in cacheRecord) {
+        if (cacheRecord[key]?.status === QueryStatus.pending) return true
+      }
     }
     return false
   }

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

@@ -1,4 +1,8 @@
-import type { QuerySubstateIdentifier, Subscribers } from '../apiState'
+import type {
+  QueryCacheKey,
+  QuerySubstateIdentifier,
+  Subscribers,
+} from '../apiState'
 import { QueryStatus } from '../apiState'
 import type {
   QueryStateMeta,
@@ -49,6 +53,20 @@ export const buildPollingHandler: InternalHandlerBuilder = ({
     }
   }
 
+  function getCacheEntrySubscriptions(
+    queryCacheKey: QueryCacheKey,
+    api: SubMiddlewareApi,
+  ) {
+    const state = api.getState()[reducerPath]
+    const querySubState = state.queries[queryCacheKey]
+    const subscriptions = internalState.currentSubscriptions[queryCacheKey]
+
+    if (!querySubState || querySubState.status === QueryStatus.uninitialized)
+      return
+
+    return subscriptions
+  }
+
   function startNextPoll(
     { queryCacheKey }: QuerySubstateIdentifier,
     api: SubMiddlewareApi,