feat(ec2-metadata-service): add retries for IMDS requests

packages/ec2-metadata-service/src/MetadataService.spec.ts

@@ -1,6 +1,6 @@
 import { NodeHttpHandler } from "@smithy/node-http-handler";
 import { Readable } from "stream";
-import { beforeEach, describe, expect, test as it, vi } from "vitest";
+import { afterEach, beforeEach, describe, expect, test as it, vi } from "vitest";
 
 import { MetadataService } from "./MetadataService";
 
@@ -24,6 +24,7 @@ describe("MetadataService Socket Leak Checks", () => {
     metadataService = new MetadataService({
       endpoint: "http://169.254.169.254",
       httpOptions: { timeout: 1000 },
+      retries: 0, // Disable retries for faster tests
     });
   });
 
@@ -276,6 +277,7 @@ describe("MetadataService Custom Ports", () => {
     metadataService = new MetadataService({
       endpoint: "http://localhost:1338",
       httpOptions: { timeout: 1000 },
+      retries: 0, // Disable retries for faster tests
     });
 
     const mockResponse = createMockResponse(200, "i-1234567890abcdef0");
@@ -294,6 +296,7 @@ describe("MetadataService Custom Ports", () => {
     metadataService = new MetadataService({
       endpoint: "http://localhost:1338",
       httpOptions: { timeout: 1000 },
+      retries: 0, // Disable retries for faster tests
     });
 
     const mockResponse = createMockResponse(200, "test-token-123");
@@ -312,6 +315,7 @@ describe("MetadataService Custom Ports", () => {
     metadataService = new MetadataService({
       endpoint: "http://169.254.169.254",
       httpOptions: { timeout: 1000 },
+      retries: 0, // Disable retries for faster tests
     });
 
     const mockResponse = createMockResponse(200, "test-token-123");
@@ -326,3 +330,68 @@ describe("MetadataService Custom Ports", () => {
     expect(requestArg.hostname).toBe("169.254.169.254");
   });
 });
+
+describe("MetadataService Retry Configuration", () => {
+  it("should use default 3 retries", () => {
+    const metadataService = new MetadataService();
+    expect((metadataService as any).retries).toBe(3);
+  });
+
+  it("should use custom retry count", () => {
+    const metadataService = new MetadataService({ retries: 5 });
+    expect((metadataService as any).retries).toBe(5);
+  });
+
+  it("should disable retries when set to 0", () => {
+    const metadataService = new MetadataService({ retries: 0 });
+    expect((metadataService as any).retries).toBe(0);
+  });
+
+  it("should create backoff function", () => {
+    const metadataService = new MetadataService();
+    const backoffFn = (metadataService as any).backoffFn;
+    expect(typeof backoffFn).toBe("function");
+  });
+
+  describe("status code handling for retries", () => {
+    it("should not retry 400 errors", async () => {
+      const metadataService = new MetadataService({ retries: 0 });
+      const shouldNotRetry = (metadataService as any).shouldNotRetry;
+
+      const error = { statusCode: 400 };
+      expect(shouldNotRetry(error)).toBe(true);
+    });
+
+    it("should not retry 403 errors", async () => {
+      const metadataService = new MetadataService({ retries: 0 });
+      const shouldNotRetry = (metadataService as any).shouldNotRetry;
+
+      const error = { statusCode: 403 };
+      expect(shouldNotRetry(error)).toBe(true);
+    });
+
+    it("should not retry 404 errors", async () => {
+      const metadataService = new MetadataService({ retries: 0 });
+      const shouldNotRetry = (metadataService as any).shouldNotRetry;
+
+      const error = { statusCode: 404 };
+      expect(shouldNotRetry(error)).toBe(true);
+    });
+
+    it("should retry 401 errors", async () => {
+      const metadataService = new MetadataService({ retries: 0 });
+      const shouldNotRetry = (metadataService as any).shouldNotRetry;
+
+      const error = { statusCode: 401 };
+      expect(shouldNotRetry(error)).toBe(false);
+    });
+
+    it("should retry 500 errors", async () => {
+      const metadataService = new MetadataService({ retries: 0 });
+      const shouldNotRetry = (metadataService as any).shouldNotRetry;
+
+      const error = { statusCode: 500 };
+      expect(shouldNotRetry(error)).toBe(false);
+    });
+  });
+});

packages/ec2-metadata-service/src/MetadataService.ts

@@ -13,6 +13,8 @@ import { MetadataServiceOptions } from "./MetadataServiceOptions";
 export class MetadataService {
   private disableFetchToken: boolean;
   private config: Promise<MetadataServiceOptions>;
+  private retries: number;
+  private backoffFn: (numFailures: number) => Promise<void> | number;
 
   /**
    * Creates a new MetadataService object with a given set of options.
@@ -30,68 +32,131 @@ export class MetadataService {
       };
     })();
     this.disableFetchToken = options?.disableFetchToken || false;
+    this.retries = options?.retries ?? 3;
+    this.backoffFn = this.createBackoffFunction(options?.backoff);
   }
 
-  async request(path: string, options: { method?: string; headers?: Record<string, string> }): Promise<string> {
-    const { endpoint, ec2MetadataV1Disabled, httpOptions } = await this.config;
-    const handler = new NodeHttpHandler({
-      requestTimeout: httpOptions?.timeout,
-      throwOnRequestTimeout: true,
-      connectionTimeout: httpOptions?.timeout,
-    });
-    const endpointUrl = new URL(endpoint!);
-    const headers = options.headers || {};
-    /**
-     * If IMDSv1 is disabled and disableFetchToken is true, throw an error
-     * Refer: https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configuring-instance-metadata-service.html
-     */
-    if (this.disableFetchToken && ec2MetadataV1Disabled) {
-      throw new Error("IMDSv1 is disabled and fetching token is disabled, cannot make the request.");
+  private createBackoffFunction(
+    backoff?: number | ((numFailures: number) => Promise<void> | number)
+  ): (numFailures: number) => Promise<void> | number {
+    // backoff in seconds
+    if (typeof backoff === "function") {
+      return backoff;
     }
-    /**
-     * Make request with token if disableFetchToken is not true (IMDSv2).
-     * Note that making the request call with token will result in an additional request to fetch the token.
-     */
-    if (!this.disableFetchToken) {
+    if (typeof backoff === "number") {
+      return () => backoff;
+    }
+    // Default exponential backoff
+    return (numFailures: number) => Math.pow(1.2, numFailures);
+  }
+
+  private sleep(ms: number): Promise<void> {
+    return new Promise((resolve) => setTimeout(resolve, ms));
+  }
+
+  private async retryWithBackoff<T>(operation: () => Promise<T>): Promise<T> {
+    let lastError: Error;
+
+    for (let attempt = 0; attempt <= this.retries; attempt++) {
       try {
-        headers["x-aws-ec2-metadata-token"] = await this.fetchMetadataToken();
-      } catch (err) {
-        if (ec2MetadataV1Disabled) {
-          // If IMDSv1 is disabled and token fetch fails (IMDSv2 fails), rethrow the error
-          throw err;
+        return await operation();
+      } catch (error) {
+        lastError = error as Error;
+
+        // Don't retry on final attempt
+        if (attempt === this.retries) {
+          break;
+        }
+
+        if (this.shouldNotRetry(error as any)) {
+          throw error;
+        }
+
+        const backoffResult = this.backoffFn(attempt);
+        if (typeof backoffResult === "number") {
+          await this.sleep(backoffResult * 1000); // seconds to milliseconds
+        } else {
+          await backoffResult;
         }
-        // If token fetch fails and IMDSv1 is not disabled, proceed without token (IMDSv1 fallback)
       }
-    } // else, IMDSv1 fallback mode
-    const request = new HttpRequest({
-      method: options.method || "GET", // Default to GET if no method is specified
-      headers: headers,
-      hostname: endpointUrl.hostname,
-      path: endpointUrl.pathname + path,
-      protocol: endpointUrl.protocol,
-      port: endpointUrl.port ? parseInt(endpointUrl.port) : undefined,
-    });
-    try {
-      const { response } = await handler.handle(request, {} as HttpHandlerOptions);
-      if (response.statusCode === 200 && response.body) {
-        // handle response.body as stream
-        return sdkStreamMixin(response.body).transformToString();
-      } else {
-        throw Object.assign(new Error(`Request failed with status code ${response.statusCode}`), {
-          $metadata: { httpStatusCode: response.statusCode },
-        });
+    }
+
+    throw lastError!;
+  }
+
+  private shouldNotRetry(error: any): boolean {
+    // 400/403 errors for token fetch MUST NOT be retried
+    // 404 errors for metadata fetch MUST NOT be retried
+    const statusCode = error.statusCode || error.$metadata?.httpStatusCode;
+    return statusCode === 400 || statusCode === 403 || statusCode === 404;
+  }
+
+  async request(path: string, options: { method?: string; headers?: Record<string, string> }): Promise<string> {
+    return this.retryWithBackoff(async () => {
+      const { endpoint, ec2MetadataV1Disabled, httpOptions } = await this.config;
+      const handler = new NodeHttpHandler({
+        requestTimeout: httpOptions?.timeout,
+        throwOnRequestTimeout: true,
+        connectionTimeout: httpOptions?.timeout,
+      });
+      const endpointUrl = new URL(endpoint!);
+      const headers = options.headers || {};
+      /**
+       * If IMDSv1 is disabled and disableFetchToken is true, throw an error
+       * Refer: https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configuring-instance-metadata-service.html
+       */
+      if (this.disableFetchToken && ec2MetadataV1Disabled) {
+        throw new Error("IMDSv1 is disabled and fetching token is disabled, cannot make the request.");
       }
-    } catch (error) {
-      const wrappedError = new Error(`Error making request to the metadata service: ${error}`);
-      const { $metadata } = error as any;
-      if ($metadata?.httpStatusCode !== undefined) {
-        Object.assign(wrappedError, { $metadata });
+      /**
+       * Make request with token if disableFetchToken is not true (IMDSv2).
+       * Note that making the request call with token will result in an additional request to fetch the token.
+       */
+      if (!this.disableFetchToken) {
+        try {
+          headers["x-aws-ec2-metadata-token"] = await this.fetchMetadataTokenInternal();
+        } catch (err) {
+          if (ec2MetadataV1Disabled) {
+            // If IMDSv1 is disabled and token fetch fails (IMDSv2 fails), rethrow the error
+            throw err;
+          }
+          // If token fetch fails and IMDSv1 is not disabled, proceed without token (IMDSv1 fallback)
+        }
+      } // else, IMDSv1 fallback mode
+      const request = new HttpRequest({
+        method: options.method || "GET", // Default to GET if no method is specified
+        headers: headers,
+        hostname: endpointUrl.hostname,
+        path: endpointUrl.pathname + path,
+        protocol: endpointUrl.protocol,
+        port: endpointUrl.port ? parseInt(endpointUrl.port) : undefined,
+      });
+      try {
+        const { response } = await handler.handle(request, {} as HttpHandlerOptions);
+        if (response.statusCode === 200 && response.body) {
+          // handle response.body as stream
+          return sdkStreamMixin(response.body).transformToString();
+        } else {
+          throw Object.assign(new Error(`Request failed with status code ${response.statusCode}`), {
+            $metadata: { httpStatusCode: response.statusCode },
+          });
+        }
+      } catch (error) {
+        const wrappedError = new Error(`Error making request to the metadata service: ${error}`);
+        const { $metadata } = error as any;
+        if ($metadata?.httpStatusCode !== undefined) {
+          Object.assign(wrappedError, { $metadata });
+        }
+        throw wrappedError;
       }
-      throw wrappedError;
-    }
+    });
   }
 
   async fetchMetadataToken(): Promise<string> {
+    return this.retryWithBackoff(() => this.fetchMetadataTokenInternal());
+  }
+
+  private async fetchMetadataTokenInternal(): Promise<string> {
     /**
      * Refer: https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instance-metadata-v2-how-it-works.html
      */
@@ -124,11 +189,14 @@ export class MetadataService {
         return bodyString;
       } else {
         throw Object.assign(new Error(`Failed to fetch metadata token with status code ${response.statusCode}`), {
-          statusCode: response.statusCode,
+          $metadata: { httpStatusCode: response.statusCode },
         });
       }
     } catch (error) {
-      if (error.message === "TimeoutError" || [403, 404, 405].includes(error.statusCode)) {
+      if (
+        error.message === "TimeoutError" ||
+        [403, 404, 405].includes((error as any).statusCode || (error as any).$metadata?.httpStatusCode)
+      ) {
         this.disableFetchToken = true; // as per JSv2 and fromInstanceMetadata implementations
         throw new Error(`Error fetching metadata token: ${error}. [disableFetchToken] is now set to true.`);
       }

packages/ec2-metadata-service/src/MetadataServiceOptions.ts

@@ -27,4 +27,12 @@ export interface MetadataServiceOptions {
    * when true, metadata service will not fetch token, which indicates usage of IMDSv1
    */
   disableFetchToken?: boolean;
+  /**
+   * the number of retry attempts for any failed request, defaulting to 3.
+   */
+  retries?: number;
+  /**
+   * the number of seconds to sleep in-between retries and/or a customer provided backoff function to call.
+   */
+  backoff?: number | ((numFailures: number) => Promise<void> | number);
 }