add cache lifecycle

src/query/core/buildMiddleware/cacheLifecycle.ts

@@ -0,0 +1,118 @@
+import { isAsyncThunkAction, isFulfilled } from '@reduxjs/toolkit'
+import { toOptionalPromise } from '../../utils/toOptionalPromise'
+import { SubMiddlewareApi, SubMiddlewareBuilder } from './types'
+
+export const build: SubMiddlewareBuilder = ({
+  api,
+  reducerPath,
+  context,
+  queryThunk,
+  mutationThunk,
+}) => {
+  type CacheLifecycle = {
+    valueResolved?(value: unknown): unknown
+    cleanup(): void
+  }
+  const lifecycleMap: Record<string, CacheLifecycle> = {}
+
+  const isQueryThunk = isAsyncThunkAction(queryThunk)
+  const isMutationThunk = isAsyncThunkAction(mutationThunk)
+  const isFullfilledThunk = isFulfilled(queryThunk, mutationThunk)
+
+  return (mwApi) => (next) => (action): any => {
+    const result = next(action)
+
+    const cacheKey = getCacheKey(action)
+
+    if (queryThunk.pending.match(action)) {
+      const state = mwApi.getState()[reducerPath].queries[cacheKey]
+      if (state?.requestId === action.meta.requestId) {
+        handleNewKey(
+          action.meta.arg.endpointName,
+          action.meta.arg.originalArgs,
+          cacheKey,
+          mwApi
+        )
+      }
+    } else if (mutationThunk.pending.match(action)) {
+      const state = mwApi.getState()[reducerPath].mutations[cacheKey]
+      if (state) {
+        handleNewKey(
+          action.meta.arg.endpointName,
+          action.meta.arg.originalArgs,
+          cacheKey,
+          mwApi
+        )
+      }
+    } else if (isFullfilledThunk(action)) {
+      const lifecycle = lifecycleMap[cacheKey]
+      if (lifecycle?.valueResolved) {
+        lifecycle.valueResolved(action.payload)
+        delete lifecycle.valueResolved
+      }
+    } else if (
+      api.internalActions.removeQueryResult.match(action) ||
+      api.internalActions.unsubscribeMutationResult.match(action)
+    ) {
+      const lifecycle = lifecycleMap[cacheKey]
+      if (lifecycle) {
+        delete lifecycleMap[cacheKey]
+        lifecycle.cleanup()
+      }
+    } else if (api.util.resetApiState.match(action)) {
+      for (const [cacheKey, lifecycle] of Object.entries(lifecycleMap)) {
+        delete lifecycleMap[cacheKey]
+        lifecycle.cleanup()
+      }
+    }
+
+    return result
+  }
+
+  function getCacheKey(action: any) {
+    if (isQueryThunk(action)) return action.meta.arg.queryCacheKey
+    if (isMutationThunk(action)) return action.meta.requestId
+    if (api.internalActions.removeQueryResult.match(action))
+      return action.payload.queryCacheKey
+    return ''
+  }
+
+  function handleNewKey(
+    endpointName: string,
+    originalArgs: any,
+    queryCacheKey: string,
+    mwApi: SubMiddlewareApi
+  ) {
+    const { onCacheEntryAdded } = context.endpointDefinitions[endpointName]
+    if (!onCacheEntryAdded) return
+
+    const neverResolvedError = new Error(
+      'Promise never resolved before cleanup.'
+    )
+    let lifecycle = {} as CacheLifecycle
+
+    const cleanup = new Promise<void>((resolve) => {
+      lifecycle.cleanup = resolve
+    })
+    const firstValueResolved = toOptionalPromise(
+      Promise.race([
+        new Promise<void>((resolve) => {
+          lifecycle.valueResolved = resolve
+        }),
+        cleanup.then(() => {
+          throw neverResolvedError
+        }),
+      ])
+    )
+    lifecycleMap[queryCacheKey] = lifecycle
+    const runningHandler = onCacheEntryAdded(originalArgs, mwApi, {
+      firstValueResolved,
+      cleanup,
+    })
+    // if a `neverResolvedError` was thrown, but not handled in the running handler, do not let it leak out further
+    Promise.resolve(runningHandler).catch((e) => {
+      if (e === neverResolvedError) return
+      throw e
+    })
+  }
+}

src/query/endpointDefinitions.ts

@@ -17,6 +17,7 @@ import {
   CastAny,
 } from './tsHelpers'
 import { NEVER } from './fakeBaseQuery'
+import { OptionalPromise } from './utils/toOptionalPromise'
 
 const resultType = Symbol()
 const baseQuery = Symbol()
@@ -55,6 +56,11 @@ interface EndpointDefinitionWithQueryFn<
   transformResponse?: never
 }
 
+export type LifecycleApi = {
+  dispatch: ThunkDispatch<any, any, any>
+  getState: () => unknown
+}
+
 export type BaseEndpointDefinition<
   QueryArg,
   BaseQuery extends BaseQueryFn,
@@ -69,6 +75,30 @@ export type BaseEndpointDefinition<
   [resultType]?: ResultType
   /* phantom type */
   [baseQuery]?: BaseQuery
+  onCacheEntryAdded?(
+    arg: QueryArg,
+    api: LifecycleApi,
+    promises: {
+      /**
+       * Promise that will resolve with the first value for this cache key.
+       * This allows you to `await` until an actual value is in cache.
+       *
+       * If the cache entry is removed from the cache before any value has ever
+       * been resolved, this Promise will reject with
+       * `new Error('Promise never resolved before cleanup.')`
+       * to prevent memory leaks.
+       * You can just re-throw that error (or not handle it at all) -
+       * it will be caught outside of `cacheEntryAdded`.
+       */
+      firstValueResolved: OptionalPromise<ResultType>
+      /**
+       * Promise that allows you to wait for the point in time when the cache entry
+       * has been removed from the cache, by not being used/subscribed to any more
+       * in the application for too long or by dispatching `api.util.resetApiState`.
+       */
+      cleanup: Promise<void>
+    }
+  ): Promise<void> | void
 } & HasRequiredProps<
     BaseQueryExtraOptions<BaseQuery>,
     { extraOptions: BaseQueryExtraOptions<BaseQuery> },

src/query/tests/optionalPromise.test.ts

@@ -0,0 +1,22 @@
+import { toOptionalPromise } from '../utils/toOptionalPromise'
+
+const id = <T>(x: T) => x
+test('should .then normally', async () => {
+  await expect(toOptionalPromise(Promise.resolve(5)).then(id)).resolves.toBe(5)
+})
+
+test('should .catch normally', async () => {
+  await expect(toOptionalPromise(Promise.reject(6)).catch(id)).resolves.toBe(6)
+})
+
+test('should .finally normally', async () => {
+  const finale = jest.fn()
+  await expect(
+    toOptionalPromise(Promise.reject(6)).finally(finale).catch(id)
+  ).resolves.toBe(6)
+  expect(finale).toHaveBeenCalled()
+})
+
+test('not interacting should not make jest freak out', () => {
+  toOptionalPromise(Promise.reject(6))
+})

src/query/utils/toOptionalPromise.ts

@@ -0,0 +1,25 @@
+/**
+ * A promise that will only throw if it has been interacted with, by calling `.then` or `.catch` on it,
+ * or by `await`ing it.
+ *
+ * This deals with unhandled promise rejections if the promise was never interacted with in any way.
+ *
+ * No interaction with the Promise => no error
+ */
+export interface OptionalPromise<T> extends Promise<T> {}
+
+/**
+ * Wraps a Promise in a new Promise that will only throw, if either `.then` or `.catch` have been accessed on it prior to throwing.
+ */
+export function toOptionalPromise<T>(promise: Promise<T>): OptionalPromise<T> {
+  let interacted = false
+  const wrapped = promise.catch((e) => {
+    if (interacted) throw e
+  })
+  const { then } = wrapped
+  wrapped.then = (...args) => {
+    interacted = true
+    return then.apply(wrapped, args) as any
+  }
+  return wrapped as Promise<T>
+}