Add automatic retry policy

Author: mattt

index.js

@@ -1,4 +1,5 @@
 const ApiError = require('./lib/error');
+const { withAutomaticRetries } = require('./lib/util');
 
 const collections = require('./lib/collections');
 const models = require('./lib/models');
@@ -201,7 +202,7 @@ class Replicate {
       body: data ? JSON.stringify(data) : undefined,
     };
 
-    const response = await this.fetch(url, init);
+    const response = await withAutomaticRetries(async () => this.fetch(url, init));
 
     if (!response.ok) {
       const request = new Request(url, init);

lib/util.js

@@ -0,0 +1,75 @@
+const ApiError = require('./error');
+
+/**
+ * Automatically retry a request if it fails with an appropriate status code.
+ *
+ * A GET request is retried if it fails with a 429 or 5xx status code.
+ * A non-GET request is retried only if it fails with a 429 status code.
+ *
+ * If the response sets a Retry-After header,
+ * the request is retried after the number of seconds specified in the header.
+ * Otherwise, the request is retried after the specified interval,
+ * with exponential backoff and jitter.
+ *
+ * @param {Function} request - A function that returns a Promise that resolves with a Response object
+ * @param {object} options
+ * @param {number} [options.retries] - Number of retries. Defaults to 5
+ * @param {number} [options.interval] - Interval between retries in milliseconds. Defaults to 500
+ * @returns {Promise<Response>} - Resolves with the response object
+ * @throws {ApiError} If the request failed
+ */
+async function withAutomaticRetries(request, options = {}) {
+  const retries = options.retries || 5;
+  const interval = options.interval || 500;
+  const jitter = options.jitter || 100;
+
+  const shouldRetry = (response) => {
+    if (request.method === 'GET') {
+      return response.status === 429 || response.status >= 500;
+    }
+
+    return response.status === 429;
+  };
+
+  // eslint-disable-next-line no-promise-executor-return
+  const sleep = (ms) => new Promise((resolve) => setTimeout(resolve, ms));
+
+  let attempts = 0;
+  do {
+    let delay = (interval * (2 ** attempts)) + (Math.random() * jitter);
+
+    /* eslint-disable no-await-in-loop */
+    try {
+      const response = await request();
+      if (response.ok || !shouldRetry(response)) {
+        return response;
+      }
+    } catch (error) {
+      if (error instanceof ApiError) {
+        const retryAfter = error.response.headers.get('Retry-After');
+        if (retryAfter) {
+          if (!Number.isInteger(retryAfter)) { // Retry-After is a date
+            const date = new Date(retryAfter);
+            if (!Number.isNaN(date.getTime())) {
+              delay = date.getTime() - new Date().getTime();
+            }
+          } else { // Retry-After is a number of seconds
+            delay = retryAfter * 1000;
+          }
+        }
+      }
+    }
+
+    if (Number.isInteger(retries) && retries > 0) {
+      if (Number.isInteger(delay) && delay > 0) {
+        await sleep(interval * 2 ** (options.retries - retries));
+      }
+      attempts += 1;
+    }
+    /* eslint-enable no-await-in-loop */
+  } while (attempts < retries);
+
+  return request();
+}
+
+module.exports = { withAutomaticRetries };