[core-rest-pipeline] Support "retry-after-ms" and "x-ms-retry-after-m…

…s" headers with the throttling policy (Azure#20817) * expand throttlingRetryStrategy with "retry-after-ms", "x-ms-retry-after-ms" headers * Expand test suite with "retry-after-ms" and "x-ms-retry-after-ms" * changelog * Update sdk/core/core-rest-pipeline/src/retryStrategies/throttlingRetryStrategy.ts * no new variable * no internal tag * refactor * new lines for readability * more refactor * Update sdk/core/core-rest-pipeline/src/retryStrategies/throttlingRetryStrategy.ts * part of the feedback * lint * Update sdk/core/core-rest-pipeline/src/retryStrategies/throttlingRetryStrategy.ts Co-authored-by: Jeff Fisher <[email protected]> * format Co-authored-by: Jeff Fisher <[email protected]>
AzureSDKAutomation · Mar 14, 2022 · a38203f · a38203f
1 parent e474d5b
commit a38203f
Show file tree

Hide file tree

Showing 5 changed files with 146 additions and 53 deletions.
diff --git a/sdk/core/core-rest-pipeline/CHANGELOG.md b/sdk/core/core-rest-pipeline/CHANGELOG.md
@@ -4,6 +4,8 @@
 
 ### Features Added
 
+- Supports the `"retry-after-ms"` and `"x-ms-retry-after-ms"` headers along with the `"Retry-After"` header from throttling retry responses from the services. [#20817](https://github.com/Azure/azure-sdk-for-js/issues/20817)
+
 ### Breaking Changes
 
 ### Bugs Fixed

diff --git a/sdk/core/core-rest-pipeline/src/policies/retryPolicy.ts b/sdk/core/core-rest-pipeline/src/policies/retryPolicy.ts
@@ -109,7 +109,7 @@ export function retryPolicy(
             throw errorToThrow;
           }
 
-          if (retryAfterInMs) {
+          if (retryAfterInMs || retryAfterInMs === 0) {
             strategyLogger.info(
               `Retry ${retryCount}: Retry strategy ${strategy.name} retries after ${retryAfterInMs}`
             );

diff --git a/sdk/core/core-rest-pipeline/src/retryStrategies/throttlingRetryStrategy.ts b/sdk/core/core-rest-pipeline/src/retryStrategies/throttlingRetryStrategy.ts
@@ -2,54 +2,77 @@
 // Licensed under the MIT license.
 
 import { PipelineResponse } from "..";
+import { parseHeaderValueAsNumber } from "../util/helpers";
 import { RetryStrategy } from "./retryStrategy";
 
 /**
- * Returns the number of milliseconds to wait based on a Retry-After header value.
- * Returns undefined if there is no valid value.
- * @param headerValue - An HTTP Retry-After header value.
+ * The header that comes back from Azure services representing
+ * the amount of time (minimum) to wait to retry (in seconds or timestamp after which we can retry).
  */
-function parseRetryAfterHeader(headerValue: string): number | undefined {
+const RetryAfterHeader = "Retry-After";
+/**
+ * The headers that come back from Azure services representing
+ * the amount of time (minimum) to wait to retry.
+ *
+ * "retry-after-ms", "x-ms-retry-after-ms" : milliseconds
+ * "Retry-After" : seconds or timestamp
+ */
+const AllRetryAfterHeaders: string[] = ["retry-after-ms", "x-ms-retry-after-ms", RetryAfterHeader];
+
+/**
+ * A response is a throttling retry response if it has a throttling status code (429 or 503),
+ * as long as one of the [ "Retry-After" or "retry-after-ms" or "x-ms-retry-after-ms" ] headers has a valid value.
+ *
+ * Returns the `retryAfterInMs` value if the response is a throttling retry response.
+ * If not throttling retry response, returns `undefined`.
+ *
+ * @internal
+ */
+function getRetryAfterInMs(response?: PipelineResponse): number | undefined {
+  if (!(response && [429, 503].includes(response.status))) return undefined;
   try {
-    const retryAfterInSeconds = Number(headerValue);
-    if (!Number.isNaN(retryAfterInSeconds)) {
-      return retryAfterInSeconds * 1000;
-    } else {
-      // It might be formatted as a date instead of a number of seconds
+    // Headers: "retry-after-ms", "x-ms-retry-after-ms", "Retry-After"
+    for (const header of AllRetryAfterHeaders) {
+      const retryAfterValue = parseHeaderValueAsNumber(response, header);
+      if (retryAfterValue === 0 || retryAfterValue) {
+        // "Retry-After" header ==> seconds
+        // "retry-after-ms", "x-ms-retry-after-ms" headers ==> milli-seconds
+        const multiplyingFactor = header === RetryAfterHeader ? 1000 : 1;
+        return retryAfterValue * multiplyingFactor; // in milli-seconds
+      }
+    }
 
-      const now: number = Date.now();
-      const date: number = Date.parse(headerValue);
-      const diff = date - now;
+    // RetryAfterHeader ("Retry-After") has a special case where it might be formatted as a date instead of a number of seconds
+    const retryAfterHeader = response.headers.get(RetryAfterHeader);
+    if (!retryAfterHeader) return;
 
-      return Number.isNaN(diff) ? undefined : diff;
-    }
+    const date = Date.parse(retryAfterHeader);
+    const diff = date - Date.now();
+    // negative diff would mean a date in the past, so retry asap with 0 milliseconds
+    return Number.isFinite(diff) ? Math.max(0, diff) : undefined;
   } catch (e) {
     return undefined;
   }
 }
 
 /**
  * A response is a retry response if it has a throttling status code (429 or 503),
- * as long as the Retry-After header has a valid value.
+ * as long as one of the [ "Retry-After" or "retry-after-ms" or "x-ms-retry-after-ms" ] headers has a valid value.
  */
 export function isThrottlingRetryResponse(response?: PipelineResponse): boolean {
-  return Boolean(
-    response &&
-      (response.status === 429 || response.status === 503) &&
-      response.headers.get("Retry-After") &&
-      parseRetryAfterHeader(response.headers.get("Retry-After")!)
-  );
+  return Number.isFinite(getRetryAfterInMs(response));
 }
 
 export function throttlingRetryStrategy(): RetryStrategy {
   return {
     name: "throttlingRetryStrategy",
     retry({ response }) {
-      if (!isThrottlingRetryResponse(response)) {
+      const retryAfterInMs = getRetryAfterInMs(response);
+      if (!Number.isFinite(retryAfterInMs)) {
         return { skipStrategy: true };
       }
       return {
-        retryAfterInMs: parseRetryAfterHeader(response!.headers.get("Retry-After")!),
+        retryAfterInMs,
       };
     },
   };

diff --git a/sdk/core/core-rest-pipeline/src/util/helpers.ts b/sdk/core/core-rest-pipeline/src/util/helpers.ts
@@ -2,6 +2,7 @@
 // Licensed under the MIT license.
 
 import { AbortError, AbortSignalLike } from "@azure/abort-controller";
+import { PipelineResponse } from "../interfaces";
 
 /**
  * A constant that indicates whether the environment the code is running is Node.JS.
@@ -106,3 +107,18 @@ export function isObject(input: unknown): input is UnknownObject {
     !(input instanceof Date)
   );
 }
+
+/**
+ * @internal
+ * @returns the parsed value or undefined if the parsed value is invalid.
+ */
+export function parseHeaderValueAsNumber(
+  response: PipelineResponse,
+  headerName: string
+): number | undefined {
+  const value = response.headers.get(headerName);
+  if (!value) return;
+  const valueAsNum = Number(value);
+  if (Number.isNaN(valueAsNum)) return;
+  return valueAsNum;
+}
diff --git a/sdk/core/core-rest-pipeline/test/throttlingRetryPolicy.spec.ts b/sdk/core/core-rest-pipeline/test/throttlingRetryPolicy.spec.ts
@@ -21,13 +21,68 @@ describe("throttlingRetryPolicy", function () {
     sinon.restore();
   });
 
-  it("It should retry after a given number of seconds on a response with status code 429", async () => {
+  const defaultDurations = [0, 10 * 1000]; // milliseconds
+
+  defaultDurations.forEach((defaultDuration) => {
+    const headersWithDefaultDuration = [
+      {
+        "Retry-After": String(defaultDuration / 1000),
+      },
+      {
+        "retry-after-ms": String(defaultDuration),
+      },
+      {
+        "x-ms-retry-after-ms": String(defaultDuration),
+      },
+    ] as const;
+    headersWithDefaultDuration.forEach((headers) => {
+      it(`(${
+        Object.keys(headers)[0]
+      }) - should retry after a given number of seconds/milliseconds on a response with status code 429`, async () => {
+        const request = createPipelineRequest({
+          url: "https://bing.com",
+        });
+        const retryResponse: PipelineResponse = {
+          headers: createHttpHeaders(headers),
+          request,
+          status: 429,
+        };
+        const successResponse: PipelineResponse = {
+          headers: createHttpHeaders(),
+          request,
+          status: 200,
+        };
+
+        const policy = throttlingRetryPolicy();
+        const next = sinon.stub<Parameters<SendRequest>, ReturnType<SendRequest>>();
+        next.onFirstCall().resolves(retryResponse);
+        next.onSecondCall().resolves(successResponse);
+
+        const clock = sinon.useFakeTimers();
+
+        const promise = policy.sendRequest(request, next);
+        assert.isTrue(next.calledOnce, "next wasn't called once");
+
+        // allow the delay to occur
+        const time = await clock.nextAsync();
+        assert.strictEqual(time, defaultDuration);
+        assert.isTrue(next.calledTwice, "next wasn't called twice");
+
+        const result = await promise;
+
+        assert.strictEqual(result, successResponse);
+        clock.restore();
+      });
+    });
+  });
+
+  it("It should retry after a given date occurs on a response with status code 429", async () => {
     const request = createPipelineRequest({
       url: "https://bing.com",
     });
     const retryResponse: PipelineResponse = {
       headers: createHttpHeaders({
-        "Retry-After": "10",
+        "Retry-After": "Wed, 21 Oct 2015 07:28:00 GMT",
       }),
       request,
       status: 429,
@@ -43,14 +98,18 @@ describe("throttlingRetryPolicy", function () {
     next.onFirstCall().resolves(retryResponse);
     next.onSecondCall().resolves(successResponse);
 
-    const clock = sinon.useFakeTimers();
+    const clock = sinon.useFakeTimers(new Date("Wed, 21 Oct 2015 07:20:00 GMT"));
 
     const promise = policy.sendRequest(request, next);
     assert.isTrue(next.calledOnce);
 
     // allow the delay to occur
     const time = await clock.nextAsync();
-    assert.strictEqual(time, 10 * 1000);
+    assert.strictEqual(
+      time,
+      new Date("Wed, 21 Oct 2015 07:28:00 GMT").getTime(),
+      "It should now be the time from the header."
+    );
     assert.isTrue(next.calledTwice);
 
     const result = await promise;
@@ -59,16 +118,16 @@ describe("throttlingRetryPolicy", function () {
     clock.restore();
   });
 
-  it("It should retry after a given date occurs on a response with status code 429", async () => {
+  it("It should retry after a given number of seconds on a response with status code 503", async () => {
     const request = createPipelineRequest({
       url: "https://bing.com",
     });
     const retryResponse: PipelineResponse = {
       headers: createHttpHeaders({
-        "Retry-After": "Wed, 21 Oct 2015 07:28:00 GMT",
+        "Retry-After": "10",
       }),
       request,
-      status: 429,
+      status: 503,
     };
     const successResponse: PipelineResponse = {
       headers: createHttpHeaders(),
@@ -81,18 +140,14 @@ describe("throttlingRetryPolicy", function () {
     next.onFirstCall().resolves(retryResponse);
     next.onSecondCall().resolves(successResponse);
 
-    const clock = sinon.useFakeTimers(new Date("Wed, 21 Oct 2015 07:20:00 GMT"));
+    const clock = sinon.useFakeTimers();
 
     const promise = policy.sendRequest(request, next);
     assert.isTrue(next.calledOnce);
 
     // allow the delay to occur
     const time = await clock.nextAsync();
-    assert.strictEqual(
-      time,
-      new Date("Wed, 21 Oct 2015 07:28:00 GMT").getTime(),
-      "It should now be the time from the header."
-    );
+    assert.strictEqual(time, 10 * 1000);
     assert.isTrue(next.calledTwice);
 
     const result = await promise;
@@ -101,13 +156,13 @@ describe("throttlingRetryPolicy", function () {
     clock.restore();
   });
 
-  it("It should retry after a given number of seconds on a response with status code 503", async () => {
+  it("It should retry after a given date occurs on a response with status code 503", async () => {
     const request = createPipelineRequest({
       url: "https://bing.com",
     });
     const retryResponse: PipelineResponse = {
       headers: createHttpHeaders({
-        "Retry-After": "10",
+        "Retry-After": "Wed, 21 Oct 2015 07:28:00 GMT",
       }),
       request,
       status: 503,
@@ -123,14 +178,18 @@ describe("throttlingRetryPolicy", function () {
     next.onFirstCall().resolves(retryResponse);
     next.onSecondCall().resolves(successResponse);
 
-    const clock = sinon.useFakeTimers();
+    const clock = sinon.useFakeTimers(new Date("Wed, 21 Oct 2015 07:20:00 GMT"));
 
     const promise = policy.sendRequest(request, next);
     assert.isTrue(next.calledOnce);
 
     // allow the delay to occur
     const time = await clock.nextAsync();
-    assert.strictEqual(time, 10 * 1000);
+    assert.strictEqual(
+      time,
+      new Date("Wed, 21 Oct 2015 07:28:00 GMT").getTime(),
+      "It should now be the time from the header."
+    );
     assert.isTrue(next.calledTwice);
 
     const result = await promise;
@@ -139,7 +198,7 @@ describe("throttlingRetryPolicy", function () {
     clock.restore();
   });
 
-  it("It should retry after a given date occurs on a response with status code 503", async () => {
+  it("It should retry after 0 seconds with status code 503 for a past date", async () => {
     const request = createPipelineRequest({
       url: "https://bing.com",
     });
@@ -161,19 +220,12 @@ describe("throttlingRetryPolicy", function () {
     next.onFirstCall().resolves(retryResponse);
     next.onSecondCall().resolves(successResponse);
 
-    const clock = sinon.useFakeTimers(new Date("Wed, 21 Oct 2015 07:20:00 GMT"));
-
     const promise = policy.sendRequest(request, next);
-    assert.isTrue(next.calledOnce);
+    const clock = sinon.useFakeTimers();
 
-    // allow the delay to occur
-    const time = await clock.nextAsync();
-    assert.strictEqual(
-      time,
-      new Date("Wed, 21 Oct 2015 07:28:00 GMT").getTime(),
-      "It should now be the time from the header."
-    );
-    assert.isTrue(next.calledTwice);
+    assert.isTrue(next.calledOnce, "next wasn't called once");
+    await clock.nextAsync();
+    assert.isTrue(next.calledTwice, "next wasn't called twice");
 
     const result = await promise;