Add shouldRetry option (#75)

pukuba · sindresorhus · web-flow · commit b993a96738c5 · 2023-12-20T21:37:20.000+01:00
Co-authored-by: Sindre Sorhus &lt;sindresorhus@gmail.com&gt;
diff --git a/index.d.ts b/index.d.ts
@@ -41,6 +41,28 @@ export type Options = {
 	*/
 	readonly onFailedAttempt?: (error: FailedAttemptError) => void | Promise<void>;
 
+	/**
+	Decide if a retry should occur based on the error. Returning true triggers a retry, false aborts with the error.
+
+	It is not called for `TypeError` (except network errors) and `AbortError`.
+
+	@param error - The error thrown by the input function.
+
+	@example
+	```
+	import pRetry from 'p-retry';
+
+	const run = async () => { … };
+
+	const result = await pRetry(run, {
+		shouldRetry: error => !(error instanceof CustomError);
+	});
+	```
+
+	In the example above, the operation will be retried unless the error is an instance of `CustomError`.
+	*/
+	readonly shouldRetry?: (error: FailedAttemptError) => boolean | Promise<boolean>;
+
 	/**
 	You can abort retrying using [`AbortController`](https://developer.mozilla.org/en-US/docs/Web/API/AbortController).
 
diff --git a/index.js b/index.js
@@ -32,6 +32,7 @@ export default async function pRetry(input, options) {
 		options = {
 			onFailedAttempt() {},
 			retries: 10,
+			shouldRetry: () => true,
 			...options,
 		};
 
@@ -70,7 +71,14 @@ export default async function pRetry(input, options) {
 						throw error;
 					}
 
-					await options.onFailedAttempt(decorateErrorWithCounts(error, attemptNumber, options));
+					decorateErrorWithCounts(error, attemptNumber, options);
+
+					if (!(await options.shouldRetry(error))) {
+						operation.stop();
+						reject(error);
+					}
+
+					await options.onFailedAttempt(error);
 
 					if (!operation.retry(error)) {
 						throw operation.mainError();
diff --git a/readme.md b/readme.md
@@ -99,6 +99,26 @@ const result = await pRetry(run, {
 
 If the `onFailedAttempt` function throws, all retries will be aborted and the original promise will reject with the thrown error.
 
+##### shouldRetry(error)
+
+Type: `Function`
+
+Decide if a retry should occur based on the error. Returning true triggers a retry, false aborts with the error.
+
+It is not called for `TypeError` (except network errors) and `AbortError`.
+
+```js
+import pRetry from 'p-retry';
+
+const run = async () => { … };
+
+const result = await pRetry(run, {
+	shouldRetry: error => !(error instanceof CustomError);
+});
+```
+
+In the example above, the operation will be retried unless the error is an instance of `CustomError`.
+
 ##### signal
 
 Type: [`AbortSignal`](https://developer.mozilla.org/en-US/docs/Web/API/AbortSignal)
diff --git a/test.js b/test.js
@@ -277,3 +277,28 @@ test('preserves the abort reason', async t => {
 
 	t.is(index, 3);
 });
+
+test('should retry only when shouldRetry returns true', async t => {
+	t.plan(2);
+
+	const shouldRetryError = new Error('should-retry');
+	const customError = new Error('custom-error');
+
+	let index = 0;
+
+	await t.throwsAsync(pRetry(async () => {
+		await delay(40);
+		index++;
+		const error = index < 3 ? shouldRetryError : customError;
+		throw error;
+	}, {
+		async shouldRetry(error) {
+			return error.message === shouldRetryError.message;
+		},
+		retries: 10,
+	}), {
+		is: customError,
+	});
+
+	t.is(index, 3);
+});
