Merge pull request #2239 from kahirokunn/retry-error

packages/toolkit/src/query/retry.ts

@@ -1,4 +1,11 @@
-import type { BaseQueryEnhancer } from './baseQueryTypes'
+import type {
+  BaseQueryApi,
+  BaseQueryArg,
+  BaseQueryEnhancer,
+  BaseQueryExtraOptions,
+  BaseQueryFn,
+} from './baseQueryTypes'
+import { FetchBaseQueryError } from './fetchBaseQuery'
 import { HandledError } from './HandledError'
 
 /**
@@ -23,31 +30,73 @@ async function defaultBackoff(attempt: number = 0, maxRetries: number = 5) {
   )
 }
 
-export interface RetryOptions {
-  /**
-   * How many times the query will be retried (default: 5)
-   */
-  maxRetries?: number
+type RetryConditionFunction = (
+  error: FetchBaseQueryError,
+  args: BaseQueryArg<BaseQueryFn>,
+  extraArgs: {
+    attempt: number
+    baseQueryApi: BaseQueryApi
+    extraOptions: BaseQueryExtraOptions<BaseQueryFn> & RetryOptions
+  }
+) => boolean
+
+export type RetryOptions = {
   /**
    * Function used to determine delay between retries
    */
   backoff?: (attempt: number, maxRetries: number) => Promise<void>
-}
+} & (
+  | {
+      /**
+       * How many times the query will be retried (default: 5)
+       */
+      maxRetries?: number
+      retryCondition?: undefined
+    }
+  | {
+      /**
+       * Callback to determine if a retry should be attempted.
+       * Return `true` for another retry and `false` to quit trying prematurely.
+       */
+      retryCondition?: RetryConditionFunction
+      maxRetries?: undefined
+    }
+)
 
 function fail(e: any): never {
   throw Object.assign(new HandledError({ error: e }), {
     throwImmediately: true,
   })
 }
 
+const EMPTY_OPTIONS = {}
+
 const retryWithBackoff: BaseQueryEnhancer<
   unknown,
   RetryOptions,
   RetryOptions | void
 > = (baseQuery, defaultOptions) => async (args, api, extraOptions) => {
-  const options = {
-    maxRetries: 5,
+  // We need to figure out `maxRetries` before we define `defaultRetryCondition.
+  // This is probably goofy, but ought to work.
+  // Put our defaults in one array, filter out undefineds, grab the last value.
+  const possibleMaxRetries: number[] = [
+    5,
+    ((defaultOptions as any) || EMPTY_OPTIONS).maxRetries,
+    ((extraOptions as any) || EMPTY_OPTIONS).maxRetries,
+  ].filter(Boolean)
+  const [maxRetries] = possibleMaxRetries.slice(-1)
+
+  const defaultRetryCondition: RetryConditionFunction = (_, __, { attempt }) =>
+    attempt <= maxRetries
+
+  const options: {
+    maxRetries: number
+    backoff: typeof defaultBackoff
+    retryCondition: typeof defaultRetryCondition
+  } = {
+    maxRetries,
     backoff: defaultBackoff,
+    retryCondition: defaultRetryCondition,
     ...defaultOptions,
     ...extraOptions,
   }
@@ -63,14 +112,26 @@ const retryWithBackoff: BaseQueryEnhancer<
       return result
     } catch (e: any) {
       retry++
-      if (e.throwImmediately || retry > options.maxRetries) {
+
+      if (e.throwImmediately) {
         if (e instanceof HandledError) {
           return e.value
         }
 
         // We don't know what this is, so we have to rethrow it
         throw e
       }
+
+      if (
+        e instanceof HandledError &&
+        !options.retryCondition(e.value as FetchBaseQueryError, args, {
+          attempt: retry,
+          baseQueryApi: api,
+          extraOptions,
+        })
+      ) {
+        return e.value
+      }
       await options.backoff(retry, options.maxRetries)
     }
   }

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

@@ -1,6 +1,7 @@
 import type { BaseQueryFn } from '@reduxjs/toolkit/query'
 import { createApi, retry } from '@reduxjs/toolkit/query'
 import { setupApiStore, waitMs } from './helpers'
+import type { RetryOptions } from '../retry'
 
 beforeEach(() => {
   jest.useFakeTimers('legacy')
@@ -339,4 +340,110 @@ describe('configuration', () => {
 
     expect(baseBaseQuery).toHaveBeenCalledTimes(9)
   })
+
+  test('accepts a custom retryCondition fn', async () => {
+    const baseBaseQuery = jest.fn<
+      ReturnType<BaseQueryFn>,
+      Parameters<BaseQueryFn>
+    >()
+    baseBaseQuery.mockResolvedValue({ error: 'rejected' })
+
+    const overrideMaxRetries = 3
+
+    const baseQuery = retry(baseBaseQuery, {
+      retryCondition: (_, __, { attempt }) => attempt <= overrideMaxRetries,
+    })
+    const api = createApi({
+      baseQuery,
+      endpoints: (build) => ({
+        q1: build.query({
+          query: () => {},
+        }),
+      }),
+    })
+
+    const storeRef = setupApiStore(api, undefined, {
+      withoutTestLifecycles: true,
+    })
+    storeRef.store.dispatch(api.endpoints.q1.initiate({}))
+
+    await loopTimers()
+
+    expect(baseBaseQuery).toHaveBeenCalledTimes(overrideMaxRetries + 1)
+  })
+
+  test('retryCondition with endpoint config that overrides baseQuery config', async () => {
+    const baseBaseQuery = jest.fn<
+      ReturnType<BaseQueryFn>,
+      Parameters<BaseQueryFn>
+    >()
+    baseBaseQuery.mockResolvedValue({ error: 'rejected' })
+
+    const baseQuery = retry(baseBaseQuery, {
+      maxRetries: 10,
+    })
+    const api = createApi({
+      baseQuery,
+      endpoints: (build) => ({
+        q1: build.query({
+          query: () => {},
+          extraOptions: {
+            retryCondition: (_, __, { attempt }) => attempt <= 5,
+          },
+        }),
+      }),
+    })
+
+    const storeRef = setupApiStore(api, undefined, {
+      withoutTestLifecycles: true,
+    })
+    storeRef.store.dispatch(api.endpoints.q1.initiate({}))
+
+    await loopTimers()
+
+    expect(baseBaseQuery).toHaveBeenCalledTimes(6)
+  })
+
+  test('retryCondition also works with mutations', async () => {
+    const baseBaseQuery = jest.fn<
+      ReturnType<BaseQueryFn>,
+      Parameters<BaseQueryFn>
+    >()
+
+    baseBaseQuery
+      .mockRejectedValueOnce(new Error('rejected'))
+      .mockRejectedValueOnce(new Error('hello retryCondition'))
+      .mockRejectedValueOnce(new Error('rejected'))
+      .mockResolvedValue({ error: 'hello retryCondition' })
+
+    const baseQuery = retry(baseBaseQuery, {})
+    const api = createApi({
+      baseQuery,
+      endpoints: (build) => ({
+        m1: build.mutation({
+          query: () => ({ method: 'PUT' }),
+          extraOptions: {
+            retryCondition: (e) => e.data === 'hello retryCondition',
+          },
+        }),
+      }),
+    })
+
+    const storeRef = setupApiStore(api, undefined, {
+      withoutTestLifecycles: true,
+    })
+    storeRef.store.dispatch(api.endpoints.m1.initiate({}))
+
+    await loopTimers()
+
+    expect(baseBaseQuery).toHaveBeenCalledTimes(4)
+  })
+
+  test.skip('RetryOptions only accepts one of maxRetries or retryCondition', () => {
+    // @ts-expect-error Should complain if both exist at once
+    const ro: RetryOptions = {
+      maxRetries: 5,
+      retryCondition: () => false,
+    }
+  })
 })