feat(rtk-query): unsubscribe prefetched queries after configurable timer #1283

Closes #1283

Automatically removes prefetch subscriptions after configurable amount of time.

Description:

Prefetch subscription are now removed after  `prefetchOptions.keepSubscriptionFor` if provided
or api.config.keepPrefetchSubscriptionsFor otherwise.

Api changes:

- adds `keepSubscriptionFor` to prefetchOptions
- adds `keepPrefetchSubscriptionsFor` to api.config (default 10s)

Internal changes:

- prefetch queries now have the same requestId and the same subscription key

Author: FaberVitale

examples/query/react/prefetching/src/features/posts/PostsManager.tsx

@@ -30,7 +30,7 @@ const getColorForStatus = (status: Post['status']) => {
 const PostList = () => {
   const [page, setPage] = useState(1)
   const { data: posts, isLoading, isFetching } = useListPostsQuery(page)
-  const prefetchPage = usePrefetch('listPosts')
+  const prefetchPage = usePrefetch('listPosts', { keepSubscriptionFor: 5 })
 
   const prefetchNext = useCallback(() => {
     prefetchPage(page + 1)

packages/toolkit/src/query/apiTypes.ts

@@ -42,6 +42,7 @@ export type Module<Name extends ModuleName> = {
       | 'reducerPath'
       | 'serializeQueryArgs'
       | 'keepUnusedDataFor'
+      | 'keepPrefetchSubscriptionsFor'
       | 'refetchOnMountOrArgChange'
       | 'refetchOnFocus'
       | 'refetchOnReconnect'

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

@@ -253,7 +253,8 @@ export type ConfigState<ReducerPath> = RefetchConfigOptions & {
 } & ModifiableConfigState
 
 export type ModifiableConfigState = {
-  keepUnusedDataFor: number
+  keepUnusedDataFor: number,
+  keepPrefetchSubscriptionsFor: number
 } & RefetchConfigOptions
 
 export type MutationState<D extends EndpointDefinitions> = {

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

@@ -33,10 +33,15 @@ declare module './module' {
   }
 }
 
+export interface PrefetchSubscribriptionOptions {
+  keepSubscriptionFor?: number;
+}
+
 export interface StartQueryActionCreatorOptions {
   subscribe?: boolean
   forceRefetch?: boolean | number
-  subscriptionOptions?: SubscriptionOptions
+  subscriptionOptions?: SubscriptionOptions,
+  prefetch?: boolean | PrefetchSubscribriptionOptions,
 }
 
 type StartQueryActionCreator<
@@ -258,7 +263,7 @@ Features like automatic cache collection, automatic refetching etc. will not be
     endpointDefinition: QueryDefinition<any, any, any, any>
   ) {
     const queryAction: StartQueryActionCreator<any> =
-      (arg, { subscribe = true, forceRefetch, subscriptionOptions } = {}) =>
+      (arg, { subscribe = true, forceRefetch, subscriptionOptions, prefetch } = {}) =>
       (dispatch, getState) => {
         const queryCacheKey = serializeQueryArgs({
           queryArgs: arg,
@@ -269,12 +274,15 @@ Features like automatic cache collection, automatic refetching etc. will not be
           type: 'query',
           subscribe,
           forceRefetch,
+          prefetch,
           subscriptionOptions,
           endpointName,
           originalArgs: arg,
           queryCacheKey,
+          reducerPath: api.reducerPath,
         })
         const thunkResult = dispatch(thunk)
+
         middlewareWarning(getState)
 
         const { requestId, abort } = thunkResult

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

@@ -2,6 +2,7 @@ import type { BaseQueryFn } from '../../baseQueryTypes'
 import type { QueryDefinition } from '../../endpointDefinitions'
 import type { ConfigState, QueryCacheKey } from '../apiState'
 import { QuerySubstateIdentifier } from '../apiState'
+import type { PrefetchSubscribriptionOptions } from '../buildInitiate'
 import type {
   QueryStateMeta,
   SubMiddlewareApi,
@@ -28,11 +29,35 @@ declare module '../../endpointDefinitions' {
   }
 }
 
-export const build: SubMiddlewareBuilder = ({ reducerPath, api, context }) => {
+/**
+ * Output is in *milliseconds*.
+ */
+const getPrefetchSubscriptionTTLMs = (
+  prefetch: true | PrefetchSubscribriptionOptions,
+  config: ConfigState<string>
+): number => {
+  if (
+    typeof prefetch === 'object' &&
+    prefetch !== null &&
+    typeof prefetch.keepSubscriptionFor === 'number'
+  ) {
+    return prefetch.keepSubscriptionFor * 1000
+  }
+
+  return config.keepPrefetchSubscriptionsFor * 1000
+}
+
+export const build: SubMiddlewareBuilder = ({
+  reducerPath,
+  api,
+  context,
+  queryThunk,
+}) => {
   const { removeQueryResult, unsubscribeQueryResult } = api.internalActions
 
   return (mwApi) => {
     const currentRemovalTimeouts: QueryStateMeta<TimeoutId> = {}
+    const autoUnsubscribeTimeouts: QueryStateMeta<TimeoutId> = {}
 
     return (next) =>
       (action): any => {
@@ -50,8 +75,31 @@ export const build: SubMiddlewareBuilder = ({ reducerPath, api, context }) => {
           )
         }
 
+        if (queryThunk.pending.match(action) && action.meta.arg.prefetch) {
+          const requestId = action.meta.requestId
+          const currentTimeout = autoUnsubscribeTimeouts[requestId]
+
+          if (currentTimeout) {
+            clearTimeout(currentTimeout)
+          }
+
+          autoUnsubscribeTimeouts[requestId] = setTimeout(
+            mwApi.dispatch,
+            getPrefetchSubscriptionTTLMs(
+              action.meta.arg.prefetch,
+              mwApi.getState()[reducerPath].config
+            ),
+            unsubscribeQueryResult({
+              requestId,
+              queryCacheKey: action.meta.arg.queryCacheKey,
+            })
+          )
+        }
+
         if (api.util.resetApiState.match(action)) {
-          for (const [key, timeout] of Object.entries(currentRemovalTimeouts)) {
+          for (const [key, timeout] of Object.entries(
+            currentRemovalTimeouts
+          ).concat(Object.entries(autoUnsubscribeTimeouts))) {
             if (timeout) clearTimeout(timeout)
             delete currentRemovalTimeouts[key]
           }

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

@@ -80,6 +80,7 @@ export function buildMiddleware<
       originalArgs: querySubState.originalArgs,
       subscribe: false,
       forceRefetch: true,
+      reducerPath: reducerPath,
       queryCacheKey: queryCacheKey as any,
       ...override,
     })

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

@@ -18,13 +18,14 @@ import type {
   ResultTypeFrom,
 } from '../endpointDefinitions'
 import { calculateProvidedBy } from '../endpointDefinitions'
-import type { AsyncThunkPayloadCreator, Draft } from '@reduxjs/toolkit'
 import {
   isAllOf,
   isFulfilled,
   isPending,
   isRejected,
   isRejectedWithValue,
+  nanoid,
+  createAsyncThunk,
 } from '@reduxjs/toolkit'
 import type { Patch } from 'immer'
 import { isDraftable, produceWithPatches } from 'immer'
@@ -33,9 +34,9 @@ import type {
   ThunkAction,
   ThunkDispatch,
   AsyncThunk,
+  AsyncThunkPayloadCreator,
+  Draft,
 } from '@reduxjs/toolkit'
-import { createAsyncThunk } from '@reduxjs/toolkit'
-
 import { HandledError } from '../HandledError'
 
 import type { ApiEndpointQuery, PrefetchOptions } from './module'
@@ -105,6 +106,7 @@ export interface QueryThunkArg
   type: 'query'
   originalArgs: unknown
   endpointName: string
+  reducerPath: string
 }
 
 export interface MutationThunkArg {
@@ -411,6 +413,13 @@ In the case of an unhandled error, no tags will be "provided" or "invalidated".`
 
       return true
     },
+    idGenerator(args): string {
+      if (args.prefetch) {
+        return `${args.reducerPath}-${args.queryCacheKey}-prefetch`
+      }
+
+      return nanoid()
+    },
     dispatchConditionRejection: true,
   })
 
@@ -443,7 +452,7 @@ In the case of an unhandled error, no tags will be "provided" or "invalidated".`
       const queryAction = (force: boolean = true) =>
         (api.endpoints[endpointName] as ApiEndpointQuery<any, any>).initiate(
           arg,
-          { forceRefetch: force }
+          { forceRefetch: force, prefetch: options || true }
         )
       const latestStateValue = (
         api.endpoints[endpointName] as ApiEndpointQuery<any, any>

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

@@ -49,16 +49,21 @@ import { enablePatches } from 'immer'
 /**
  * `ifOlderThan` - (default: `false` | `number`) - _number is value in seconds_
  * - If specified, it will only run the query if the difference between `new Date()` and the last `fulfilledTimeStamp` is greater than the given value
- *
+ * 
+ * - `keepSubscriptionFor`: how long before the data is considered unused;
+ *   defaults to `api.config.keepPrefetchSubscriptionsFor`.  - _number is value in seconds_
+ * 
+ * 
  * @overloadSummary
  * `force`
  * - If `force: true`, it will ignore the `ifOlderThan` value if it is set and the query will be run even if it exists in the cache.
  */
 export type PrefetchOptions =
   | {
-      ifOlderThan?: false | number
+      ifOlderThan?: false | number,
+      keepSubscriptionFor?: number,
     }
-  | { force?: boolean }
+  | { force?: boolean, keepSubscriptionFor?: number, }
 
 export const coreModuleName = /* @__PURE__ */ Symbol()
 export type CoreModule =
@@ -365,6 +370,7 @@ export const coreModule = (): Module<CoreModule> => ({
       reducerPath,
       serializeQueryArgs,
       keepUnusedDataFor,
+      keepPrefetchSubscriptionsFor,
       refetchOnMountOrArgChange,
       refetchOnFocus,
       refetchOnReconnect,
@@ -427,6 +433,7 @@ export const coreModule = (): Module<CoreModule> => ({
         refetchOnReconnect,
         refetchOnMountOrArgChange,
         keepUnusedDataFor,
+        keepPrefetchSubscriptionsFor,
         reducerPath,
       },
     })

packages/toolkit/src/query/createApi.ts

@@ -126,6 +126,36 @@ export interface CreateApiOptions<
    * ```
    */
   keepUnusedDataFor?: number
+
+  /**
+   * Defaults to `10` _(this value is in seconds)_.
+   * 
+   * The default time to live of prefetch subscriptions.
+   *
+   * ```ts
+   * // codeblock-meta title="keepPrefetchSubscriptionsFor example"
+   *
+   * import { createApi, fetchBaseQuery } from '@reduxjs/toolkit/query/react'
+   * interface Post {
+   *   id: number
+   *   name: string
+   * }
+   * type PostsResponse = Post[]
+   *
+   * const api = createApi({
+   *   baseQuery: fetchBaseQuery({ baseUrl: '/' }),
+   *   keepPrefetchSubscriptionsFor: 5,
+   *   endpoints: (build) => ({
+   *     getPosts: build.query<PostsResponse, void>({
+   *       query: () => 'posts',
+   *       // highlight-start
+   *       // highlight-end
+   *     })
+   *   })
+   * })
+   * ```
+   */
+  keepPrefetchSubscriptionsFor?: number
   /**
    * Defaults to `false`. This setting allows you to control whether if a cached result is already available RTK Query will only serve a cached result, or if it should `refetch` when set to `true` or if an adequate amount of time has passed since the last successful query result.
    * - `false` - Will not cause a query to be performed _unless_ it does not exist yet.
@@ -240,6 +270,7 @@ export function buildCreateApi<Modules extends [Module<any>, ...Module<any>[]]>(
       reducerPath: 'api',
       serializeQueryArgs: defaultSerializeQueryArgs,
       keepUnusedDataFor: 60,
+      keepPrefetchSubscriptionsFor: 10,
       refetchOnMountOrArgChange: false,
       refetchOnFocus: false,
       refetchOnReconnect: false,

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

@@ -42,6 +42,7 @@ it('only resets the api state when resetApiState is dispatched', async () => {
       config: {
         focused: true,
         keepUnusedDataFor: 60,
+        keepPrefetchSubscriptionsFor: 10,
         middlewareRegistered: true,
         online: true,
         reducerPath: 'api',