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 <30836115+duncanista@users.noreply.github.com> --------- Co-authored-by: jordan gonzález <30836115+duncanista@users.noreply.github.com>

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") {