feat: add retry quotas to StandardRetryStrategy

package.json

@@ -49,7 +49,7 @@
     "jmespath": "^0.15.0",
     "lerna": "3.20.2",
     "lint-staged": "^10.0.1",
-    "prettier": "2.0.4",
+    "prettier": "2.0.5",
     "typescript": "~3.8.3",
     "yarn": "1.22.4"
   },

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

@@ -0,0 +1,47 @@
+import { resolveRetryConfig } from "./configurations";
+import { StandardRetryStrategy } from "./defaultStrategy";
+
+describe("resolveRetryConfig", () => {
+  describe("maxAttempts", () => {
+    it("uses passed maxAttempts value if present", () => {
+      [1, 2, 3].forEach(maxAttempts => {
+        expect(resolveRetryConfig({ maxAttempts }).maxAttempts).toEqual(
+          maxAttempts
+        );
+      });
+    });
+
+    it("assigns default value of 3 if maxAttempts not passed", () => {
+      expect(resolveRetryConfig({}).maxAttempts).toEqual(3);
+    });
+  });
+
+  describe("retryStrategy", () => {
+    it("uses passed retryStrategy if present", () => {
+      const mockRetryStrategy = {
+        maxAttempts: 2,
+        retry: jest.fn()
+      };
+      const { retryStrategy } = resolveRetryConfig({
+        retryStrategy: mockRetryStrategy
+      });
+      expect(retryStrategy).toEqual(mockRetryStrategy);
+    });
+
+    describe("creates StandardRetryStrategy if retryStrategy not present", () => {
+      describe("uses maxAttempts if present", () => {
+        [1, 2, 3].forEach(maxAttempts => {
+          const { retryStrategy } = resolveRetryConfig({ maxAttempts });
+          expect(retryStrategy).toBeInstanceOf(StandardRetryStrategy);
+          expect(retryStrategy.maxAttempts).toBe(maxAttempts);
+        });
+      });
+
+      it("uses default 3 if maxAttempts is not present", () => {
+        const { retryStrategy } = resolveRetryConfig({});
+        expect(retryStrategy).toBeInstanceOf(StandardRetryStrategy);
+        expect(retryStrategy.maxAttempts).toBe(3);
+      });
+    });
+  });
+});

packages/middleware-retry/src/configurations.ts

@@ -17,13 +17,13 @@ export interface RetryResolvedConfig {
   retryStrategy: RetryStrategy;
 }
 
-export function resolveRetryConfig<T>(
+export const resolveRetryConfig = <T>(
   input: T & RetryInputConfig
-): T & RetryResolvedConfig {
-  const maxAttempts = input.maxAttempts === undefined ? 3 : input.maxAttempts;
+): T & RetryResolvedConfig => {
+  const maxAttempts = input.maxAttempts ?? 3;
   return {
     ...input,
     maxAttempts,
     retryStrategy: input.retryStrategy || new StandardRetryStrategy(maxAttempts)
   };
-}
+};

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,149 @@
+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) => {
+    const retryQuota = getDefaultRetryQuota();
+    let availableCapacity = INITIAL_RETRY_TOKENS;
+    while (availableCapacity >= targetCapacity) {
+      retryQuota.retrieveRetryTokens(error);
+      availableCapacity -= targetCapacity;
+    }
+    return retryQuota;
+  };
+
+  describe("hasRetryTokens", () => {
+    describe("returns true if capacity is available", () => {
+      it("when it's TimeoutError", () => {
+        const timeoutError = getMockTimeoutError();
+        expect(getDefaultRetryQuota().hasRetryTokens(timeoutError)).toBe(true);
+      });
+
+      it("when it's not TimeoutError", () => {
+        expect(getDefaultRetryQuota().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().retrieveRetryTokens(timeoutError)).toBe(
+          TIMEOUT_RETRY_COST
+        );
+      });
+
+      it("when it's not TimeoutError", () => {
+        expect(getDefaultRetryQuota().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();
+
+      // 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,40 @@
+import { RetryQuota } from "./defaultStrategy";
+import { SdkError } from "@aws-sdk/smithy-client";
+import {
+  INITIAL_RETRY_TOKENS,
+  RETRY_COST,
+  TIMEOUT_RETRY_COST,
+  NO_RETRY_INCREMENT
+} from "./constants";
+
+export const getDefaultRetryQuota = (): RetryQuota => {
+  const MAX_CAPACITY = INITIAL_RETRY_TOKENS;
+  let availableCapacity = INITIAL_RETRY_TOKENS;
+
+  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
+  });
+};