From 1155483dd0371b0997b78a401e7edd299538ee12 Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Tue, 9 Jun 2020 18:31:58 +0000
Subject: [PATCH 1/9] chore: add optional RetryQuota to StandardRetryStrategy

---
 .../middleware-retry/src/defaultStrategy.ts   | 24 +++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/packages/middleware-retry/src/defaultStrategy.ts b/packages/middleware-retry/src/defaultStrategy.ts
index c2b5d798ed8c0..7b91e9563589b 100644
--- a/packages/middleware-retry/src/defaultStrategy.ts
+++ b/packages/middleware-retry/src/defaultStrategy.ts
@@ -33,17 +33,40 @@ export interface DelayDecider {
   (delayBase: number, attempts: number): number;
 }
 
+/**
+ * Interface that specifies the retry quota behavior.
+ */
+export interface RetryQuota {
+  /**
+   * returns true if retry tokens are available from the retry quota bucket.
+   */
+  hasRetryTokens: (error: SdkError) => boolean;
+
+  /**
+   * returns token amount from the retry quota bucket.
+   * throws error is retry tokens are not available.
+   */
+  retrieveRetryToken: (error: SdkError) => number;
+
+  /**
+   * releases tokens back to the retry quota.
+   */
+  releaseRetryToken: (releaseCapacityAmount?: number) => void;
+}
+
 /**
  * Strategy options to be passed to StandardRetryStrategy
  */
 export interface StandardRetryStrategyOptions {
   retryDecider?: RetryDecider;
   delayDecider?: DelayDecider;
+  retryQuota?: RetryQuota;
 }
 
 export class StandardRetryStrategy implements RetryStrategy {
   private retryDecider: RetryDecider;
   private delayDecider: DelayDecider;
+  private retryQuota?: RetryQuota;
 
   constructor(
     public readonly maxAttempts: number,
@@ -51,6 +74,7 @@ export class StandardRetryStrategy implements RetryStrategy {
   ) {
     this.retryDecider = options?.retryDecider ?? defaultRetryDecider;
     this.delayDecider = options?.delayDecider ?? defaultDelayDecider;
+    this.retryQuota = options?.retryQuota;
   }
 
   private shouldRetry(error: SdkError, attempts: number) {

From f1064c51ea015a5bd1e6107e9673ec8c27cb1abf Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Tue, 9 Jun 2020 21:07:48 +0000
Subject: [PATCH 2/9] chore: add defaultRetryQuota

---
 packages/middleware-retry/src/constants.ts    |  22 +++
 .../src/defaultRetryQuota.spec.ts             | 147 ++++++++++++++++++
 .../middleware-retry/src/defaultRetryQuota.ts |  40 +++++
 .../middleware-retry/src/defaultStrategy.ts   |   4 +-
 4 files changed, 211 insertions(+), 2 deletions(-)
 create mode 100644 packages/middleware-retry/src/defaultRetryQuota.spec.ts
 create mode 100644 packages/middleware-retry/src/defaultRetryQuota.ts

diff --git a/packages/middleware-retry/src/constants.ts b/packages/middleware-retry/src/constants.ts
index 5f976928e5141..c14baa7b87c33 100644
--- a/packages/middleware-retry/src/constants.ts
+++ b/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;
diff --git a/packages/middleware-retry/src/defaultRetryQuota.spec.ts b/packages/middleware-retry/src/defaultRetryQuota.spec.ts
new file mode 100644
index 0000000000000..eb6e8df7d6395
--- /dev/null
+++ b/packages/middleware-retry/src/defaultRetryQuota.spec.ts
@@ -0,0 +1,147 @@
+import { defaultRetryQuota } 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 = defaultRetryQuota();
+    let availableCapacity = INITIAL_RETRY_TOKENS;
+    while (availableCapacity >= targetCapacity) {
+      retryQuota.retrieveRetryToken(error);
+      availableCapacity -= targetCapacity;
+    }
+    return retryQuota;
+  };
+
+  describe("hasRetryTokens", () => {
+    describe("returns true if capacity is available", () => {
+      it("when it's TimeoutError", () => {
+        const timeoutError = getMockTimeoutError();
+        expect(defaultRetryQuota().hasRetryTokens(timeoutError)).toBe(true);
+      });
+
+      it("when it's not TimeoutError", () => {
+        expect(defaultRetryQuota().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(defaultRetryQuota().retrieveRetryToken(timeoutError)).toBe(
+          TIMEOUT_RETRY_COST
+        );
+      });
+
+      it("when it's not TimeoutError", () => {
+        expect(defaultRetryQuota().retrieveRetryToken(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.retrieveRetryToken(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.retrieveRetryToken(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.releaseRetryToken(RETRY_COST);
+      expect(retryQuota.hasRetryTokens(error)).toBe(true);
+      expect(retryQuota.retrieveRetryToken(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.releaseRetryToken();
+        tokensReleased += NO_RETRY_INCREMENT;
+      }
+      expect(retryQuota.hasRetryTokens(error)).toBe(true);
+    });
+
+    it("ensures availableCapacity is maxed at INITIAL_RETRY_TOKENS", () => {
+      const error = getMockError();
+      const retryQuota = defaultRetryQuota();
+
+      // release 100 tokens.
+      [...Array(100).keys()].forEach(key => {
+        retryQuota.releaseRetryToken();
+      });
+
+      // 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.retrieveRetryToken(error);
+        }
+      );
+      expect(retryQuota.hasRetryTokens(error)).toBe(false);
+    });
+  });
+});
diff --git a/packages/middleware-retry/src/defaultRetryQuota.ts b/packages/middleware-retry/src/defaultRetryQuota.ts
new file mode 100644
index 0000000000000..a2044b02ffdd6
--- /dev/null
+++ b/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 defaultRetryQuota = (): 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 retrieveRetryToken = (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 releaseRetryToken = (capacityReleaseAmount?: number) => {
+    availableCapacity += capacityReleaseAmount ?? NO_RETRY_INCREMENT;
+    availableCapacity = Math.min(availableCapacity, MAX_CAPACITY);
+  };
+
+  return Object.freeze({
+    hasRetryTokens,
+    retrieveRetryToken,
+    releaseRetryToken
+  });
+};
diff --git a/packages/middleware-retry/src/defaultStrategy.ts b/packages/middleware-retry/src/defaultStrategy.ts
index 7b91e9563589b..e0d8c6f68ec2f 100644
--- a/packages/middleware-retry/src/defaultStrategy.ts
+++ b/packages/middleware-retry/src/defaultStrategy.ts
@@ -66,7 +66,7 @@ export interface StandardRetryStrategyOptions {
 export class StandardRetryStrategy implements RetryStrategy {
   private retryDecider: RetryDecider;
   private delayDecider: DelayDecider;
-  private retryQuota?: RetryQuota;
+  // private retryQuota?: RetryQuota;
 
   constructor(
     public readonly maxAttempts: number,
@@ -74,7 +74,7 @@ export class StandardRetryStrategy implements RetryStrategy {
   ) {
     this.retryDecider = options?.retryDecider ?? defaultRetryDecider;
     this.delayDecider = options?.delayDecider ?? defaultDelayDecider;
-    this.retryQuota = options?.retryQuota;
+    // this.retryQuota = options?.retryQuota;
   }
 
   private shouldRetry(error: SdkError, attempts: number) {

From 592b08c26cf2d46a7ff4374d716283c551bce602 Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Tue, 9 Jun 2020 23:12:38 +0000
Subject: [PATCH 3/9] feat: add defaultRetryQuota to defaultStrategy

---
 .../src/defaultRetryQuota.spec.ts             | 32 ++++++++++---------
 .../middleware-retry/src/defaultRetryQuota.ts | 10 +++---
 .../middleware-retry/src/defaultStrategy.ts   | 19 ++++++++---
 3 files changed, 36 insertions(+), 25 deletions(-)

diff --git a/packages/middleware-retry/src/defaultRetryQuota.spec.ts b/packages/middleware-retry/src/defaultRetryQuota.spec.ts
index eb6e8df7d6395..975cdae4218df 100644
--- a/packages/middleware-retry/src/defaultRetryQuota.spec.ts
+++ b/packages/middleware-retry/src/defaultRetryQuota.spec.ts
@@ -1,4 +1,4 @@
-import { defaultRetryQuota } from "./defaultRetryQuota";
+import { getDefaultRetryQuota } from "./defaultRetryQuota";
 import { SdkError } from "@aws-sdk/smithy-client";
 import {
   INITIAL_RETRY_TOKENS,
@@ -15,10 +15,10 @@ describe("defaultRetryQuota", () => {
     }) as SdkError;
 
   const getDrainedRetryQuota = (targetCapacity: number, error: SdkError) => {
-    const retryQuota = defaultRetryQuota();
+    const retryQuota = getDefaultRetryQuota();
     let availableCapacity = INITIAL_RETRY_TOKENS;
     while (availableCapacity >= targetCapacity) {
-      retryQuota.retrieveRetryToken(error);
+      retryQuota.retrieveRetryTokens(error);
       availableCapacity -= targetCapacity;
     }
     return retryQuota;
@@ -28,11 +28,13 @@ describe("defaultRetryQuota", () => {
     describe("returns true if capacity is available", () => {
       it("when it's TimeoutError", () => {
         const timeoutError = getMockTimeoutError();
-        expect(defaultRetryQuota().hasRetryTokens(timeoutError)).toBe(true);
+        expect(getDefaultRetryQuota().hasRetryTokens(timeoutError)).toBe(true);
       });
 
       it("when it's not TimeoutError", () => {
-        expect(defaultRetryQuota().hasRetryTokens(getMockError())).toBe(true);
+        expect(getDefaultRetryQuota().hasRetryTokens(getMockError())).toBe(
+          true
+        );
       });
     });
 
@@ -58,13 +60,13 @@ describe("defaultRetryQuota", () => {
     describe("returns retry tokens amount if available", () => {
       it("when it's TimeoutError", () => {
         const timeoutError = getMockTimeoutError();
-        expect(defaultRetryQuota().retrieveRetryToken(timeoutError)).toBe(
+        expect(getDefaultRetryQuota().retrieveRetryTokens(timeoutError)).toBe(
           TIMEOUT_RETRY_COST
         );
       });
 
       it("when it's not TimeoutError", () => {
-        expect(defaultRetryQuota().retrieveRetryToken(getMockError())).toBe(
+        expect(getDefaultRetryQuota().retrieveRetryTokens(getMockError())).toBe(
           RETRY_COST
         );
       });
@@ -78,7 +80,7 @@ describe("defaultRetryQuota", () => {
           timeoutError
         );
         expect(() => {
-          retryQuota.retrieveRetryToken(timeoutError);
+          retryQuota.retrieveRetryTokens(timeoutError);
         }).toThrowError(new Error("No retry token available"));
       });
 
@@ -86,7 +88,7 @@ describe("defaultRetryQuota", () => {
         const error = getMockError();
         const retryQuota = getDrainedRetryQuota(RETRY_COST, error);
         expect(() => {
-          retryQuota.retrieveRetryToken(error);
+          retryQuota.retrieveRetryTokens(error);
         }).toThrowError(new Error("No retry token available"));
       });
     });
@@ -101,9 +103,9 @@ describe("defaultRetryQuota", () => {
       expect(retryQuota.hasRetryTokens(error)).toBe(false);
 
       // Release RETRY_COST tokens.
-      retryQuota.releaseRetryToken(RETRY_COST);
+      retryQuota.releaseRetryTokens(RETRY_COST);
       expect(retryQuota.hasRetryTokens(error)).toBe(true);
-      expect(retryQuota.retrieveRetryToken(error)).toBe(RETRY_COST);
+      expect(retryQuota.retrieveRetryTokens(error)).toBe(RETRY_COST);
       expect(retryQuota.hasRetryTokens(error)).toBe(false);
     });
 
@@ -118,7 +120,7 @@ describe("defaultRetryQuota", () => {
         RETRY_COST - (INITIAL_RETRY_TOKENS % RETRY_COST);
       while (tokensReleased < tokensToBeReleased) {
         expect(retryQuota.hasRetryTokens(error)).toBe(false);
-        retryQuota.releaseRetryToken();
+        retryQuota.releaseRetryTokens();
         tokensReleased += NO_RETRY_INCREMENT;
       }
       expect(retryQuota.hasRetryTokens(error)).toBe(true);
@@ -126,11 +128,11 @@ describe("defaultRetryQuota", () => {
 
     it("ensures availableCapacity is maxed at INITIAL_RETRY_TOKENS", () => {
       const error = getMockError();
-      const retryQuota = defaultRetryQuota();
+      const retryQuota = getDefaultRetryQuota();
 
       // release 100 tokens.
       [...Array(100).keys()].forEach(key => {
-        retryQuota.releaseRetryToken();
+        retryQuota.releaseRetryTokens();
       });
 
       // availableCapacity is still maxed at INITIAL_RETRY_TOKENS
@@ -138,7 +140,7 @@ describe("defaultRetryQuota", () => {
       [...Array(Math.floor(INITIAL_RETRY_TOKENS / RETRY_COST)).keys()].forEach(
         key => {
           expect(retryQuota.hasRetryTokens(error)).toBe(true);
-          retryQuota.retrieveRetryToken(error);
+          retryQuota.retrieveRetryTokens(error);
         }
       );
       expect(retryQuota.hasRetryTokens(error)).toBe(false);
diff --git a/packages/middleware-retry/src/defaultRetryQuota.ts b/packages/middleware-retry/src/defaultRetryQuota.ts
index a2044b02ffdd6..20250d304b34d 100644
--- a/packages/middleware-retry/src/defaultRetryQuota.ts
+++ b/packages/middleware-retry/src/defaultRetryQuota.ts
@@ -7,7 +7,7 @@ import {
   NO_RETRY_INCREMENT
 } from "./constants";
 
-export const defaultRetryQuota = (): RetryQuota => {
+export const getDefaultRetryQuota = (): RetryQuota => {
   const MAX_CAPACITY = INITIAL_RETRY_TOKENS;
   let availableCapacity = INITIAL_RETRY_TOKENS;
 
@@ -17,7 +17,7 @@ export const defaultRetryQuota = (): RetryQuota => {
   const hasRetryTokens = (error: SdkError) =>
     getCapacityAmount(error) <= availableCapacity;
 
-  const retrieveRetryToken = (error: SdkError) => {
+  const retrieveRetryTokens = (error: SdkError) => {
     if (!hasRetryTokens(error)) {
       // retryStrategy should stop retrying, and return last error
       throw new Error("No retry token available");
@@ -27,14 +27,14 @@ export const defaultRetryQuota = (): RetryQuota => {
     return capacityAmount;
   };
 
-  const releaseRetryToken = (capacityReleaseAmount?: number) => {
+  const releaseRetryTokens = (capacityReleaseAmount?: number) => {
     availableCapacity += capacityReleaseAmount ?? NO_RETRY_INCREMENT;
     availableCapacity = Math.min(availableCapacity, MAX_CAPACITY);
   };
 
   return Object.freeze({
     hasRetryTokens,
-    retrieveRetryToken,
-    releaseRetryToken
+    retrieveRetryTokens,
+    releaseRetryTokens
   });
 };
diff --git a/packages/middleware-retry/src/defaultStrategy.ts b/packages/middleware-retry/src/defaultStrategy.ts
index e0d8c6f68ec2f..8c7e467b784d2 100644
--- a/packages/middleware-retry/src/defaultStrategy.ts
+++ b/packages/middleware-retry/src/defaultStrategy.ts
@@ -12,6 +12,7 @@ import {
   FinalizeHandlerArguments,
   RetryStrategy
 } from "@aws-sdk/types";
+import { getDefaultRetryQuota } from "./defaultRetryQuota";
 
 /**
  * Determines whether an error is retryable based on the number of retries
@@ -46,12 +47,12 @@ export interface RetryQuota {
    * returns token amount from the retry quota bucket.
    * throws error is retry tokens are not available.
    */
-  retrieveRetryToken: (error: SdkError) => number;
+  retrieveRetryTokens: (error: SdkError) => number;
 
   /**
    * releases tokens back to the retry quota.
    */
-  releaseRetryToken: (releaseCapacityAmount?: number) => void;
+  releaseRetryTokens: (releaseCapacityAmount?: number) => void;
 }
 
 /**
@@ -66,7 +67,7 @@ export interface StandardRetryStrategyOptions {
 export class StandardRetryStrategy implements RetryStrategy {
   private retryDecider: RetryDecider;
   private delayDecider: DelayDecider;
-  // private retryQuota?: RetryQuota;
+  private retryQuota: RetryQuota;
 
   constructor(
     public readonly maxAttempts: number,
@@ -74,22 +75,29 @@ export class StandardRetryStrategy implements RetryStrategy {
   ) {
     this.retryDecider = options?.retryDecider ?? defaultRetryDecider;
     this.delayDecider = options?.delayDecider ?? defaultDelayDecider;
-    // this.retryQuota = options?.retryQuota;
+    this.retryQuota = options?.retryQuota ?? getDefaultRetryQuota();
   }
 
   private shouldRetry(error: SdkError, attempts: number) {
-    return attempts < this.maxAttempts && this.retryDecider(error);
+    return (
+      attempts < this.maxAttempts &&
+      this.retryDecider(error) &&
+      this.retryQuota.hasRetryTokens(error)
+    );
   }
 
   async retry<Input extends object, Ouput extends MetadataBearer>(
     next: FinalizeHandler<Input, Ouput>,
     args: FinalizeHandlerArguments<Input>
   ) {
+    let retryTokenAmount;
     let attempts = 0;
     let totalDelay = 0;
     while (true) {
       try {
         const { response, output } = await next(args);
+
+        this.retryQuota.releaseRetryTokens(retryTokenAmount);
         output.$metadata.attempts = attempts + 1;
         output.$metadata.totalRetryDelay = totalDelay;
 
@@ -97,6 +105,7 @@ export class StandardRetryStrategy implements RetryStrategy {
       } catch (err) {
         attempts++;
         if (this.shouldRetry(err as SdkError, attempts)) {
+          retryTokenAmount = this.retryQuota.retrieveRetryTokens(err);
           const delay = this.delayDecider(
             isThrottlingError(err)
               ? THROTTLING_RETRY_DELAY_BASE

From e106e630d001edc7451eaa82815528af38874c2f Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Wed, 10 Jun 2020 01:04:52 +0000
Subject: [PATCH 4/9] test: defaultRetryQuota in defaultStrategy

---
 .../src/defaultStrategy.spec.ts               | 236 ++++++++++++++----
 1 file changed, 191 insertions(+), 45 deletions(-)

diff --git a/packages/middleware-retry/src/defaultStrategy.spec.ts b/packages/middleware-retry/src/defaultStrategy.spec.ts
index 0b9207a34eebc..c4ca6a00030a2 100644
--- a/packages/middleware-retry/src/defaultStrategy.spec.ts
+++ b/packages/middleware-retry/src/defaultStrategy.spec.ts
@@ -6,6 +6,7 @@ import { isThrottlingError } from "@aws-sdk/service-error-classification";
 import { defaultDelayDecider } from "./delayDecider";
 import { defaultRetryDecider } from "./retryDecider";
 import { StandardRetryStrategy } from "./defaultStrategy";
+import { getDefaultRetryQuota } from "./defaultRetryQuota";
 
 jest.mock("@aws-sdk/service-error-classification", () => ({
   isThrottlingError: jest.fn().mockReturnValue(true)
@@ -19,8 +20,80 @@ jest.mock("./retryDecider", () => ({
   defaultRetryDecider: jest.fn().mockReturnValue(true)
 }));
 
+jest.mock("./defaultRetryQuota", () => {
+  const mockDefaultRetryQuota = {
+    hasRetryTokens: jest.fn().mockReturnValue(true),
+    retrieveRetryTokens: jest.fn().mockReturnValue(1),
+    releaseRetryTokens: jest.fn()
+  };
+  return { getDefaultRetryQuota: () => mockDefaultRetryQuota };
+});
+
 describe("defaultStrategy", () => {
-  const maxAttempts = 2;
+  const maxAttempts = 3;
+
+  const mockSuccessfulOperation = (maxAttempts: number, response?: string) => {
+    const next = jest.fn().mockResolvedValueOnce({
+      response,
+      output: { $metadata: {} }
+    });
+
+    const retryStrategy = new StandardRetryStrategy(maxAttempts);
+    return retryStrategy.retry(next, {} as any);
+  };
+
+  const mockFailedOperation = async (maxAttempts: number, error?: Error) => {
+    const mockError = error ?? new Error("mockError");
+    const next = jest.fn().mockRejectedValue(mockError);
+
+    const retryStrategy = new StandardRetryStrategy(maxAttempts);
+    try {
+      await retryStrategy.retry(next, {} as any);
+    } catch (error) {
+      expect(error).toStrictEqual(mockError);
+    }
+  };
+
+  const mockSuccessAfterOneFail = (
+    maxAttempts: number,
+    error?: Error,
+    response?: string
+  ) => {
+    const mockError = error ?? new Error("mockError");
+    const mockResponse = {
+      response,
+      output: { $metadata: {} }
+    };
+
+    const next = jest
+      .fn()
+      .mockRejectedValueOnce(mockError)
+      .mockResolvedValueOnce(mockResponse);
+
+    const retryStrategy = new StandardRetryStrategy(maxAttempts);
+    return retryStrategy.retry(next, {} as any);
+  };
+
+  const mockSuccessAfterTwoFails = (
+    maxAttempts: number,
+    error?: Error,
+    response?: string
+  ) => {
+    const mockError = error ?? new Error("mockError");
+    const mockResponse = {
+      response,
+      output: { $metadata: {} }
+    };
+
+    const next = jest
+      .fn()
+      .mockRejectedValueOnce(mockError)
+      .mockRejectedValueOnce(mockError)
+      .mockResolvedValueOnce(mockResponse);
+
+    const retryStrategy = new StandardRetryStrategy(maxAttempts);
+    return retryStrategy.retry(next, {} as any);
+  };
 
   afterEach(() => {
     jest.clearAllMocks();
@@ -80,19 +153,8 @@ describe("defaultStrategy", () => {
     ) => {
       (isThrottlingError as jest.Mock).mockReturnValueOnce(mockThrottlingError);
 
-      const mockError = new Error("mockError");
-      const mockResponse = {
-        response: "mockResponse",
-        output: { $metadata: {} }
-      };
-
-      const next = jest
-        .fn()
-        .mockRejectedValueOnce(mockError)
-        .mockResolvedValueOnce(mockResponse);
-
-      const retryStrategy = new StandardRetryStrategy(maxAttempts);
-      await retryStrategy.retry(next, {} as any);
+      const mockError = new Error();
+      await mockSuccessAfterOneFail(maxAttempts, mockError);
 
       expect(isThrottlingError as jest.Mock).toHaveBeenCalledTimes(1);
       expect(isThrottlingError as jest.Mock).toHaveBeenCalledWith(mockError);
@@ -112,19 +174,103 @@ describe("defaultStrategy", () => {
     });
   });
 
-  describe("should not retry", () => {
-    it("when the handler completes successfully", async () => {
-      const mockResponse = {
-        response: "mockResponse",
-        output: { $metadata: {} }
-      };
+  describe("retryQuota", () => {
+    describe("hasRetryTokens", () => {
+      it("not called on successful operation", async () => {
+        const { hasRetryTokens } = getDefaultRetryQuota();
+        await mockSuccessfulOperation(maxAttempts);
+        expect(hasRetryTokens).not.toHaveBeenCalled();
+      });
 
-      const next = jest.fn().mockResolvedValueOnce(mockResponse);
+      it("called once in case of single failure", async () => {
+        const { hasRetryTokens } = getDefaultRetryQuota();
+        await mockSuccessAfterOneFail(maxAttempts);
+        expect(hasRetryTokens).toHaveBeenCalledTimes(1);
+      });
 
-      const retryStrategy = new StandardRetryStrategy(maxAttempts);
-      const { response, output } = await retryStrategy.retry(next, {} as any);
+      it("called once on each retry request", async () => {
+        const { hasRetryTokens } = getDefaultRetryQuota();
+        await mockFailedOperation(maxAttempts);
+        expect(hasRetryTokens).toHaveBeenCalledTimes(maxAttempts - 1);
+      });
+    });
+
+    describe("releaseRetryTokens", () => {
+      it("called once without param on successful operation", async () => {
+        const { releaseRetryTokens } = getDefaultRetryQuota();
+        await mockSuccessfulOperation(maxAttempts);
+        expect(releaseRetryTokens).toHaveBeenCalledTimes(1);
+        expect(releaseRetryTokens).toHaveBeenCalledWith(undefined);
+      });
 
-      expect(response).toStrictEqual(mockResponse.response);
+      it("called once with retryTokenAmount in case of single failure", async () => {
+        const retryTokens = 15;
+        const {
+          releaseRetryTokens,
+          retrieveRetryTokens
+        } = getDefaultRetryQuota();
+        (retrieveRetryTokens as jest.Mock).mockReturnValueOnce(retryTokens);
+
+        await mockSuccessAfterOneFail(maxAttempts);
+        expect(releaseRetryTokens).toHaveBeenCalledTimes(1);
+        expect(releaseRetryTokens).toHaveBeenCalledWith(retryTokens);
+      });
+
+      it("called once with second retryTokenAmount in case of two failures", async () => {
+        const retryTokensFirst = 15;
+        const retryTokensSecond = 30;
+
+        const {
+          releaseRetryTokens,
+          retrieveRetryTokens
+        } = getDefaultRetryQuota();
+
+        (retrieveRetryTokens as jest.Mock)
+          .mockReturnValueOnce(retryTokensFirst)
+          .mockReturnValueOnce(retryTokensSecond);
+
+        await mockSuccessAfterTwoFails(maxAttempts);
+        expect(releaseRetryTokens).toHaveBeenCalledTimes(1);
+        expect(releaseRetryTokens).toHaveBeenCalledWith(retryTokensSecond);
+      });
+
+      it("not called on unsuccessful operation", async () => {
+        const { releaseRetryTokens } = getDefaultRetryQuota();
+        await mockFailedOperation(maxAttempts);
+        expect(releaseRetryTokens).not.toHaveBeenCalled();
+      });
+    });
+
+    describe("retrieveRetryTokens", () => {
+      it("not called on successful operation", async () => {
+        const { retrieveRetryTokens } = getDefaultRetryQuota();
+        await mockSuccessfulOperation(maxAttempts);
+        expect(retrieveRetryTokens).not.toHaveBeenCalled();
+      });
+
+      it("called once in case of single failure", async () => {
+        const { retrieveRetryTokens } = getDefaultRetryQuota();
+        await mockSuccessAfterOneFail(maxAttempts);
+        expect(retrieveRetryTokens).toHaveBeenCalledTimes(1);
+      });
+
+      it("called once on each retry request", async () => {
+        const { retrieveRetryTokens } = getDefaultRetryQuota();
+        await mockFailedOperation(maxAttempts);
+        expect(retrieveRetryTokens).toHaveBeenCalledTimes(maxAttempts - 1);
+      });
+    });
+  });
+
+  describe("should not retry", () => {
+    it("when the handler completes successfully", async () => {
+      const mockResponse = "mockResponse";
+      const { response, output } = await mockSuccessfulOperation(
+        maxAttempts,
+        mockResponse
+      );
+
+      expect(response).toStrictEqual(mockResponse);
       expect(output.$metadata.attempts).toBe(1);
       expect(output.$metadata.totalRetryDelay).toBe(0);
       expect(defaultRetryDecider as jest.Mock).not.toHaveBeenCalled();
@@ -133,35 +279,35 @@ describe("defaultStrategy", () => {
 
     it("when retryDecider returns false", async () => {
       (defaultRetryDecider as jest.Mock).mockReturnValueOnce(false);
-
-      const mockError = new Error("mockError");
-      const next = jest.fn().mockRejectedValueOnce(mockError);
-
-      const retryStrategy = new StandardRetryStrategy(maxAttempts);
-      try {
-        await retryStrategy.retry(next, {} as any);
-      } catch (error) {
-        expect(error).toStrictEqual(mockError);
-      }
-
+      const mockError = new Error();
+      await mockFailedOperation(maxAttempts, mockError);
       expect(defaultRetryDecider as jest.Mock).toHaveBeenCalledTimes(1);
       expect(defaultRetryDecider as jest.Mock).toHaveBeenCalledWith(mockError);
     });
 
-    it("when the the maximum number of attempts is reached", async () => {
-      const mockError = new Error("mockError");
-      const next = jest.fn().mockRejectedValue(mockError);
-
-      const retryStrategy = new StandardRetryStrategy(maxAttempts);
-      try {
-        await retryStrategy.retry(next, {} as any);
-      } catch (error) {
-        expect(error).toStrictEqual(mockError);
-      }
+    it("when the maximum number of attempts is reached", async () => {
+      await mockFailedOperation(maxAttempts);
       expect(defaultRetryDecider as jest.Mock).toHaveBeenCalledTimes(
         maxAttempts - 1
       );
     });
+
+    it("when retryQuota.hasRetryTokens returns false", async () => {
+      const {
+        hasRetryTokens,
+        retrieveRetryTokens,
+        releaseRetryTokens
+      } = getDefaultRetryQuota();
+      (hasRetryTokens as jest.Mock).mockReturnValueOnce(false);
+
+      const mockError = new Error();
+      await mockFailedOperation(maxAttempts, mockError);
+
+      expect(hasRetryTokens).toHaveBeenCalledTimes(1);
+      expect(hasRetryTokens).toHaveBeenCalledWith(mockError);
+      expect(retrieveRetryTokens).not.toHaveBeenCalled();
+      expect(releaseRetryTokens).not.toHaveBeenCalled();
+    });
   });
 
   it("should delay equal to the value returned by delayDecider", async () => {

From dc9db8d14032d34aea379931b2f1dd273d6af255 Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Wed, 10 Jun 2020 02:11:34 +0000
Subject: [PATCH 5/9] test: add options.retryQuota tests

---
 .../src/defaultStrategy.spec.ts               | 22 ++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/packages/middleware-retry/src/defaultStrategy.spec.ts b/packages/middleware-retry/src/defaultStrategy.spec.ts
index c4ca6a00030a2..579ce90478da3 100644
--- a/packages/middleware-retry/src/defaultStrategy.spec.ts
+++ b/packages/middleware-retry/src/defaultStrategy.spec.ts
@@ -5,7 +5,7 @@ import {
 import { isThrottlingError } from "@aws-sdk/service-error-classification";
 import { defaultDelayDecider } from "./delayDecider";
 import { defaultRetryDecider } from "./retryDecider";
-import { StandardRetryStrategy } from "./defaultStrategy";
+import { StandardRetryStrategy, RetryQuota } from "./defaultStrategy";
 import { getDefaultRetryQuota } from "./defaultRetryQuota";
 
 jest.mock("@aws-sdk/service-error-classification", () => ({
@@ -146,6 +146,26 @@ describe("defaultStrategy", () => {
     });
   });
 
+  describe("retryQuota", () => {
+    it("sets getDefaultRetryQuota if options is undefined", () => {
+      const retryStrategy = new StandardRetryStrategy(maxAttempts);
+      expect(retryStrategy["retryQuota"]).toBe(getDefaultRetryQuota());
+    });
+
+    it("sets getDefaultRetryQuota if options.delayDecider undefined", () => {
+      const retryStrategy = new StandardRetryStrategy(maxAttempts, {});
+      expect(retryStrategy["retryQuota"]).toBe(getDefaultRetryQuota());
+    });
+
+    it("sets options.retryQuota if defined", () => {
+      const retryQuota = {} as RetryQuota;
+      const retryStrategy = new StandardRetryStrategy(maxAttempts, {
+        retryQuota
+      });
+      expect(retryStrategy["retryQuota"]).toBe(retryQuota);
+    });
+  });
+
   describe("delayBase passed to delayDecider", () => {
     const testDelayBasePassed = async (
       delayBaseToTest: number,

From df323f4661ee0593f140b2454ca201261e109ada Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Thu, 11 Jun 2020 16:18:53 +0000
Subject: [PATCH 6/9] chore: pass options as second object of utility functions

---
 .../src/defaultStrategy.spec.ts               | 39 ++++++++++---------
 1 file changed, 21 insertions(+), 18 deletions(-)

diff --git a/packages/middleware-retry/src/defaultStrategy.spec.ts b/packages/middleware-retry/src/defaultStrategy.spec.ts
index 579ce90478da3..d8076aa329302 100644
--- a/packages/middleware-retry/src/defaultStrategy.spec.ts
+++ b/packages/middleware-retry/src/defaultStrategy.spec.ts
@@ -32,9 +32,12 @@ jest.mock("./defaultRetryQuota", () => {
 describe("defaultStrategy", () => {
   const maxAttempts = 3;
 
-  const mockSuccessfulOperation = (maxAttempts: number, response?: string) => {
+  const mockSuccessfulOperation = (
+    maxAttempts: number,
+    options?: { mockResponse?: string }
+  ) => {
     const next = jest.fn().mockResolvedValueOnce({
-      response,
+      response: options?.mockResponse,
       output: { $metadata: {} }
     });
 
@@ -42,8 +45,11 @@ describe("defaultStrategy", () => {
     return retryStrategy.retry(next, {} as any);
   };
 
-  const mockFailedOperation = async (maxAttempts: number, error?: Error) => {
-    const mockError = error ?? new Error("mockError");
+  const mockFailedOperation = async (
+    maxAttempts: number,
+    options?: { mockError?: Error }
+  ) => {
+    const mockError = options?.mockError ?? new Error("mockError");
     const next = jest.fn().mockRejectedValue(mockError);
 
     const retryStrategy = new StandardRetryStrategy(maxAttempts);
@@ -56,12 +62,11 @@ describe("defaultStrategy", () => {
 
   const mockSuccessAfterOneFail = (
     maxAttempts: number,
-    error?: Error,
-    response?: string
+    options?: { mockError?: Error; mockResponse?: string }
   ) => {
-    const mockError = error ?? new Error("mockError");
+    const mockError = options?.mockError ?? new Error("mockError");
     const mockResponse = {
-      response,
+      response: options?.mockResponse,
       output: { $metadata: {} }
     };
 
@@ -76,12 +81,11 @@ describe("defaultStrategy", () => {
 
   const mockSuccessAfterTwoFails = (
     maxAttempts: number,
-    error?: Error,
-    response?: string
+    options?: { mockError?: Error; mockResponse?: string }
   ) => {
-    const mockError = error ?? new Error("mockError");
+    const mockError = options?.mockError ?? new Error("mockError");
     const mockResponse = {
-      response,
+      response: options?.mockResponse,
       output: { $metadata: {} }
     };
 
@@ -174,7 +178,7 @@ describe("defaultStrategy", () => {
       (isThrottlingError as jest.Mock).mockReturnValueOnce(mockThrottlingError);
 
       const mockError = new Error();
-      await mockSuccessAfterOneFail(maxAttempts, mockError);
+      await mockSuccessAfterOneFail(maxAttempts, { mockError });
 
       expect(isThrottlingError as jest.Mock).toHaveBeenCalledTimes(1);
       expect(isThrottlingError as jest.Mock).toHaveBeenCalledWith(mockError);
@@ -285,10 +289,9 @@ describe("defaultStrategy", () => {
   describe("should not retry", () => {
     it("when the handler completes successfully", async () => {
       const mockResponse = "mockResponse";
-      const { response, output } = await mockSuccessfulOperation(
-        maxAttempts,
+      const { response, output } = await mockSuccessfulOperation(maxAttempts, {
         mockResponse
-      );
+      });
 
       expect(response).toStrictEqual(mockResponse);
       expect(output.$metadata.attempts).toBe(1);
@@ -300,7 +303,7 @@ describe("defaultStrategy", () => {
     it("when retryDecider returns false", async () => {
       (defaultRetryDecider as jest.Mock).mockReturnValueOnce(false);
       const mockError = new Error();
-      await mockFailedOperation(maxAttempts, mockError);
+      await mockFailedOperation(maxAttempts, { mockError });
       expect(defaultRetryDecider as jest.Mock).toHaveBeenCalledTimes(1);
       expect(defaultRetryDecider as jest.Mock).toHaveBeenCalledWith(mockError);
     });
@@ -321,7 +324,7 @@ describe("defaultStrategy", () => {
       (hasRetryTokens as jest.Mock).mockReturnValueOnce(false);
 
       const mockError = new Error();
-      await mockFailedOperation(maxAttempts, mockError);
+      await mockFailedOperation(maxAttempts, { mockError });
 
       expect(hasRetryTokens).toHaveBeenCalledTimes(1);
       expect(hasRetryTokens).toHaveBeenCalledWith(mockError);

From 01ad868ebf0a49a6a65bbea419cfb2aff0b4948b Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Thu, 11 Jun 2020 16:56:52 +0000
Subject: [PATCH 7/9] test: add more tests for delayDecider

---
 .../src/defaultStrategy.spec.ts               | 137 ++++++++++--------
 1 file changed, 76 insertions(+), 61 deletions(-)

diff --git a/packages/middleware-retry/src/defaultStrategy.spec.ts b/packages/middleware-retry/src/defaultStrategy.spec.ts
index d8076aa329302..8e1ad58e2b7fd 100644
--- a/packages/middleware-retry/src/defaultStrategy.spec.ts
+++ b/packages/middleware-retry/src/defaultStrategy.spec.ts
@@ -57,6 +57,7 @@ describe("defaultStrategy", () => {
       await retryStrategy.retry(next, {} as any);
     } catch (error) {
       expect(error).toStrictEqual(mockError);
+      return error;
     }
   };
 
@@ -110,7 +111,7 @@ describe("defaultStrategy", () => {
     });
   });
 
-  describe("retryDecider", () => {
+  describe("retryDecider init", () => {
     it("sets defaultRetryDecider if options is undefined", () => {
       const retryStrategy = new StandardRetryStrategy(maxAttempts);
       expect(retryStrategy["retryDecider"]).toBe(defaultRetryDecider);
@@ -130,7 +131,7 @@ describe("defaultStrategy", () => {
     });
   });
 
-  describe("delayDecider", () => {
+  describe("delayDecider init", () => {
     it("sets defaultDelayDecider if options is undefined", () => {
       const retryStrategy = new StandardRetryStrategy(maxAttempts);
       expect(retryStrategy["delayDecider"]).toBe(defaultDelayDecider);
@@ -150,7 +151,7 @@ describe("defaultStrategy", () => {
     });
   });
 
-  describe("retryQuota", () => {
+  describe("retryQuota init", () => {
     it("sets getDefaultRetryQuota if options is undefined", () => {
       const retryStrategy = new StandardRetryStrategy(maxAttempts);
       expect(retryStrategy["retryQuota"]).toBe(getDefaultRetryQuota());
@@ -170,31 +171,82 @@ describe("defaultStrategy", () => {
     });
   });
 
-  describe("delayBase passed to delayDecider", () => {
-    const testDelayBasePassed = async (
-      delayBaseToTest: number,
-      mockThrottlingError: boolean
-    ) => {
-      (isThrottlingError as jest.Mock).mockReturnValueOnce(mockThrottlingError);
+  describe("delayDecider", () => {
+    describe("delayBase value passed", () => {
+      const testDelayBasePassed = async (
+        delayBaseToTest: number,
+        mockThrottlingError: boolean
+      ) => {
+        (isThrottlingError as jest.Mock).mockReturnValueOnce(
+          mockThrottlingError
+        );
+
+        const mockError = new Error();
+        await mockSuccessAfterOneFail(maxAttempts, { mockError });
+
+        expect(isThrottlingError as jest.Mock).toHaveBeenCalledTimes(1);
+        expect(isThrottlingError as jest.Mock).toHaveBeenCalledWith(mockError);
+        expect(defaultDelayDecider as jest.Mock).toHaveBeenCalledTimes(1);
+        expect((defaultDelayDecider as jest.Mock).mock.calls[0][0]).toBe(
+          delayBaseToTest
+        );
+      };
+
+      it("should be equal to THROTTLING_RETRY_DELAY_BASE if error is throttling error", async () => {
+        return testDelayBasePassed(THROTTLING_RETRY_DELAY_BASE, true);
+      });
 
-      const mockError = new Error();
-      await mockSuccessAfterOneFail(maxAttempts, { mockError });
-
-      expect(isThrottlingError as jest.Mock).toHaveBeenCalledTimes(1);
-      expect(isThrottlingError as jest.Mock).toHaveBeenCalledWith(mockError);
-      expect(defaultDelayDecider as jest.Mock).toHaveBeenCalledTimes(1);
-      expect(defaultDelayDecider as jest.Mock).toHaveBeenCalledWith(
-        delayBaseToTest,
-        1
-      );
-    };
+      it("should be equal to DEFAULT_RETRY_DELAY_BASE in error is not a throttling error", async () => {
+        return testDelayBasePassed(DEFAULT_RETRY_DELAY_BASE, false);
+      });
+    });
+
+    describe("attempts value passed", () => {
+      it("on successful operation", async () => {
+        await mockSuccessfulOperation(maxAttempts);
+        expect(defaultDelayDecider as jest.Mock).not.toHaveBeenCalled();
+      });
+
+      it("in case of single failure", async () => {
+        await mockSuccessAfterOneFail(maxAttempts);
+        expect(defaultDelayDecider as jest.Mock).toHaveBeenCalledTimes(1);
+        expect((defaultDelayDecider as jest.Mock).mock.calls[0][1]).toBe(1);
+      });
 
-    it("should be equal to THROTTLING_RETRY_DELAY_BASE if error is throttling error", async () => {
-      return testDelayBasePassed(THROTTLING_RETRY_DELAY_BASE, true);
+      it("on all fails", async () => {
+        await mockFailedOperation(maxAttempts);
+        expect(defaultDelayDecider as jest.Mock).toHaveBeenCalledTimes(2);
+        expect((defaultDelayDecider as jest.Mock).mock.calls[0][1]).toBe(1);
+        expect((defaultDelayDecider as jest.Mock).mock.calls[1][1]).toBe(2);
+      });
     });
 
-    it("should be equal to DEFAULT_RETRY_DELAY_BASE in error is not a throttling error", async () => {
-      return testDelayBasePassed(DEFAULT_RETRY_DELAY_BASE, false);
+    it("delay value returned", async () => {
+      jest.spyOn(global, "setTimeout");
+
+      const FIRST_DELAY = 100;
+      const SECOND_DELAY = 200;
+
+      (defaultDelayDecider as jest.Mock)
+        .mockReturnValueOnce(FIRST_DELAY)
+        .mockReturnValueOnce(SECOND_DELAY);
+
+      const maxAttempts = 3;
+      const error = await mockFailedOperation(maxAttempts);
+      expect(error.$metadata.totalRetryDelay).toEqual(
+        FIRST_DELAY + SECOND_DELAY
+      );
+
+      expect(defaultDelayDecider as jest.Mock).toHaveBeenCalledTimes(
+        maxAttempts - 1
+      );
+      expect(setTimeout).toHaveBeenCalledTimes(maxAttempts - 1);
+      expect(((setTimeout as unknown) as jest.Mock).mock.calls[0][1]).toBe(
+        FIRST_DELAY
+      );
+      expect(((setTimeout as unknown) as jest.Mock).mock.calls[1][1]).toBe(
+        SECOND_DELAY
+      );
     });
   });
 
@@ -332,41 +384,4 @@ describe("defaultStrategy", () => {
       expect(releaseRetryTokens).not.toHaveBeenCalled();
     });
   });
-
-  it("should delay equal to the value returned by delayDecider", async () => {
-    jest.spyOn(global, "setTimeout");
-
-    const FIRST_DELAY = 100;
-    const SECOND_DELAY = 200;
-
-    (defaultDelayDecider as jest.Mock)
-      .mockReturnValueOnce(FIRST_DELAY)
-      .mockReturnValueOnce(SECOND_DELAY);
-
-    const mockError = new Error("mockError");
-    const next = jest.fn().mockRejectedValue(mockError);
-
-    const retryStrategy = new StandardRetryStrategy(3);
-    try {
-      await retryStrategy.retry(next, {} as any);
-    } catch (error) {
-      expect(error).toStrictEqual(mockError);
-      expect(error.$metadata.totalRetryDelay).toEqual(
-        FIRST_DELAY + SECOND_DELAY
-      );
-    }
-
-    expect(defaultDelayDecider as jest.Mock).toHaveBeenCalledTimes(2);
-    expect(setTimeout).toHaveBeenCalledTimes(2);
-    expect(setTimeout).toHaveBeenNthCalledWith(
-      1,
-      expect.any(Function),
-      FIRST_DELAY
-    );
-    expect(setTimeout).toHaveBeenNthCalledWith(
-      2,
-      expect.any(Function),
-      SECOND_DELAY
-    );
-  });
 });

From fb44543e63b43873977c8b1a7cbef211e9c724db Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Mon, 15 Jun 2020 19:38:41 +0000
Subject: [PATCH 8/9] chore: modified getDefaultRetryQuota to accept
 initialRetryTokens

---
 .../src/defaultRetryQuota.spec.ts             | 68 +++++++++++++++----
 .../middleware-retry/src/defaultRetryQuota.ts |  9 +--
 .../src/defaultStrategy.spec.ts               | 43 ++++++++----
 .../middleware-retry/src/defaultStrategy.ts   |  6 +-
 4 files changed, 92 insertions(+), 34 deletions(-)

diff --git a/packages/middleware-retry/src/defaultRetryQuota.spec.ts b/packages/middleware-retry/src/defaultRetryQuota.spec.ts
index 975cdae4218df..985c591771acc 100644
--- a/packages/middleware-retry/src/defaultRetryQuota.spec.ts
+++ b/packages/middleware-retry/src/defaultRetryQuota.spec.ts
@@ -14,9 +14,13 @@ describe("defaultRetryQuota", () => {
       name: "TimeoutError"
     }) as SdkError;
 
-  const getDrainedRetryQuota = (targetCapacity: number, error: SdkError) => {
-    const retryQuota = getDefaultRetryQuota();
-    let availableCapacity = INITIAL_RETRY_TOKENS;
+  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;
@@ -24,17 +28,49 @@ describe("defaultRetryQuota", () => {
     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().hasRetryTokens(timeoutError)).toBe(true);
+        expect(
+          getDefaultRetryQuota(INITIAL_RETRY_TOKENS).hasRetryTokens(
+            timeoutError
+          )
+        ).toBe(true);
       });
 
       it("when it's not TimeoutError", () => {
-        expect(getDefaultRetryQuota().hasRetryTokens(getMockError())).toBe(
-          true
-        );
+        expect(
+          getDefaultRetryQuota(INITIAL_RETRY_TOKENS).hasRetryTokens(
+            getMockError()
+          )
+        ).toBe(true);
       });
     });
 
@@ -60,15 +96,19 @@ describe("defaultRetryQuota", () => {
     describe("returns retry tokens amount if available", () => {
       it("when it's TimeoutError", () => {
         const timeoutError = getMockTimeoutError();
-        expect(getDefaultRetryQuota().retrieveRetryTokens(timeoutError)).toBe(
-          TIMEOUT_RETRY_COST
-        );
+        expect(
+          getDefaultRetryQuota(INITIAL_RETRY_TOKENS).retrieveRetryTokens(
+            timeoutError
+          )
+        ).toBe(TIMEOUT_RETRY_COST);
       });
 
       it("when it's not TimeoutError", () => {
-        expect(getDefaultRetryQuota().retrieveRetryTokens(getMockError())).toBe(
-          RETRY_COST
-        );
+        expect(
+          getDefaultRetryQuota(INITIAL_RETRY_TOKENS).retrieveRetryTokens(
+            getMockError()
+          )
+        ).toBe(RETRY_COST);
       });
     });
 
@@ -128,7 +168,7 @@ describe("defaultRetryQuota", () => {
 
     it("ensures availableCapacity is maxed at INITIAL_RETRY_TOKENS", () => {
       const error = getMockError();
-      const retryQuota = getDefaultRetryQuota();
+      const retryQuota = getDefaultRetryQuota(INITIAL_RETRY_TOKENS);
 
       // release 100 tokens.
       [...Array(100).keys()].forEach(key => {
diff --git a/packages/middleware-retry/src/defaultRetryQuota.ts b/packages/middleware-retry/src/defaultRetryQuota.ts
index 20250d304b34d..3a78abf2a3aea 100644
--- a/packages/middleware-retry/src/defaultRetryQuota.ts
+++ b/packages/middleware-retry/src/defaultRetryQuota.ts
@@ -1,15 +1,16 @@
 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;
+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;
diff --git a/packages/middleware-retry/src/defaultStrategy.spec.ts b/packages/middleware-retry/src/defaultStrategy.spec.ts
index 8e1ad58e2b7fd..36039bbdf1784 100644
--- a/packages/middleware-retry/src/defaultStrategy.spec.ts
+++ b/packages/middleware-retry/src/defaultStrategy.spec.ts
@@ -1,6 +1,7 @@
 import {
   DEFAULT_RETRY_DELAY_BASE,
-  THROTTLING_RETRY_DELAY_BASE
+  THROTTLING_RETRY_DELAY_BASE,
+  INITIAL_RETRY_TOKENS
 } from "./constants";
 import { isThrottlingError } from "@aws-sdk/service-error-classification";
 import { defaultDelayDecider } from "./delayDecider";
@@ -154,12 +155,16 @@ describe("defaultStrategy", () => {
   describe("retryQuota init", () => {
     it("sets getDefaultRetryQuota if options is undefined", () => {
       const retryStrategy = new StandardRetryStrategy(maxAttempts);
-      expect(retryStrategy["retryQuota"]).toBe(getDefaultRetryQuota());
+      expect(retryStrategy["retryQuota"]).toBe(
+        getDefaultRetryQuota(INITIAL_RETRY_TOKENS)
+      );
     });
 
     it("sets getDefaultRetryQuota if options.delayDecider undefined", () => {
       const retryStrategy = new StandardRetryStrategy(maxAttempts, {});
-      expect(retryStrategy["retryQuota"]).toBe(getDefaultRetryQuota());
+      expect(retryStrategy["retryQuota"]).toBe(
+        getDefaultRetryQuota(INITIAL_RETRY_TOKENS)
+      );
     });
 
     it("sets options.retryQuota if defined", () => {
@@ -253,19 +258,19 @@ describe("defaultStrategy", () => {
   describe("retryQuota", () => {
     describe("hasRetryTokens", () => {
       it("not called on successful operation", async () => {
-        const { hasRetryTokens } = getDefaultRetryQuota();
+        const { hasRetryTokens } = getDefaultRetryQuota(INITIAL_RETRY_TOKENS);
         await mockSuccessfulOperation(maxAttempts);
         expect(hasRetryTokens).not.toHaveBeenCalled();
       });
 
       it("called once in case of single failure", async () => {
-        const { hasRetryTokens } = getDefaultRetryQuota();
+        const { hasRetryTokens } = getDefaultRetryQuota(INITIAL_RETRY_TOKENS);
         await mockSuccessAfterOneFail(maxAttempts);
         expect(hasRetryTokens).toHaveBeenCalledTimes(1);
       });
 
       it("called once on each retry request", async () => {
-        const { hasRetryTokens } = getDefaultRetryQuota();
+        const { hasRetryTokens } = getDefaultRetryQuota(INITIAL_RETRY_TOKENS);
         await mockFailedOperation(maxAttempts);
         expect(hasRetryTokens).toHaveBeenCalledTimes(maxAttempts - 1);
       });
@@ -273,7 +278,9 @@ describe("defaultStrategy", () => {
 
     describe("releaseRetryTokens", () => {
       it("called once without param on successful operation", async () => {
-        const { releaseRetryTokens } = getDefaultRetryQuota();
+        const { releaseRetryTokens } = getDefaultRetryQuota(
+          INITIAL_RETRY_TOKENS
+        );
         await mockSuccessfulOperation(maxAttempts);
         expect(releaseRetryTokens).toHaveBeenCalledTimes(1);
         expect(releaseRetryTokens).toHaveBeenCalledWith(undefined);
@@ -284,7 +291,7 @@ describe("defaultStrategy", () => {
         const {
           releaseRetryTokens,
           retrieveRetryTokens
-        } = getDefaultRetryQuota();
+        } = getDefaultRetryQuota(INITIAL_RETRY_TOKENS);
         (retrieveRetryTokens as jest.Mock).mockReturnValueOnce(retryTokens);
 
         await mockSuccessAfterOneFail(maxAttempts);
@@ -299,7 +306,7 @@ describe("defaultStrategy", () => {
         const {
           releaseRetryTokens,
           retrieveRetryTokens
-        } = getDefaultRetryQuota();
+        } = getDefaultRetryQuota(INITIAL_RETRY_TOKENS);
 
         (retrieveRetryTokens as jest.Mock)
           .mockReturnValueOnce(retryTokensFirst)
@@ -311,7 +318,9 @@ describe("defaultStrategy", () => {
       });
 
       it("not called on unsuccessful operation", async () => {
-        const { releaseRetryTokens } = getDefaultRetryQuota();
+        const { releaseRetryTokens } = getDefaultRetryQuota(
+          INITIAL_RETRY_TOKENS
+        );
         await mockFailedOperation(maxAttempts);
         expect(releaseRetryTokens).not.toHaveBeenCalled();
       });
@@ -319,19 +328,25 @@ describe("defaultStrategy", () => {
 
     describe("retrieveRetryTokens", () => {
       it("not called on successful operation", async () => {
-        const { retrieveRetryTokens } = getDefaultRetryQuota();
+        const { retrieveRetryTokens } = getDefaultRetryQuota(
+          INITIAL_RETRY_TOKENS
+        );
         await mockSuccessfulOperation(maxAttempts);
         expect(retrieveRetryTokens).not.toHaveBeenCalled();
       });
 
       it("called once in case of single failure", async () => {
-        const { retrieveRetryTokens } = getDefaultRetryQuota();
+        const { retrieveRetryTokens } = getDefaultRetryQuota(
+          INITIAL_RETRY_TOKENS
+        );
         await mockSuccessAfterOneFail(maxAttempts);
         expect(retrieveRetryTokens).toHaveBeenCalledTimes(1);
       });
 
       it("called once on each retry request", async () => {
-        const { retrieveRetryTokens } = getDefaultRetryQuota();
+        const { retrieveRetryTokens } = getDefaultRetryQuota(
+          INITIAL_RETRY_TOKENS
+        );
         await mockFailedOperation(maxAttempts);
         expect(retrieveRetryTokens).toHaveBeenCalledTimes(maxAttempts - 1);
       });
@@ -372,7 +387,7 @@ describe("defaultStrategy", () => {
         hasRetryTokens,
         retrieveRetryTokens,
         releaseRetryTokens
-      } = getDefaultRetryQuota();
+      } = getDefaultRetryQuota(INITIAL_RETRY_TOKENS);
       (hasRetryTokens as jest.Mock).mockReturnValueOnce(false);
 
       const mockError = new Error();
diff --git a/packages/middleware-retry/src/defaultStrategy.ts b/packages/middleware-retry/src/defaultStrategy.ts
index 8c7e467b784d2..cb8c242d13a9b 100644
--- a/packages/middleware-retry/src/defaultStrategy.ts
+++ b/packages/middleware-retry/src/defaultStrategy.ts
@@ -1,6 +1,7 @@
 import {
   DEFAULT_RETRY_DELAY_BASE,
-  THROTTLING_RETRY_DELAY_BASE
+  THROTTLING_RETRY_DELAY_BASE,
+  INITIAL_RETRY_TOKENS
 } from "./constants";
 import { defaultDelayDecider } from "./delayDecider";
 import { defaultRetryDecider } from "./retryDecider";
@@ -75,7 +76,8 @@ export class StandardRetryStrategy implements RetryStrategy {
   ) {
     this.retryDecider = options?.retryDecider ?? defaultRetryDecider;
     this.delayDecider = options?.delayDecider ?? defaultDelayDecider;
-    this.retryQuota = options?.retryQuota ?? getDefaultRetryQuota();
+    this.retryQuota =
+      options?.retryQuota ?? getDefaultRetryQuota(INITIAL_RETRY_TOKENS);
   }
 
   private shouldRetry(error: SdkError, attempts: number) {

From 8a1c222d6a81e216cae6871f706698aa68620b3c Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Mon, 15 Jun 2020 20:46:09 +0000
Subject: [PATCH 9/9] test: retryQuota.hasRetryTokens false after first retry

---
 .../src/defaultStrategy.spec.ts               | 50 ++++++++++++++-----
 1 file changed, 37 insertions(+), 13 deletions(-)

diff --git a/packages/middleware-retry/src/defaultStrategy.spec.ts b/packages/middleware-retry/src/defaultStrategy.spec.ts
index 36039bbdf1784..8809f55b6dce6 100644
--- a/packages/middleware-retry/src/defaultStrategy.spec.ts
+++ b/packages/middleware-retry/src/defaultStrategy.spec.ts
@@ -382,21 +382,45 @@ describe("defaultStrategy", () => {
       );
     });
 
-    it("when retryQuota.hasRetryTokens returns false", async () => {
-      const {
-        hasRetryTokens,
-        retrieveRetryTokens,
-        releaseRetryTokens
-      } = getDefaultRetryQuota(INITIAL_RETRY_TOKENS);
-      (hasRetryTokens as jest.Mock).mockReturnValueOnce(false);
+    describe("when retryQuota.hasRetryTokens returns false", () => {
+      it("in the first request", async () => {
+        const {
+          hasRetryTokens,
+          retrieveRetryTokens,
+          releaseRetryTokens
+        } = getDefaultRetryQuota(INITIAL_RETRY_TOKENS);
+        (hasRetryTokens as jest.Mock).mockReturnValueOnce(false);
 
-      const mockError = new Error();
-      await mockFailedOperation(maxAttempts, { mockError });
+        const mockError = new Error();
+        await mockFailedOperation(maxAttempts, { mockError });
+
+        expect(hasRetryTokens).toHaveBeenCalledTimes(1);
+        expect(hasRetryTokens).toHaveBeenCalledWith(mockError);
+        expect(retrieveRetryTokens).not.toHaveBeenCalled();
+        expect(releaseRetryTokens).not.toHaveBeenCalled();
+      });
 
-      expect(hasRetryTokens).toHaveBeenCalledTimes(1);
-      expect(hasRetryTokens).toHaveBeenCalledWith(mockError);
-      expect(retrieveRetryTokens).not.toHaveBeenCalled();
-      expect(releaseRetryTokens).not.toHaveBeenCalled();
+      it("after the first retry", async () => {
+        const {
+          hasRetryTokens,
+          retrieveRetryTokens,
+          releaseRetryTokens
+        } = getDefaultRetryQuota(INITIAL_RETRY_TOKENS);
+        (hasRetryTokens as jest.Mock)
+          .mockReturnValueOnce(true)
+          .mockReturnValueOnce(false);
+
+        const mockError = new Error();
+        await mockFailedOperation(maxAttempts, { mockError });
+
+        expect(hasRetryTokens).toHaveBeenCalledTimes(2);
+        [1, 2].forEach(n => {
+          expect(hasRetryTokens).toHaveBeenNthCalledWith(n, mockError);
+        });
+        expect(retrieveRetryTokens).toHaveBeenCalledTimes(1);
+        expect(retrieveRetryTokens).toHaveBeenCalledWith(mockError);
+        expect(releaseRetryTokens).not.toHaveBeenCalled();
+      });
     });
   });
 });