Step Functions Legacy Lambda Span Linking (#567)

* logic to parse payload from legacy lambda event

* fix linter

* add legacy lambda unit test case

* fix lint

* added more legacy lambda unit test cases

* fix lint

* Remove constant

Co-authored-by: jordan gonzález <[email protected]>

---------

Co-authored-by: jordan gonzález <[email protected]>

Author: avedmala

src/trace/context/extractor.spec.ts

@@ -907,6 +907,43 @@ describe("TraceContextExtractor", () => {
 
       expect(extractor).toBeInstanceOf(StepFunctionEventTraceExtractor);
     });
+
+    it("returns StepFunctionEventTraceExtractor when event contains legacy lambda StepFunctionContext", () => {
+      const event = {
+        Payload: {
+          Execution: {
+            Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
+            Input: {
+              MyInput: "MyValue",
+            },
+            Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be",
+            RoleArn:
+              "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03",
+            StartTime: "2022-12-08T21:08:17.924Z",
+          },
+          State: {
+            Name: "step-one",
+            EnteredTime: "2022-12-08T21:08:19.224Z",
+            RetryCount: 2,
+          },
+          StateMachine: {
+            Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
+            Name: "my-state-machine",
+          },
+        },
+      };
+
+      const tracerWrapper = new TracerWrapper();
+      const traceContextExtractor = new TraceContextExtractor(tracerWrapper, {} as TraceConfig);
+
+      // Mimick TraceContextService.extract initialization
+      const instance = StepFunctionContextService.instance(event);
+      traceContextExtractor["stepFunctionContextService"] = instance;
+
+      const extractor = traceContextExtractor["getTraceEventExtractor"](event);
+
+      expect(extractor).toBeInstanceOf(StepFunctionEventTraceExtractor);
+    });
   });
 
   describe("addTraceContexToXray", () => {

src/trace/context/extractors/step-function.spec.ts

@@ -55,6 +55,22 @@ describe("StepFunctionEventTraceExtractor", () => {
       expect(traceContext?.source).toBe("event");
     });
 
+    it("extracts trace context with valid legacy lambda payload", () => {
+      // Mimick TraceContextService.extract initialization
+      StepFunctionContextService.instance({ Payload: payload });
+
+      const extractor = new StepFunctionEventTraceExtractor();
+
+      // Payload is sent again for safety in case the instance wasn't previously initialized
+      const traceContext = extractor.extract({ Payload: payload });
+      expect(traceContext).not.toBeNull();
+
+      expect(traceContext?.toTraceId()).toBe("1139193989631387307");
+      expect(traceContext?.toSpanId()).toBe("5892738536804826142");
+      expect(traceContext?.sampleMode()).toBe("1");
+      expect(traceContext?.source).toBe("event");
+    });
+
     it("returns null when StepFunctionContextService.context is undefined", async () => {
       const extractor = new StepFunctionEventTraceExtractor();
 

src/trace/step-function-service.spec.ts

@@ -209,6 +209,21 @@ describe("StepFunctionContextService", () => {
 
       expect(spanContext).toBeNull();
     });
+
+    it("returns a SpanContextWrapper when event is from legacy lambda", () => {
+      const instance = StepFunctionContextService.instance();
+      // Force setting event
+      instance["setContext"]({ Payload: stepFunctionEvent });
+
+      const spanContext = instance.spanContext;
+
+      expect(spanContext).not.toBeNull();
+
+      expect(spanContext?.toTraceId()).toBe("1139193989631387307");
+      expect(spanContext?.toSpanId()).toBe("5892738536804826142");
+      expect(spanContext?.sampleMode()).toBe("1");
+      expect(spanContext?.source).toBe("event");
+    });
   });
 
   describe("deterministicSha256HashToBigIntString", () => {

src/trace/step-function-service.ts

@@ -41,6 +41,11 @@ export class StepFunctionContextService {
     // always triggered by the same event.
     if (typeof event !== "object") return;
 
+    // Legacy lambda parsing
+    if (typeof event.Payload === "object") {
+      event = event.Payload;
+    }
+
     // Execution
     const execution = event.Execution;
     if (typeof execution !== "object") {