chore(middleware-recursion-detection): read trace id from Labda Invoke Store

Author: trivikr

packages/middleware-recursion-detection/src/index.spec.ts

@@ -1,3 +1,4 @@
+import { InvokeStore } from "@aws/lambda-invoke-store";
 import { HttpRequest } from "@smithy/protocol-http";
 import { afterAll, beforeEach, describe, expect, test as it, vi } from "vitest";
 
@@ -6,31 +7,71 @@ import { recursionDetectionMiddleware } from "./index";
 describe(recursionDetectionMiddleware.name, () => {
   const mockNextHandler = vi.fn();
   const originEnv = process.env;
+
   const TRACE_ID_HEADER_NAME = "X-Amzn-Trace-Id";
 
   beforeEach(() => {
     vi.clearAllMocks();
+    vi.spyOn(InvokeStore, "getXRayTraceId").mockImplementation(() => undefined);
     process.env = {};
   });
 
   afterAll(() => {
     process.env = originEnv;
   });
 
-  it(`sets ${TRACE_ID_HEADER_NAME} header when function name and trace id environmental variables are set`, async () => {
-    process.env = {
-      AWS_LAMBDA_FUNCTION_NAME: "some-function",
-      _X_AMZN_TRACE_ID: "some-trace-id",
-    };
-    const handler = recursionDetectionMiddleware({ runtime: "node" })(mockNextHandler, {} as any);
-    await handler({
-      input: {},
-      request: new HttpRequest({}),
+  describe(`sets ${TRACE_ID_HEADER_NAME} header when function name and`, () => {
+    const mockTraceIdEnv = "trace-id-from-env";
+    const mockTraceIdInvokeStore = "trace-id-from-invoke-store";
+
+    it("trace id environmental variables is set", async () => {
+      process.env = {
+        AWS_LAMBDA_FUNCTION_NAME: "some-function",
+        _X_AMZN_TRACE_ID: mockTraceIdEnv,
+      };
+      const handler = recursionDetectionMiddleware({ runtime: "node" })(mockNextHandler, {} as any);
+      await handler({
+        input: {},
+        request: new HttpRequest({}),
+      });
+      const { calls } = (mockNextHandler as any).mock;
+      expect(calls.length).toBe(1);
+      const { request } = mockNextHandler.mock.calls[0][0];
+      expect(request.headers[TRACE_ID_HEADER_NAME]).toBe(mockTraceIdEnv);
+    });
+
+    it("trace id value is set in InvokeStore", async () => {
+      vi.spyOn(InvokeStore, "getXRayTraceId").mockImplementation(() => mockTraceIdInvokeStore);
+      process.env = {
+        AWS_LAMBDA_FUNCTION_NAME: "some-function",
+      };
+      const handler = recursionDetectionMiddleware({ runtime: "node" })(mockNextHandler, {} as any);
+      await handler({
+        input: {},
+        request: new HttpRequest({}),
+      });
+      const { calls } = (mockNextHandler as any).mock;
+      expect(calls.length).toBe(1);
+      const { request } = mockNextHandler.mock.calls[0][0];
+      expect(request.headers[TRACE_ID_HEADER_NAME]).toBe(mockTraceIdInvokeStore);
+    });
+
+    it("favors trace id value from InvokeStore over that from env variable", async () => {
+      vi.spyOn(InvokeStore, "getXRayTraceId").mockImplementation(() => mockTraceIdInvokeStore);
+      process.env = {
+        AWS_LAMBDA_FUNCTION_NAME: "some-function",
+        _X_AMZN_TRACE_ID: mockTraceIdEnv,
+      };
+      const handler = recursionDetectionMiddleware({ runtime: "node" })(mockNextHandler, {} as any);
+      await handler({
+        input: {},
+        request: new HttpRequest({}),
+      });
+      const { calls } = (mockNextHandler as any).mock;
+      expect(calls.length).toBe(1);
+      const { request } = mockNextHandler.mock.calls[0][0];
+      expect(request.headers[TRACE_ID_HEADER_NAME]).toBe(mockTraceIdInvokeStore);
     });
-    const { calls } = (mockNextHandler as any).mock;
-    expect(calls.length).toBe(1);
-    const { request } = mockNextHandler.mock.calls[0][0];
-    expect(request.headers[TRACE_ID_HEADER_NAME]).toBe("some-trace-id");
   });
 
   it(`should NOT set ${TRACE_ID_HEADER_NAME} header when function name environmental variable is NOT set`, async () => {

packages/middleware-recursion-detection/src/index.ts

@@ -1,3 +1,4 @@
+import { InvokeStore } from "@aws/lambda-invoke-store";
 import { HttpRequest } from "@smithy/protocol-http";
 import {
   AbsoluteLocation,
@@ -38,7 +39,11 @@ export const recursionDetectionMiddleware =
       return next(args);
     }
     const functionName = process.env[ENV_LAMBDA_FUNCTION_NAME];
-    const traceId = process.env[ENV_TRACE_ID];
+
+    const traceIdFromEnv = process.env[ENV_TRACE_ID];
+    const traceIdFromInvokeStore = InvokeStore.getXRayTraceId();
+    const traceId = traceIdFromInvokeStore ?? traceIdFromEnv;
+
     const nonEmptyString = (str: unknown): str is string => typeof str === "string" && str.length > 0;
     if (nonEmptyString(functionName) && nonEmptyString(traceId)) {
       request.headers[TRACE_ID_HEADER_NAME] = traceId;