feat(middleware-retry): call retry strategy based on value in retryMode (#2456)

trivikr · web-flow · commit 095b9906889b · 2021-06-03T11:36:00.000-07:00
diff --git a/packages/middleware-retry/src/configurations.spec.ts b/packages/middleware-retry/src/configurations.spec.ts
@@ -1,3 +1,4 @@
+import { AdaptiveRetryStrategy } from "./AdaptiveRetryStrategy";
 import { DEFAULT_MAX_ATTEMPTS } from "./config";
 import {
   CONFIG_MAX_ATTEMPTS,
@@ -7,17 +8,19 @@ import {
 } from "./configurations";
 import { StandardRetryStrategy } from "./StandardRetryStrategy";
 
+jest.mock("./AdaptiveRetryStrategy");
 jest.mock("./StandardRetryStrategy");
 
-describe("resolveRetryConfig", () => {
+describe(resolveRetryConfig.name, () => {
+  const retryModeProvider = jest.fn();
   afterEach(() => {
     jest.clearAllMocks();
   });
 
   describe("maxAttempts", () => {
     it("assigns maxAttempts value if present", async () => {
       for (const maxAttempts of [1, 2, 3]) {
-        const output = await resolveRetryConfig({ maxAttempts }).maxAttempts();
+        const output = await resolveRetryConfig({ maxAttempts, retryModeProvider }).maxAttempts();
         expect(output).toStrictEqual(maxAttempts);
       }
     });
@@ -29,21 +32,85 @@ describe("resolveRetryConfig", () => {
         retry: jest.fn(),
       };
       const { retryStrategy } = resolveRetryConfig({
+        retryModeProvider,
         retryStrategy: mockRetryStrategy,
       });
-      expect(retryStrategy).toEqual(mockRetryStrategy);
+      expect(retryStrategy()).resolves.toEqual(mockRetryStrategy);
     });
 
-    describe("creates StandardRetryStrategy if retryStrategy not present", () => {
-      describe("passes maxAttempts if present", () => {
-        for (const maxAttempts of [1, 2, 3]) {
-          it(`when maxAttempts=${maxAttempts}`, async () => {
-            resolveRetryConfig({ maxAttempts });
-            expect(StandardRetryStrategy as jest.Mock).toHaveBeenCalledTimes(1);
-            const output = await (StandardRetryStrategy as jest.Mock).mock.calls[0][0]();
-            expect(output).toStrictEqual(maxAttempts);
+    describe("creates RetryStrategy if retryStrategy not present", () => {
+      describe("StandardRetryStrategy", () => {
+        describe("when retryMode=standard", () => {
+          describe("passes maxAttempts if present", () => {
+            const retryMode = "standard";
+            for (const maxAttempts of [1, 2, 3]) {
+              it(`when maxAttempts=${maxAttempts}`, async () => {
+                const { retryStrategy } = resolveRetryConfig({ maxAttempts, retryMode, retryModeProvider });
+                await retryStrategy();
+                expect(StandardRetryStrategy as jest.Mock).toHaveBeenCalledTimes(1);
+                expect(AdaptiveRetryStrategy as jest.Mock).not.toHaveBeenCalled();
+                const output = await (StandardRetryStrategy as jest.Mock).mock.calls[0][0]();
+                expect(output).toStrictEqual(maxAttempts);
+              });
+            }
           });
-        }
+        });
+
+        describe("when retryModeProvider returns 'standard'", () => {
+          describe("passes maxAttempts if present", () => {
+            beforeEach(() => {
+              retryModeProvider.mockResolvedValueOnce("standard");
+            });
+            for (const maxAttempts of [1, 2, 3]) {
+              it(`when maxAttempts=${maxAttempts}`, async () => {
+                const { retryStrategy } = resolveRetryConfig({ maxAttempts, retryModeProvider });
+                await retryStrategy();
+                expect(retryModeProvider).toHaveBeenCalledTimes(1);
+                expect(StandardRetryStrategy as jest.Mock).toHaveBeenCalledTimes(1);
+                expect(AdaptiveRetryStrategy as jest.Mock).not.toHaveBeenCalled();
+                const output = await (StandardRetryStrategy as jest.Mock).mock.calls[0][0]();
+                expect(output).toStrictEqual(maxAttempts);
+              });
+            }
+          });
+        });
+      });
+
+      describe("AdaptiveRetryStrategy", () => {
+        describe("when retryMode=adaptive", () => {
+          describe("passes maxAttempts if present", () => {
+            const retryMode = "adaptive";
+            for (const maxAttempts of [1, 2, 3]) {
+              it(`when maxAttempts=${maxAttempts}`, async () => {
+                const { retryStrategy } = resolveRetryConfig({ maxAttempts, retryMode, retryModeProvider });
+                await retryStrategy();
+                expect(StandardRetryStrategy as jest.Mock).not.toHaveBeenCalled();
+                expect(AdaptiveRetryStrategy as jest.Mock).toHaveBeenCalledTimes(1);
+                const output = await (AdaptiveRetryStrategy as jest.Mock).mock.calls[0][0]();
+                expect(output).toStrictEqual(maxAttempts);
+              });
+            }
+          });
+        });
+
+        describe("when retryModeProvider returns 'adaptive'", () => {
+          describe("passes maxAttempts if present", () => {
+            beforeEach(() => {
+              retryModeProvider.mockResolvedValueOnce("adaptive");
+            });
+            for (const maxAttempts of [1, 2, 3]) {
+              it(`when maxAttempts=${maxAttempts}`, async () => {
+                const { retryStrategy } = resolveRetryConfig({ maxAttempts, retryModeProvider });
+                await retryStrategy();
+                expect(retryModeProvider).toHaveBeenCalledTimes(1);
+                expect(StandardRetryStrategy as jest.Mock).not.toHaveBeenCalled();
+                expect(AdaptiveRetryStrategy as jest.Mock).toHaveBeenCalledTimes(1);
+                const output = await (AdaptiveRetryStrategy as jest.Mock).mock.calls[0][0]();
+                expect(output).toStrictEqual(maxAttempts);
+              });
+            }
+          });
+        });
       });
     });
   });
diff --git a/packages/middleware-retry/src/configurations.ts b/packages/middleware-retry/src/configurations.ts
@@ -1,7 +1,8 @@
 import { LoadedConfigSelectors } from "@aws-sdk/node-config-provider";
 import { Provider, RetryStrategy } from "@aws-sdk/types";
 
-import { DEFAULT_MAX_ATTEMPTS, DEFAULT_RETRY_MODE } from "./config";
+import { AdaptiveRetryStrategy } from "./AdaptiveRetryStrategy";
+import { DEFAULT_MAX_ATTEMPTS, DEFAULT_RETRY_MODE, RETRY_MODES } from "./config";
 import { StandardRetryStrategy } from "./StandardRetryStrategy";
 
 export const ENV_MAX_ATTEMPTS = "AWS_MAX_ATTEMPTS";
@@ -38,9 +39,20 @@ export interface RetryInputConfig {
    * The strategy to retry the request. Using built-in exponential backoff strategy by default.
    */
   retryStrategy?: RetryStrategy;
+  /**
+   * Specifies which retry algorithm to use.
+   */
+  retryMode?: string;
+}
+
+interface PreviouslyResolved {
+  /**
+   * Specifies provider for retry algorithm to use.
+   * @internal
+   */
+  retryModeProvider: Provider<string>;
 }
 
-interface PreviouslyResolved {}
 export interface RetryResolvedConfig {
   /**
    * Resolved value for input config {@link RetryInputConfig.maxAttempts}
@@ -49,15 +61,24 @@ export interface RetryResolvedConfig {
   /**
    * Resolved value for input config {@link RetryInputConfig.retryStrategy}
    */
-  retryStrategy: RetryStrategy;
+  retryStrategy: Provider<RetryStrategy>;
 }
 
 export const resolveRetryConfig = <T>(input: T & PreviouslyResolved & RetryInputConfig): T & RetryResolvedConfig => {
   const maxAttempts = normalizeMaxAttempts(input.maxAttempts);
   return {
     ...input,
     maxAttempts,
-    retryStrategy: input.retryStrategy || new StandardRetryStrategy(maxAttempts),
+    retryStrategy: async () => {
+      if (input.retryStrategy) {
+        return input.retryStrategy;
+      }
+      const retryMode = input.retryMode || (await input.retryModeProvider());
+      if (retryMode === RETRY_MODES.ADAPTIVE) {
+        return new AdaptiveRetryStrategy(maxAttempts);
+      }
+      return new StandardRetryStrategy(maxAttempts);
+    },
   };
 };
 
diff --git a/packages/middleware-retry/src/retryMiddleware.spec.ts b/packages/middleware-retry/src/retryMiddleware.spec.ts
@@ -2,10 +2,14 @@ import { FinalizeHandlerArguments, HandlerExecutionContext, MiddlewareStack, Ret
 
 import { getRetryPlugin, retryMiddleware, retryMiddlewareOptions } from "./retryMiddleware";
 
-describe("getRetryPlugin", () => {
+describe(getRetryPlugin.name, () => {
   const mockClientStack = {
     add: jest.fn(),
   };
+  const mockRetryStrategy = {
+    mode: "mock",
+    retry: jest.fn(),
+  };
 
   afterEach(() => {
     jest.clearAllMocks();
@@ -16,7 +20,7 @@ describe("getRetryPlugin", () => {
       it(`when maxAttempts=${maxAttempts}`, () => {
         getRetryPlugin({
           maxAttempts: () => Promise.resolve(maxAttempts),
-          retryStrategy: {} as RetryStrategy,
+          retryStrategy: jest.fn().mockResolvedValue(mockRetryStrategy),
         }).applyToStack(mockClientStack as unknown as MiddlewareStack<any, any>);
         expect(mockClientStack.add).toHaveBeenCalledTimes(1);
         expect(mockClientStack.add.mock.calls[0][1]).toEqual(retryMiddlewareOptions);
@@ -25,7 +29,11 @@ describe("getRetryPlugin", () => {
   });
 });
 
-describe("retryMiddleware", () => {
+describe(retryMiddleware.name, () => {
+  const mockRetryStrategy = {
+    mode: "mock",
+    retry: jest.fn(),
+  };
   afterEach(() => {
     jest.clearAllMocks();
   });
@@ -36,22 +44,17 @@ describe("retryMiddleware", () => {
     const args = {
       request: {},
     };
-    const mockRetryStrategy = {
-      mode: "mock",
-      maxAttempts,
-      retry: jest.fn(),
-    };
     const context: HandlerExecutionContext = {};
 
     await retryMiddleware({
       maxAttempts: () => Promise.resolve(maxAttempts),
-      retryStrategy: mockRetryStrategy,
+      retryStrategy: jest.fn().mockResolvedValue({ ...mockRetryStrategy, maxAttempts }),
     })(
       next,
       context
     )(args as FinalizeHandlerArguments<any>);
     expect(mockRetryStrategy.retry).toHaveBeenCalledTimes(1);
     expect(mockRetryStrategy.retry).toHaveBeenCalledWith(next, args);
-    expect(context.userAgent).toContainEqual(["cfg/retry-mode", "mock"]);
+    expect(context.userAgent).toContainEqual(["cfg/retry-mode", mockRetryStrategy.mode]);
   });
 });
diff --git a/packages/middleware-retry/src/retryMiddleware.ts b/packages/middleware-retry/src/retryMiddleware.ts
@@ -18,9 +18,9 @@ export const retryMiddleware =
     context: HandlerExecutionContext
   ): FinalizeHandler<any, Output> =>
   async (args: FinalizeHandlerArguments<any>): Promise<FinalizeHandlerOutput<Output>> => {
-    if (options?.retryStrategy?.mode)
-      context.userAgent = [...(context.userAgent || []), ["cfg/retry-mode", options.retryStrategy.mode]];
-    return options.retryStrategy.retry(next, args);
+    const retryStrategy = await options.retryStrategy();
+    if (retryStrategy?.mode) context.userAgent = [...(context.userAgent || []), ["cfg/retry-mode", retryStrategy.mode]];
+    return retryStrategy.retry(next, args);
   };
 
 export const retryMiddlewareOptions: FinalizeRequestHandlerOptions & AbsoluteLocation = {
