fix(middleware-retry): use delay from response header if it's higher (#3916)

Author: trivikr

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

@@ -1,4 +1,4 @@
-import { HttpRequest } from "@aws-sdk/protocol-http";
+import { HttpRequest, HttpResponse } from "@aws-sdk/protocol-http";
 import { isThrottlingError } from "@aws-sdk/service-error-classification";
 import { v4 } from "uuid";
 
@@ -88,6 +88,9 @@ describe("defaultStrategy", () => {
     (HttpRequest as unknown as jest.Mock).mockReturnValue({
       isInstance: jest.fn().mockReturnValue(false),
     });
+    (HttpResponse as unknown as jest.Mock).mockReturnValue({
+      isInstance: jest.fn().mockReturnValue(false),
+    });
     (v4 as jest.Mock).mockReturnValue("42");
   });
 
@@ -220,22 +223,101 @@ describe("defaultStrategy", () => {
       });
     });
 
-    it("delay value returned", async () => {
-      jest.spyOn(global, "setTimeout");
+    describe("totalRetryDelay", () => {
+      describe("when retry-after is not set", () => {
+        it("should be equal to sum of values computed 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 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);
+        });
+      });
 
-      const FIRST_DELAY = 100;
-      const SECOND_DELAY = 200;
+      describe("when retry-after is set", () => {
+        const getErrorWithValues = async (
+          delayDeciderInMs: number,
+          retryAfter: number | string,
+          retryAfterHeaderName?: string
+        ) => {
+          (defaultDelayDecider as jest.Mock).mockReturnValueOnce(delayDeciderInMs);
+
+          const maxAttempts = 2;
+          const mockError = new Error();
+          Object.defineProperty(mockError, "$response", {
+            value: {
+              headers: { [retryAfterHeaderName ? retryAfterHeaderName : "retry-after"]: String(retryAfter) },
+            },
+          });
+          const error = await mockFailedOperation(maxAttempts, { mockError });
+          expect(defaultDelayDecider as jest.Mock).toHaveBeenCalledTimes(maxAttempts - 1);
+          expect(setTimeout).toHaveBeenCalledTimes(maxAttempts - 1);
+
+          return error;
+        };
+
+        beforeEach(() => {
+          jest.spyOn(global, "setTimeout");
+        });
 
-      (defaultDelayDecider as jest.Mock).mockReturnValueOnce(FIRST_DELAY).mockReturnValueOnce(SECOND_DELAY);
+        describe("uses retry-after value if it's greater than that from delayDecider", () => {
+          beforeEach(() => {
+            const { isInstance } = HttpResponse;
+            (isInstance as unknown as jest.Mock).mockReturnValueOnce(true);
+          });
+
+          describe("when value is in seconds", () => {
+            const testWithHeaderName = async (retryAfterHeaderName: string) => {
+              const delayDeciderInMs = 2000;
+              const retryAfterInSeconds = 3;
+
+              const error = await getErrorWithValues(delayDeciderInMs, retryAfterInSeconds, retryAfterHeaderName);
+              expect(error.$metadata.totalRetryDelay).toEqual(retryAfterInSeconds * 1000);
+              expect((setTimeout as unknown as jest.Mock).mock.calls[0][1]).toBe(retryAfterInSeconds * 1000);
+            };
+
+            it("with header in small case", async () => {
+              testWithHeaderName("retry-after");
+            });
+
+            it("with header with first letter capital", async () => {
+              testWithHeaderName("Retry-After");
+            });
+          });
+
+          it("when value is a Date", async () => {
+            const mockDateNow = Date.now();
+            jest.spyOn(Date, "now").mockReturnValue(mockDateNow);
+
+            const delayDeciderInMs = 2000;
+            const retryAfterInSeconds = 3;
+            const retryAfterDate = new Date(mockDateNow + retryAfterInSeconds * 1000);
+
+            const error = await getErrorWithValues(delayDeciderInMs, retryAfterDate.toISOString());
+            expect(error.$metadata.totalRetryDelay).toEqual(retryAfterInSeconds * 1000);
+            expect((setTimeout as unknown as jest.Mock).mock.calls[0][1]).toBe(retryAfterInSeconds * 1000);
+          });
+        });
 
-      const maxAttempts = 3;
-      const error = await mockFailedOperation(maxAttempts);
-      expect(error.$metadata.totalRetryDelay).toEqual(FIRST_DELAY + SECOND_DELAY);
+        it("ignores retry-after value if it's smaller than that from delayDecider", async () => {
+          const delayDeciderInMs = 3000;
+          const retryAfterInSeconds = 2;
 
-      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);
+          const error = await getErrorWithValues(delayDeciderInMs, retryAfterInSeconds);
+          expect(error.$metadata.totalRetryDelay).toEqual(delayDeciderInMs);
+          expect((setTimeout as unknown as jest.Mock).mock.calls[0][1]).toBe(delayDeciderInMs);
+        });
+      });
     });
   });
 

packages/middleware-retry/src/StandardRetryStrategy.ts

@@ -1,4 +1,4 @@
-import { HttpRequest } from "@aws-sdk/protocol-http";
+import { HttpRequest, HttpResponse } from "@aws-sdk/protocol-http";
 import { isThrottlingError } from "@aws-sdk/service-error-classification";
 import { SdkError } from "@aws-sdk/types";
 import { FinalizeHandler, FinalizeHandlerArguments, MetadataBearer, Provider, RetryStrategy } from "@aws-sdk/types";
@@ -95,10 +95,14 @@ export class StandardRetryStrategy implements RetryStrategy {
         attempts++;
         if (this.shouldRetry(err as SdkError, attempts, maxAttempts)) {
           retryTokenAmount = this.retryQuota.retrieveRetryTokens(err);
-          const delay = this.delayDecider(
+          const delayFromDecider = this.delayDecider(
             isThrottlingError(err) ? THROTTLING_RETRY_DELAY_BASE : DEFAULT_RETRY_DELAY_BASE,
             attempts
           );
+
+          const delayFromResponse = getDelayFromRetryAfterHeader(err.$response);
+          const delay = Math.max(delayFromResponse || 0, delayFromDecider);
+
           totalDelay += delay;
 
           await new Promise((resolve) => setTimeout(resolve, delay));
@@ -117,6 +121,24 @@ export class StandardRetryStrategy implements RetryStrategy {
   }
 }
 
+/**
+ * Returns number of milliseconds to wait based on "Retry-After" header value.
+ */
+const getDelayFromRetryAfterHeader = (response: unknown): number | undefined => {
+  if (!HttpResponse.isInstance(response)) return;
+
+  const retryAfterHeaderName = Object.keys(response.headers).find((key) => key.toLowerCase() === "retry-after");
+  if (!retryAfterHeaderName) return;
+
+  const retryAfter = response.headers[retryAfterHeaderName];
+
+  const retryAfterSeconds = Number(retryAfter);
+  if (!Number.isNaN(retryAfterSeconds)) return retryAfterSeconds * 1000;
+
+  const retryAfterDate = new Date(retryAfter);
+  return retryAfterDate.getTime() - Date.now();
+};
+
 const asSdkError = (error: unknown): SdkError => {
   if (error instanceof Error) return error;
   if (error instanceof Object) return Object.assign(new Error(), error);