feat(middleware-retry): add Adaptive Retry Strategy

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

@@ -0,0 +1,97 @@
+import { AdaptiveRetryStrategy } from "./AdaptiveRetryStrategy";
+import { RETRY_MODES } from "./config";
+import { DefaultRateLimiter } from "./DefaultRateLimiter";
+import { StandardRetryStrategy } from "./StandardRetryStrategy";
+import { RateLimiter, RetryQuota } from "./types";
+
+jest.mock("./StandardRetryStrategy");
+jest.mock("./DefaultRateLimiter");
+
+describe(AdaptiveRetryStrategy.name, () => {
+  const maxAttemptsProvider = jest.fn();
+  const mockDefaultRateLimiter = {
+    getSendToken: jest.fn(),
+    updateClientSendingRate: jest.fn(),
+  };
+
+  beforeEach(() => {
+    (DefaultRateLimiter as jest.Mock).mockReturnValue(mockDefaultRateLimiter);
+  });
+
+  afterEach(() => {
+    jest.clearAllMocks();
+  });
+
+  describe("constructor", () => {
+    it("calls super constructor", () => {
+      const retryDecider = jest.fn();
+      const delayDecider = jest.fn();
+      const retryQuota = {} as RetryQuota;
+      const rateLimiter = {} as RateLimiter;
+
+      new AdaptiveRetryStrategy(maxAttemptsProvider, {
+        retryDecider,
+        delayDecider,
+        retryQuota,
+        rateLimiter,
+      });
+      expect(StandardRetryStrategy).toHaveBeenCalledWith(maxAttemptsProvider, {
+        retryDecider,
+        delayDecider,
+        retryQuota,
+      });
+    });
+
+    it(`sets mode=${RETRY_MODES.ADAPTIVE}`, () => {
+      const retryStrategy = new AdaptiveRetryStrategy(maxAttemptsProvider);
+      expect(retryStrategy.mode).toStrictEqual(RETRY_MODES.ADAPTIVE);
+    });
+
+    describe("rateLimiter init", () => {
+      it("sets getDefaultrateLimiter if options is undefined", () => {
+        const retryStrategy = new AdaptiveRetryStrategy(maxAttemptsProvider);
+        expect(retryStrategy["rateLimiter"]).toBe(mockDefaultRateLimiter);
+      });
+
+      it("sets getDefaultrateLimiter if options.delayDecider undefined", () => {
+        const retryStrategy = new AdaptiveRetryStrategy(maxAttemptsProvider, {});
+        expect(retryStrategy["rateLimiter"]).toBe(mockDefaultRateLimiter);
+      });
+
+      it("sets options.rateLimiter if defined", () => {
+        const rateLimiter = {} as RateLimiter;
+        const retryStrategy = new AdaptiveRetryStrategy(maxAttemptsProvider, {
+          rateLimiter,
+        });
+        expect(retryStrategy["rateLimiter"]).toBe(rateLimiter);
+      });
+    });
+  });
+
+  describe("retry", () => {
+    const mockedSuperRetry = jest.spyOn(StandardRetryStrategy.prototype, "retry");
+
+    beforeEach(async () => {
+      const next = jest.fn();
+      const retryStrategy = new AdaptiveRetryStrategy(maxAttemptsProvider);
+      await retryStrategy.retry(next, { request: { headers: {} } } as any);
+      expect(mockedSuperRetry).toHaveBeenCalledTimes(1);
+    });
+
+    afterEach(() => {
+      jest.clearAllMocks();
+    });
+
+    it("calls rateLimiter.getSendToken in beforeRequest", async () => {
+      expect(mockDefaultRateLimiter.getSendToken).toHaveBeenCalledTimes(0);
+      await mockedSuperRetry.mock.calls[0][2].beforeRequest();
+      expect(mockDefaultRateLimiter.getSendToken).toHaveBeenCalledTimes(1);
+    });
+
+    it("calls rateLimiter.updateClientSendingRate in afterRequest", async () => {
+      expect(mockDefaultRateLimiter.updateClientSendingRate).toHaveBeenCalledTimes(0);
+      await mockedSuperRetry.mock.calls[0][2].afterRequest();
+      expect(mockDefaultRateLimiter.updateClientSendingRate).toHaveBeenCalledTimes(1);
+    });
+  });
+});

packages/middleware-retry/src/AdaptiveRetryStrategy.ts

@@ -0,0 +1,38 @@
+import { FinalizeHandler, FinalizeHandlerArguments, MetadataBearer, Provider } from "@aws-sdk/types";
+
+import { RETRY_MODES } from "./config";
+import { DefaultRateLimiter } from "./DefaultRateLimiter";
+import { StandardRetryStrategy, StandardRetryStrategyOptions } from "./StandardRetryStrategy";
+import { RateLimiter } from "./types";
+
+/**
+ * Strategy options to be passed to AdaptiveRetryStrategy
+ */
+export interface AdaptiveRetryStrategyOptions extends StandardRetryStrategyOptions {
+  rateLimiter?: RateLimiter;
+}
+
+export class AdaptiveRetryStrategy extends StandardRetryStrategy {
+  private rateLimiter: RateLimiter;
+
+  constructor(maxAttemptsProvider: Provider<number>, options?: AdaptiveRetryStrategyOptions) {
+    const { rateLimiter, ...superOptions } = options ?? {};
+    super(maxAttemptsProvider, superOptions);
+    this.rateLimiter = rateLimiter ?? new DefaultRateLimiter();
+    this.mode = RETRY_MODES.ADAPTIVE;
+  }
+
+  async retry<Input extends object, Ouput extends MetadataBearer>(
+    next: FinalizeHandler<Input, Ouput>,
+    args: FinalizeHandlerArguments<Input>
+  ) {
+    return super.retry(next, args, {
+      beforeRequest: async () => {
+        return this.rateLimiter.getSendToken();
+      },
+      afterRequest: (response: any) => {
+        this.rateLimiter.updateClientSendingRate(response);
+      },
+    });
+  }
+}

packages/middleware-retry/src/defaultStrategy.spec.ts → packages/middleware-retry/src/StandardRetryStrategy.spec.ts

@@ -2,12 +2,12 @@ import { HttpRequest } from "@aws-sdk/protocol-http";
 import { isThrottlingError } from "@aws-sdk/service-error-classification";
 import { v4 } from "uuid";
 
-import { DEFAULT_MAX_ATTEMPTS } from "./configurations";
+import { DEFAULT_MAX_ATTEMPTS, RETRY_MODES } from "./config";
 import { DEFAULT_RETRY_DELAY_BASE, INITIAL_RETRY_TOKENS, THROTTLING_RETRY_DELAY_BASE } from "./constants";
 import { getDefaultRetryQuota } from "./defaultRetryQuota";
-import { StandardRetryStrategy } from "./defaultStrategy";
 import { defaultDelayDecider } from "./delayDecider";
 import { defaultRetryDecider } from "./retryDecider";
+import { StandardRetryStrategy } from "./StandardRetryStrategy";
 import { RetryQuota } from "./types";
 
 jest.mock("@aws-sdk/service-error-classification");
@@ -102,6 +102,11 @@ describe("defaultStrategy", () => {
     });
   });
 
+  it(`sets mode=${RETRY_MODES.STANDARD}`, () => {
+    const retryStrategy = new StandardRetryStrategy(() => Promise.resolve(maxAttempts));
+    expect(retryStrategy.mode).toStrictEqual(RETRY_MODES.STANDARD);
+  });
+
   it("handles non-standard errors", () => {
     const nonStandardErrors = [undefined, "foo", { foo: "bar" }, 123, false, null];
     const maxAttempts = 1;

packages/middleware-retry/src/defaultStrategy.ts → packages/middleware-retry/src/StandardRetryStrategy.ts

@@ -4,7 +4,7 @@ import { SdkError } from "@aws-sdk/smithy-client";
 import { FinalizeHandler, FinalizeHandlerArguments, MetadataBearer, Provider, RetryStrategy } from "@aws-sdk/types";
 import { v4 } from "uuid";
 
-import { DEFAULT_MAX_ATTEMPTS, DEFAULT_RETRY_MODE } from "./configurations";
+import { DEFAULT_MAX_ATTEMPTS, RETRY_MODES } from "./config";
 import {
   DEFAULT_RETRY_DELAY_BASE,
   INITIAL_RETRY_TOKENS,
@@ -30,7 +30,7 @@ export class StandardRetryStrategy implements RetryStrategy {
   private retryDecider: RetryDecider;
   private delayDecider: DelayDecider;
   private retryQuota: RetryQuota;
-  public readonly mode = DEFAULT_RETRY_MODE;
+  public mode: string = RETRY_MODES.STANDARD;
 
   constructor(private readonly maxAttemptsProvider: Provider<number>, options?: StandardRetryStrategyOptions) {
     this.retryDecider = options?.retryDecider ?? defaultRetryDecider;
@@ -54,7 +54,11 @@ export class StandardRetryStrategy implements RetryStrategy {
 
   async retry<Input extends object, Ouput extends MetadataBearer>(
     next: FinalizeHandler<Input, Ouput>,
-    args: FinalizeHandlerArguments<Input>
+    args: FinalizeHandlerArguments<Input>,
+    options?: {
+      beforeRequest: Function;
+      afterRequest: Function;
+    }
   ) {
     let retryTokenAmount;
     let attempts = 0;
@@ -72,7 +76,14 @@ export class StandardRetryStrategy implements RetryStrategy {
         if (HttpRequest.isInstance(request)) {
           request.headers[REQUEST_HEADER] = `attempt=${attempts + 1}; max=${maxAttempts}`;
         }
+
+        if (options?.beforeRequest) {
+          await options.beforeRequest();
+        }
         const { response, output } = await next(args);
+        if (options?.afterRequest) {
+          options.afterRequest(response);
+        }
 
         this.retryQuota.releaseRetryTokens(retryTokenAmount);
         output.$metadata.attempts = attempts + 1;

packages/middleware-retry/src/config.ts

@@ -0,0 +1,15 @@
+export enum RETRY_MODES {
+  STANDARD = "standard",
+  ADAPTIVE = "adaptive",
+}
+
+/**
+ * The default value for how many HTTP requests an SDK should make for a
+ * single SDK operation invocation before giving up
+ */
+export const DEFAULT_MAX_ATTEMPTS = 3;
+
+/**
+ * The default retry algorithm to use.
+ */
+export const DEFAULT_RETRY_MODE = RETRY_MODES.STANDARD;

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

@@ -1,13 +1,13 @@
+import { DEFAULT_MAX_ATTEMPTS } from "./config";
 import {
   CONFIG_MAX_ATTEMPTS,
-  DEFAULT_MAX_ATTEMPTS,
   ENV_MAX_ATTEMPTS,
   NODE_MAX_ATTEMPT_CONFIG_OPTIONS,
   resolveRetryConfig,
 } from "./configurations";
-import { StandardRetryStrategy } from "./defaultStrategy";
+import { StandardRetryStrategy } from "./StandardRetryStrategy";
 
-jest.mock("./defaultStrategy");
+jest.mock("./StandardRetryStrategy");
 
 describe("resolveRetryConfig", () => {
   afterEach(() => {

packages/middleware-retry/src/configurations.ts

@@ -1,22 +1,12 @@
 import { LoadedConfigSelectors } from "@aws-sdk/node-config-provider";
 import { Provider, RetryStrategy } from "@aws-sdk/types";
 
-import { StandardRetryStrategy } from "./defaultStrategy";
+import { DEFAULT_MAX_ATTEMPTS, DEFAULT_RETRY_MODE } from "./config";
+import { StandardRetryStrategy } from "./StandardRetryStrategy";
 
 export const ENV_MAX_ATTEMPTS = "AWS_MAX_ATTEMPTS";
 export const CONFIG_MAX_ATTEMPTS = "max_attempts";
 
-/**
- * The default value for how many HTTP requests an SDK should make for a
- * single SDK operation invocation before giving up
- */
-export const DEFAULT_MAX_ATTEMPTS = 3;
-
-/**
- * The default retry algorithm to use.
- */
-export const DEFAULT_RETRY_MODE = "standard";
-
 export const NODE_MAX_ATTEMPT_CONFIG_OPTIONS: LoadedConfigSelectors<number> = {
   environmentVariableSelector: (env) => {
     const value = env[ENV_MAX_ATTEMPTS];

packages/middleware-retry/src/index.ts

@@ -1,6 +1,8 @@
 export * from "./retryMiddleware";
 export * from "./omitRetryHeadersMiddleware";
-export * from "./defaultStrategy";
+export * from "./StandardRetryStrategy";
+export * from "./AdaptiveRetryStrategy";
+export * from "./config";
 export * from "./configurations";
 export * from "./delayDecider";
 export * from "./retryDecider";