From fdca8d0c6746cd0035f8fce011459a131933f022 Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Wed, 2 Jun 2021 05:33:19 +0000
Subject: [PATCH] chore(middleware-retry): add options to StandardRetryStrategy

---
 .../middleware-retry/src/defaultRetryQuota.ts | 29 +++++++++++++++++--
 1 file changed, 26 insertions(+), 3 deletions(-)

diff --git a/packages/middleware-retry/src/defaultRetryQuota.ts b/packages/middleware-retry/src/defaultRetryQuota.ts
index 31010ef63664..822136192de0 100644
--- a/packages/middleware-retry/src/defaultRetryQuota.ts
+++ b/packages/middleware-retry/src/defaultRetryQuota.ts
@@ -3,11 +3,34 @@ import { SdkError } from "@aws-sdk/smithy-client";
 import { NO_RETRY_INCREMENT, RETRY_COST, TIMEOUT_RETRY_COST } from "./constants";
 import { RetryQuota } from "./types";
 
-export const getDefaultRetryQuota = (initialRetryTokens: number): RetryQuota => {
+export interface DefaultRetryQuotaOptions {
+  /**
+   * The total amount of retry token to be incremented from retry token balance
+   * if an SDK operation invocation succeeds without requiring a retry request.
+   */
+  noRetryIncrement?: number;
+
+  /**
+   * The total amount of retry tokens to be decremented from retry token balance.
+   */
+  retryCost?: number;
+
+  /**
+   * The total amount of retry tokens to be decremented from retry token balance
+   * when a throttling error is encountered.
+   */
+  timeoutRetryCost?: number;
+}
+
+export const getDefaultRetryQuota = (initialRetryTokens: number, options?: DefaultRetryQuotaOptions): RetryQuota => {
   const MAX_CAPACITY = initialRetryTokens;
+  const noRetryIncrement = options?.noRetryIncrement ?? NO_RETRY_INCREMENT;
+  const retryCost = options?.retryCost ?? RETRY_COST;
+  const timeoutRetryCost = options?.timeoutRetryCost ?? TIMEOUT_RETRY_COST;
+
   let availableCapacity = initialRetryTokens;
 
-  const getCapacityAmount = (error: SdkError) => (error.name === "TimeoutError" ? TIMEOUT_RETRY_COST : RETRY_COST);
+  const getCapacityAmount = (error: SdkError) => (error.name === "TimeoutError" ? timeoutRetryCost : retryCost);
 
   const hasRetryTokens = (error: SdkError) => getCapacityAmount(error) <= availableCapacity;
 
@@ -22,7 +45,7 @@ export const getDefaultRetryQuota = (initialRetryTokens: number): RetryQuota =>
   };
 
   const releaseRetryTokens = (capacityReleaseAmount?: number) => {
-    availableCapacity += capacityReleaseAmount ?? NO_RETRY_INCREMENT;
+    availableCapacity += capacityReleaseAmount ?? noRetryIncrement;
     availableCapacity = Math.min(availableCapacity, MAX_CAPACITY);
   };