feat: add retry quotas to StandardRetryStrategy

packages/middleware-retry/src/constants.ts

@@ -15,3 +15,25 @@ export const MAXIMUM_RETRY_DELAY = 20 * 1000;
  * encountered.
  */
 export const THROTTLING_RETRY_DELAY_BASE = 500;
+
+/**
+ * Initial number of retry tokens in Retry Quota
+ */
+export const INITIAL_RETRY_TOKENS = 500;
+
+/**
+ * The total amount of retry tokens to be decremented from retry token balance.
+ */
+export const RETRY_COST = 5;
+
+/**
+ * The total amount of retry tokens to be decremented from retry token balance
+ * when a throttling error is encountered.
+ */
+export const TIMEOUT_RETRY_COST = 10;
+
+/**
+ * The total amount of retry token to be incremented from retry token balance
+ * if an SDK operation invocation succeeds without requiring a retry request.
+ */
+export const NO_RETRY_INCREMENT = 1;

packages/middleware-retry/src/defaultRetryQuota.spec.ts

@@ -0,0 +1,189 @@
+import { getDefaultRetryQuota } from "./defaultRetryQuota";
+import { SdkError } from "@aws-sdk/smithy-client";
+import {
+  INITIAL_RETRY_TOKENS,
+  TIMEOUT_RETRY_COST,
+  RETRY_COST,
+  NO_RETRY_INCREMENT
+} from "./constants";
+
+describe("defaultRetryQuota", () => {
+  const getMockError = () => new Error() as SdkError;
+  const getMockTimeoutError = () =>
+    Object.assign(new Error(), {
+      name: "TimeoutError"
+    }) as SdkError;
+
+  const getDrainedRetryQuota = (
+    targetCapacity: number,
+    error: SdkError,
+    initialRetryTokens: number = INITIAL_RETRY_TOKENS
+  ) => {
+    const retryQuota = getDefaultRetryQuota(initialRetryTokens);
+    let availableCapacity = initialRetryTokens;
+    while (availableCapacity >= targetCapacity) {
+      retryQuota.retrieveRetryTokens(error);
+      availableCapacity -= targetCapacity;
+    }
+    return retryQuota;
+  };
+
+  describe("custom initial retry tokens", () => {
+    it("hasRetryTokens returns false if capacity is not available", () => {
+      const customRetryTokens = 100;
+      const error = getMockError();
+      const retryQuota = getDrainedRetryQuota(
+        RETRY_COST,
+        error,
+        customRetryTokens
+      );
+      expect(retryQuota.hasRetryTokens(error)).toBe(false);
+    });
+
+    it("retrieveRetryToken throws error if retry tokens not available", () => {
+      const customRetryTokens = 100;
+      const error = getMockError();
+      const retryQuota = getDrainedRetryQuota(
+        RETRY_COST,
+        error,
+        customRetryTokens
+      );
+      expect(() => {
+        retryQuota.retrieveRetryTokens(error);
+      }).toThrowError(new Error("No retry token available"));
+    });
+  });
+
+  describe("hasRetryTokens", () => {
+    describe("returns true if capacity is available", () => {
+      it("when it's TimeoutError", () => {
+        const timeoutError = getMockTimeoutError();
+        expect(
+          getDefaultRetryQuota(INITIAL_RETRY_TOKENS).hasRetryTokens(
+            timeoutError
+          )
+        ).toBe(true);
+      });
+
+      it("when it's not TimeoutError", () => {
+        expect(
+          getDefaultRetryQuota(INITIAL_RETRY_TOKENS).hasRetryTokens(
+            getMockError()
+          )
+        ).toBe(true);
+      });
+    });
+
+    describe("returns false if capacity is not available", () => {
+      it("when it's TimeoutError", () => {
+        const timeoutError = getMockTimeoutError();
+        const retryQuota = getDrainedRetryQuota(
+          TIMEOUT_RETRY_COST,
+          timeoutError
+        );
+        expect(retryQuota.hasRetryTokens(timeoutError)).toBe(false);
+      });
+
+      it("when it's not TimeoutError", () => {
+        const error = getMockError();
+        const retryQuota = getDrainedRetryQuota(RETRY_COST, error);
+        expect(retryQuota.hasRetryTokens(error)).toBe(false);
+      });
+    });
+  });
+
+  describe("retrieveRetryToken", () => {
+    describe("returns retry tokens amount if available", () => {
+      it("when it's TimeoutError", () => {
+        const timeoutError = getMockTimeoutError();
+        expect(
+          getDefaultRetryQuota(INITIAL_RETRY_TOKENS).retrieveRetryTokens(
+            timeoutError
+          )
+        ).toBe(TIMEOUT_RETRY_COST);
+      });
+
+      it("when it's not TimeoutError", () => {
+        expect(
+          getDefaultRetryQuota(INITIAL_RETRY_TOKENS).retrieveRetryTokens(
+            getMockError()
+          )
+        ).toBe(RETRY_COST);
+      });
+    });
+
+    describe("throws error if retry tokens not available", () => {
+      it("when it's TimeoutError", () => {
+        const timeoutError = getMockTimeoutError();
+        const retryQuota = getDrainedRetryQuota(
+          TIMEOUT_RETRY_COST,
+          timeoutError
+        );
+        expect(() => {
+          retryQuota.retrieveRetryTokens(timeoutError);
+        }).toThrowError(new Error("No retry token available"));
+      });
+
+      it("when it's not TimeoutError", () => {
+        const error = getMockError();
+        const retryQuota = getDrainedRetryQuota(RETRY_COST, error);
+        expect(() => {
+          retryQuota.retrieveRetryTokens(error);
+        }).toThrowError(new Error("No retry token available"));
+      });
+    });
+  });
+
+  describe("releaseRetryToken", () => {
+    it("adds capacityReleaseAmount if passed", () => {
+      const error = getMockError();
+      const retryQuota = getDrainedRetryQuota(RETRY_COST, error);
+
+      // Ensure that retry tokens are not available.
+      expect(retryQuota.hasRetryTokens(error)).toBe(false);
+
+      // Release RETRY_COST tokens.
+      retryQuota.releaseRetryTokens(RETRY_COST);
+      expect(retryQuota.hasRetryTokens(error)).toBe(true);
+      expect(retryQuota.retrieveRetryTokens(error)).toBe(RETRY_COST);
+      expect(retryQuota.hasRetryTokens(error)).toBe(false);
+    });
+
+    it("adds NO_RETRY_INCREMENT if capacityReleaseAmount not passed", () => {
+      const error = getMockError();
+      const retryQuota = getDrainedRetryQuota(RETRY_COST, error);
+
+      // retry tokens will not be available till NO_RETRY_INCREMENT is added
+      // till it's equal to RETRY_COST - (INITIAL_RETRY_TOKENS % RETRY_COST)
+      let tokensReleased = 0;
+      const tokensToBeReleased =
+        RETRY_COST - (INITIAL_RETRY_TOKENS % RETRY_COST);
+      while (tokensReleased < tokensToBeReleased) {
+        expect(retryQuota.hasRetryTokens(error)).toBe(false);
+        retryQuota.releaseRetryTokens();
+        tokensReleased += NO_RETRY_INCREMENT;
+      }
+      expect(retryQuota.hasRetryTokens(error)).toBe(true);
+    });
+
+    it("ensures availableCapacity is maxed at INITIAL_RETRY_TOKENS", () => {
+      const error = getMockError();
+      const retryQuota = getDefaultRetryQuota(INITIAL_RETRY_TOKENS);
+
+      // release 100 tokens.
+      [...Array(100).keys()].forEach(key => {
+        retryQuota.releaseRetryTokens();
+      });
+
+      // availableCapacity is still maxed at INITIAL_RETRY_TOKENS
+      // hasRetryTokens would be true only till INITIAL_RETRY_TOKENS/RETRY_COST times
+      [...Array(Math.floor(INITIAL_RETRY_TOKENS / RETRY_COST)).keys()].forEach(
+        key => {
+          expect(retryQuota.hasRetryTokens(error)).toBe(true);
+          retryQuota.retrieveRetryTokens(error);
+        }
+      );
+      expect(retryQuota.hasRetryTokens(error)).toBe(false);
+    });
+  });
+});

packages/middleware-retry/src/defaultRetryQuota.ts

@@ -0,0 +1,41 @@
+import { RetryQuota } from "./defaultStrategy";
+import { SdkError } from "@aws-sdk/smithy-client";
+import {
+  RETRY_COST,
+  TIMEOUT_RETRY_COST,
+  NO_RETRY_INCREMENT
+} from "./constants";
+
+export const getDefaultRetryQuota = (
+  initialRetryTokens: number
+): RetryQuota => {
+  const MAX_CAPACITY = initialRetryTokens;
+  let availableCapacity = initialRetryTokens;
+
+  const getCapacityAmount = (error: SdkError) =>
+    error.name === "TimeoutError" ? TIMEOUT_RETRY_COST : RETRY_COST;
+
+  const hasRetryTokens = (error: SdkError) =>
+    getCapacityAmount(error) <= availableCapacity;
+
+  const retrieveRetryTokens = (error: SdkError) => {
+    if (!hasRetryTokens(error)) {
+      // retryStrategy should stop retrying, and return last error
+      throw new Error("No retry token available");
+    }
+    const capacityAmount = getCapacityAmount(error);
+    availableCapacity -= capacityAmount;
+    return capacityAmount;
+  };
+
+  const releaseRetryTokens = (capacityReleaseAmount?: number) => {
+    availableCapacity += capacityReleaseAmount ?? NO_RETRY_INCREMENT;
+    availableCapacity = Math.min(availableCapacity, MAX_CAPACITY);
+  };
+
+  return Object.freeze({
+    hasRetryTokens,
+    retrieveRetryTokens,
+    releaseRetryTokens
+  });
+};