fix(util-retry): make standard retry tokens immutable (#4755)

* fix(util-retry): use token instances with shared token availability

* fix(util-retry): remove mutations in standard retry tokens

* fix(util-retry): make standard retry token immutable, deprecate some methods on standard retry token

* fix(util-retry): remove deprecated parts of StandardRetryToken

* chore: remove redundant variables retryCost and timeoutRetryCost

* fix(util-retry): packages/util-retry/src/StandardRetryStrategy.ts

Co-authored-by: Trivikram Kamat <[email protected]>

* fix(util-retry): remove unnecessary input

---------

Co-authored-by: Trivikram Kamat <[email protected]>

Author: kuhe

packages/middleware-retry/src/types.ts

@@ -4,7 +4,7 @@ import { SdkError } from "@aws-sdk/types";
  * Determines whether an error is retryable based on the number of retries
  * already attempted, the HTTP status code, and the error received (if any).
  *
- * @param error         The error encountered.
+ * @param error - The error encountered.
  */
 export interface RetryDecider {
   (error: SdkError): boolean;
@@ -13,8 +13,8 @@ export interface RetryDecider {
 /**
  * Determines the number of milliseconds to wait before retrying an action.
  *
- * @param delayBase The base delay (in milliseconds).
- * @param attempts  The number of times the action has already been tried.
+ * @param delayBase - The base delay (in milliseconds).
+ * @param attempts - The number of times the action has already been tried.
  */
 export interface DelayDecider {
   (delayBase: number, attempts: number): number;

packages/types/src/retry.ts

@@ -95,26 +95,9 @@ export interface RetryToken {
  */
 export interface StandardRetryToken extends RetryToken {
   /**
-   * @returns wheather token has remaining tokens.
+   * @returns the cost of the last retry attempt.
    */
-  hasRetryTokens(errorType: RetryErrorType): boolean;
-
-  /**
-   * @returns the number of available tokens.
-   */
-  getRetryTokenCount(errorInfo: RetryErrorInfo): number;
-
-  /**
-   * @returns the cost of the last retry attemp.
-   */
-  getLastRetryCost(): number | undefined;
-
-  /**
-   * Releases a number of tokens.
-   *
-   * @param amount - of tokens to release.
-   */
-  releaseRetryTokens(amount?: number): void;
+  getRetryCost(): number | undefined;
 }
 
 /**

packages/util-retry/src/AdaptiveRetryStrategy.spec.ts

@@ -1,4 +1,4 @@
-import { RetryErrorInfo, RetryErrorType, StandardRetryToken } from "@aws-sdk/types";
+import { RetryErrorInfo, StandardRetryToken } from "@aws-sdk/types";
 
 import { AdaptiveRetryStrategy } from "./AdaptiveRetryStrategy";
 import { RETRY_MODES } from "./config";
@@ -17,10 +17,7 @@ describe(AdaptiveRetryStrategy.name, () => {
     updateClientSendingRate: jest.fn(),
   };
   const mockRetryToken: StandardRetryToken = {
-    hasRetryTokens: (errorType: RetryErrorType) => true,
-    getLastRetryCost: () => 1,
-    getRetryTokenCount: (errorInfo: RetryErrorInfo) => 1,
-    releaseRetryTokens: (amount: number) => {},
+    getRetryCost: () => 1,
     getRetryCount: () => 1,
     getRetryDelay: () => 1,
   };

packages/util-retry/src/ConfiguredRetryStrategy.spec.ts

@@ -11,7 +11,7 @@ describe(ConfiguredRetryStrategy.name, () => {
       errorType: "TRANSIENT",
     });
 
-    expect(retryToken.getRetryCount()).toBe(4);
-    expect(retryToken.getRetryDelay()).toBe(4000);
+    expect(retryToken.getRetryCount()).toBe(5);
+    expect(retryToken.getRetryDelay()).toBe(5000);
   });
 });

packages/util-retry/src/StandardRetryStrategy.spec.ts

@@ -2,7 +2,7 @@ import { RetryErrorInfo, RetryErrorType } from "@aws-sdk/types";
 
 import { RETRY_MODES } from "./config";
 import { DEFAULT_RETRY_DELAY_BASE, INITIAL_RETRY_TOKENS } from "./constants";
-import { getDefaultRetryToken } from "./defaultRetryToken";
+import { createDefaultRetryToken } from "./defaultRetryToken";
 import { StandardRetryStrategy } from "./StandardRetryStrategy";
 
 jest.mock("./defaultRetryToken");
@@ -18,7 +18,7 @@ describe(StandardRetryStrategy.name, () => {
   const errorInfo = { errorType: "TRANSIENT" } as RetryErrorInfo;
 
   beforeEach(() => {
-    (getDefaultRetryToken as jest.Mock).mockReturnValue(mockRetryToken);
+    (createDefaultRetryToken as jest.Mock).mockReturnValue(mockRetryToken);
   });
 
   afterEach(() => {
@@ -37,48 +37,40 @@ describe(StandardRetryStrategy.name, () => {
     expect(retryStrategy.mode).toStrictEqual(RETRY_MODES.STANDARD);
   });
 
-  describe("retryToken init", () => {
-    it("sets retryToken", () => {
-      const retryStrategy = new StandardRetryStrategy(() => Promise.resolve(maxAttempts));
-      expect(retryStrategy["retryToken"]).toBe(getDefaultRetryToken(INITIAL_RETRY_TOKENS, DEFAULT_RETRY_DELAY_BASE));
-    });
-  });
-
   describe("acquireInitialRetryToken", () => {
     it("returns default retryToken", async () => {
       const retryStrategy = new StandardRetryStrategy(() => Promise.resolve(maxAttempts));
       const retryToken = await retryStrategy.acquireInitialRetryToken(retryTokenScope);
-      expect(retryToken).toEqual(getDefaultRetryToken(INITIAL_RETRY_TOKENS, DEFAULT_RETRY_DELAY_BASE));
+      expect(retryToken).toEqual(
+        createDefaultRetryToken({
+          retryDelay: DEFAULT_RETRY_DELAY_BASE,
+          retryCount: 0,
+        })
+      );
     });
   });
 
   describe("refreshRetryTokenForRetry", () => {
     it("refreshes the token", async () => {
-      const getRetryTokenCount = jest.fn().mockReturnValue(1);
       const getRetryCount = jest.fn().mockReturnValue(0);
       const hasRetryTokens = jest.fn().mockReturnValue(true);
       const mockRetryToken = {
         getRetryCount,
-        getRetryTokenCount,
         hasRetryTokens,
       };
-      (getDefaultRetryToken as jest.Mock).mockReturnValue(mockRetryToken);
+      (createDefaultRetryToken as jest.Mock).mockReturnValue(mockRetryToken);
       const retryStrategy = new StandardRetryStrategy(() => Promise.resolve(maxAttempts));
       const token = await retryStrategy.acquireInitialRetryToken(retryTokenScope);
-      const refreshedToken = await retryStrategy.refreshRetryTokenForRetry(token, errorInfo);
-      expect(getRetryTokenCount).toHaveBeenCalledTimes(1);
-      expect(getRetryTokenCount).toHaveBeenCalledWith(errorInfo);
-      expect(getRetryCount).toHaveBeenCalledTimes(1);
-      expect(hasRetryTokens).toHaveBeenCalledTimes(1);
-      expect(hasRetryTokens).toHaveBeenCalledWith(errorInfo.errorType);
+      await retryStrategy.refreshRetryTokenForRetry(token, errorInfo);
+      expect(getRetryCount).toHaveBeenCalledTimes(3);
     });
 
     it("throws when attempts exceeds maxAttempts", async () => {
       const mockRetryToken = {
         getRetryCount: () => 2,
         getRetryTokenCount: (errorInfo: any) => 1,
       };
-      (getDefaultRetryToken as jest.Mock).mockReturnValue(mockRetryToken);
+      (createDefaultRetryToken as jest.Mock).mockReturnValue(mockRetryToken);
       const retryStrategy = new StandardRetryStrategy(() => Promise.resolve(1));
       const token = await retryStrategy.acquireInitialRetryToken(retryTokenScope);
       try {
@@ -93,7 +85,7 @@ describe(StandardRetryStrategy.name, () => {
         getRetryCount: () => 5,
         getRetryTokenCount: (errorInfo: any) => 1,
       };
-      (getDefaultRetryToken as jest.Mock).mockReturnValue(mockRetryToken);
+      (createDefaultRetryToken as jest.Mock).mockReturnValue(mockRetryToken);
       const retryStrategy = new StandardRetryStrategy(() => Promise.resolve(5));
       const token = await retryStrategy.acquireInitialRetryToken(retryTokenScope);
       try {
@@ -109,7 +101,7 @@ describe(StandardRetryStrategy.name, () => {
         getRetryTokenCount: (errorInfo: any) => 1,
         hasRetryTokens: (errorType: RetryErrorType) => false,
       };
-      (getDefaultRetryToken as jest.Mock).mockReturnValue(mockRetryToken);
+      (createDefaultRetryToken as jest.Mock).mockReturnValue(mockRetryToken);
       const retryStrategy = new StandardRetryStrategy(() => Promise.resolve(maxAttempts));
       const token = await retryStrategy.acquireInitialRetryToken(retryTokenScope);
       try {
@@ -125,7 +117,7 @@ describe(StandardRetryStrategy.name, () => {
         getRetryTokenCount: (errorInfo: any) => 1,
         hasRetryTokens: (errorType: RetryErrorType) => true,
       };
-      (getDefaultRetryToken as jest.Mock).mockReturnValue(mockRetryToken);
+      (createDefaultRetryToken as jest.Mock).mockReturnValue(mockRetryToken);
       const retryStrategy = new StandardRetryStrategy(() => Promise.resolve(maxAttempts));
       const token = await retryStrategy.acquireInitialRetryToken(retryTokenScope);
       const errorInfo = {
@@ -137,24 +129,5 @@ describe(StandardRetryStrategy.name, () => {
         expect(error).toStrictEqual(noRetryTokenAvailableError);
       }
     });
-
-    describe("recordSuccess", () => {
-      it("releases tokens", async () => {
-        const retryCost = 1;
-        const releaseRetryTokens = jest.fn();
-        const getLastRetryCost = jest.fn().mockReturnValue(retryCost);
-        const mockRetryToken = {
-          releaseRetryTokens,
-          getLastRetryCost,
-        };
-        (getDefaultRetryToken as jest.Mock).mockReturnValue(mockRetryToken);
-        const retryStrategy = new StandardRetryStrategy(() => Promise.resolve(maxAttempts));
-        const token = await retryStrategy.acquireInitialRetryToken(retryTokenScope);
-        retryStrategy.recordSuccess(token);
-        expect(releaseRetryTokens).toHaveBeenCalledTimes(1);
-        expect(releaseRetryTokens).toHaveBeenCalledWith(retryCost);
-        expect(getLastRetryCost).toHaveBeenCalledTimes(1);
-      });
-    });
   });
 });

packages/util-retry/src/StandardRetryStrategy.ts

@@ -1,47 +1,82 @@
 import { Provider, RetryErrorInfo, RetryErrorType, RetryStrategyV2, StandardRetryToken } from "@aws-sdk/types";
 
 import { DEFAULT_MAX_ATTEMPTS, RETRY_MODES } from "./config";
-import { DEFAULT_RETRY_DELAY_BASE, INITIAL_RETRY_TOKENS } from "./constants";
-import { getDefaultRetryToken } from "./defaultRetryToken";
+import {
+  DEFAULT_RETRY_DELAY_BASE,
+  INITIAL_RETRY_TOKENS,
+  NO_RETRY_INCREMENT,
+  RETRY_COST,
+  THROTTLING_RETRY_DELAY_BASE,
+  TIMEOUT_RETRY_COST,
+} from "./constants";
+import { getDefaultRetryBackoffStrategy } from "./defaultRetryBackoffStrategy";
+import { createDefaultRetryToken } from "./defaultRetryToken";
 
 /**
  * @public
  */
 export class StandardRetryStrategy implements RetryStrategyV2 {
   public readonly mode: string = RETRY_MODES.STANDARD;
-  private retryToken: StandardRetryToken;
+  private capacity: number = INITIAL_RETRY_TOKENS;
+  private readonly retryBackoffStrategy = getDefaultRetryBackoffStrategy();
   private readonly maxAttemptsProvider: Provider<number>;
 
   constructor(maxAttempts: number);
   constructor(maxAttemptsProvider: Provider<number>);
   constructor(private readonly maxAttempts: number | Provider<number>) {
-    this.retryToken = getDefaultRetryToken(INITIAL_RETRY_TOKENS, DEFAULT_RETRY_DELAY_BASE);
     this.maxAttemptsProvider = typeof maxAttempts === "function" ? maxAttempts : async () => maxAttempts;
   }
 
   public async acquireInitialRetryToken(retryTokenScope: string): Promise<StandardRetryToken> {
-    return this.retryToken;
+    return createDefaultRetryToken({
+      retryDelay: DEFAULT_RETRY_DELAY_BASE,
+      retryCount: 0,
+    });
   }
 
   public async refreshRetryTokenForRetry(
-    tokenToRenew: StandardRetryToken,
+    token: StandardRetryToken,
     errorInfo: RetryErrorInfo
   ): Promise<StandardRetryToken> {
     const maxAttempts = await this.getMaxAttempts();
 
-    if (this.shouldRetry(tokenToRenew, errorInfo, maxAttempts)) {
-      tokenToRenew.getRetryTokenCount(errorInfo);
-      return tokenToRenew;
+    if (this.shouldRetry(token, errorInfo, maxAttempts)) {
+      const errorType = errorInfo.errorType;
+      this.retryBackoffStrategy.setDelayBase(
+        errorType === "THROTTLING" ? THROTTLING_RETRY_DELAY_BASE : DEFAULT_RETRY_DELAY_BASE
+      );
+
+      const delayFromErrorType = this.retryBackoffStrategy.computeNextBackoffDelay(token.getRetryCount());
+      const retryDelay = errorInfo.retryAfterHint
+        ? Math.max(errorInfo.retryAfterHint.getTime() - Date.now() || 0, delayFromErrorType)
+        : delayFromErrorType;
+
+      const capacityCost = this.getCapacityCost(errorType);
+      this.capacity -= capacityCost;
+      return createDefaultRetryToken({
+        retryDelay,
+        retryCount: token.getRetryCount() + 1,
+        retryCost: capacityCost,
+      });
     }
+
     throw new Error("No retry token available");
   }
 
   public recordSuccess(token: StandardRetryToken): void {
-    this.retryToken.releaseRetryTokens(token.getLastRetryCost());
+    this.capacity = Math.max(INITIAL_RETRY_TOKENS, this.capacity + (token.getRetryCost() ?? NO_RETRY_INCREMENT));
+  }
+
+  /**
+   * @returns the current available retry capacity.
+   *
+   * This number decreases when retries are executed and refills when requests or retries succeed.
+   */
+  public getCapacity(): number {
+    return this.capacity;
   }
 
   private async getMaxAttempts() {
-    let maxAttempts: number;
     try {
       return await this.maxAttemptsProvider();
     } catch (error) {
@@ -52,13 +87,18 @@ export class StandardRetryStrategy implements RetryStrategyV2 {
 
   private shouldRetry(tokenToRenew: StandardRetryToken, errorInfo: RetryErrorInfo, maxAttempts: number): boolean {
     const attempts = tokenToRenew.getRetryCount();
+
     return (
       attempts < maxAttempts &&
-      tokenToRenew.hasRetryTokens(errorInfo.errorType) &&
+      this.capacity >= this.getCapacityCost(errorInfo.errorType) &&
       this.isRetryableError(errorInfo.errorType)
     );
   }
 
+  private getCapacityCost(errorType: RetryErrorType) {
+    return errorType === "TRANSIENT" ? TIMEOUT_RETRY_COST : RETRY_COST;
+  }
+
   private isRetryableError(errorType: RetryErrorType): boolean {
     return errorType === "THROTTLING" || errorType === "TRANSIENT";
   }