feat: extract trace context from MSK events

src/trace/context/extractor.spec.ts

@@ -15,6 +15,7 @@ import {
   EventBridgeSQSEventTraceExtractor,
   HTTPEventTraceExtractor,
   KinesisEventTraceExtractor,
+  MSKEventTraceExtractor,
   SNSEventTraceExtractor,
   SNSSQSEventTraceExtractor,
   SQSEventTraceExtractor,
@@ -812,6 +813,11 @@ describe("getTraceEventExtractor", () => {
     ["a string", "some-value"],
     ["a number", 1234],
     ["an object which doesn't match any expected event", { custom: "event" }],
+    ["MSK without records", { eventSource: "aws:kafka" }],
+    ["MSK with null records", { eventSource: "aws:kafka", records: null }],
+    ["MSK with array records", { eventSource: "aws:kafka", records: [] }],
+    ["MSK with string records", { eventSource: "aws:kafka", records: "invalid" }],
+    ["records without an MSK event source", { records: { "topic-0": [] } }],
   ])("returns undefined when event is '%s'", (_, event) => {
     const tracerWrapper = new TracerWrapper();
     const traceContextExtractor = new TraceContextExtractor(tracerWrapper, {} as TraceConfig);
@@ -894,6 +900,12 @@ describe("getTraceEventExtractor", () => {
         ],
       },
     ],
+    [
+      "MSKEventTraceExtractor",
+      "MSK event",
+      MSKEventTraceExtractor,
+      { eventSource: "aws:kafka", records: { "topic-0": [] } },
+    ],
     [
       "KinesisEventTraceExtractor",
       "Kinesis stream event",

src/trace/context/extractor.ts

@@ -10,6 +10,7 @@ import {
   EventBridgeSQSEventTraceExtractor,
   HTTPEventTraceExtractor,
   KinesisEventTraceExtractor,
+  MSKEventTraceExtractor,
   LambdaContextTraceExtractor,
   SNSEventTraceExtractor,
   SNSSQSEventTraceExtractor,
@@ -95,6 +96,8 @@ export class TraceContextExtractor {
     if (EventValidator.isEventBridgeSQSEvent(event))
       return new EventBridgeSQSEventTraceExtractor(this.tracerWrapper, this.config);
     if (EventValidator.isAppSyncResolverEvent(event)) return new AppSyncEventTraceExtractor(this.tracerWrapper);
+    if (EventValidator.isMSKEvent(event)) return new MSKEventTraceExtractor(this.tracerWrapper);
+
     if (EventValidator.isSQSEvent(event)) return new SQSEventTraceExtractor(this.tracerWrapper, this.config);
     if (EventValidator.isKinesisStreamEvent(event))
       return new KinesisEventTraceExtractor(this.tracerWrapper, this.config);

src/trace/context/extractors/index.ts

@@ -10,3 +10,4 @@ export { StepFunctionEventTraceExtractor } from "./step-function";
 export { LambdaContextTraceExtractor } from "./lambda-context";
 export { CustomTraceExtractor } from "./custom";
 export { DurableExecutionEventTraceExtractor } from "./durable-execution";
+export { MSKEventTraceExtractor } from "./msk";

src/trace/context/extractors/msk.spec.ts

@@ -0,0 +1,122 @@
+import { MSKEvent } from "aws-lambda";
+import { MSKEventTraceExtractor } from "./msk";
+import { TracerWrapper } from "../../tracer-wrapper";
+import { SpanContextWrapper } from "../../span-context-wrapper";
+import { TraceSource } from "../../trace-context-service";
+
+const datadogHeaders = {
+  "x-datadog-trace-id": "667309514221035538",
+  "x-datadog-parent-id": "1350735035497811828",
+  "x-datadog-sampling-priority": "1",
+  "x-datadog-tags": "_dd.p.tid=0123456789abcdef",
+};
+const w3cHeaders = {
+  traceparent: "00-0123456789abcdef0942c4f3b84c8812-12becbb3b2d5bd74-01",
+  tracestate: "dd=s:1;t.dm:-0",
+};
+const record = (headers: Record<string, string>) => ({
+  headers: Object.entries(headers).map(([name, value]) => ({ [name]: Array.from(Buffer.from(value)) })),
+});
+const event = (records: any): MSKEvent => ({ eventSource: "aws:kafka", records } as MSKEvent);
+
+// Mock the tracer boundary so these tests exercise MSK decoding and record selection.
+describe("MSKEventTraceExtractor", () => {
+  const spanContext = new SpanContextWrapper(
+    { toTraceId: () => datadogHeaders["x-datadog-trace-id"], toSpanId: () => datadogHeaders["x-datadog-parent-id"] },
+    TraceSource.Event,
+  );
+  let extract: jest.Mock;
+  let extractor: MSKEventTraceExtractor;
+
+  beforeEach(() => {
+    extract = jest.fn().mockReturnValue(null);
+    extractor = new MSKEventTraceExtractor({ extract } as unknown as TracerWrapper);
+  });
+
+  it.each([
+    ["Datadog", datadogHeaders],
+    ["W3C", w3cHeaders],
+    ["combined", { ...datadogHeaders, ...w3cHeaders }],
+  ])("decodes and forwards all %s propagation headers", (_, headers) => {
+    extract.mockReturnValue(spanContext);
+    expect(extractor.extract(event({ "topic-0": [record(headers)] }))).toBe(spanContext);
+    expect(extract).toHaveBeenCalledWith(headers);
+    expect(extract).toHaveBeenCalledTimes(1);
+  });
+
+  it("normalizes header names and decodes UTF-8", () => {
+    extractor.extract(
+      event({ "topic-0": [record({ TraceParent: w3cHeaders.traceparent, baggage: "city=São Paulo" })] }),
+    );
+    expect(extract).toHaveBeenCalledWith({ traceparent: w3cHeaders.traceparent, baggage: "city=São Paulo" });
+  });
+
+  it("uses the last value for duplicate headers", () => {
+    const message = { headers: [...record({ traceparent: "old" }).headers, ...record(w3cHeaders).headers] };
+    extractor.extract(event({ "topic-0": [message] }));
+    expect(extract).toHaveBeenCalledWith(w3cHeaders);
+  });
+
+  it("skips untraced records across partitions and stops at the first valid context", () => {
+    extract.mockReturnValueOnce(null).mockReturnValueOnce(spanContext);
+    expect(
+      extractor.extract(
+        event({
+          "topic-0": [{ headers: [] }, record({ traceparent: "invalid" })],
+          "topic-1": [record(datadogHeaders), record(w3cHeaders)],
+        }),
+      ),
+    ).toBe(spanContext);
+    expect(extract).toHaveBeenCalledTimes(2);
+    expect(extract).toHaveBeenNthCalledWith(2, datadogHeaders);
+  });
+
+  it("never combines partial trace headers from different records", () => {
+    extractor.extract(
+      event({ "topic-0": [record({ "x-datadog-trace-id": "123" }), record({ "x-datadog-parent-id": "456" })] }),
+    );
+    expect(extract).toHaveBeenNthCalledWith(1, { "x-datadog-trace-id": "123" });
+    expect(extract).toHaveBeenNthCalledWith(2, { "x-datadog-parent-id": "456" });
+  });
+
+  it.each([null, "123", { "0": 49 }].map((value) => [value]))("skips non-array header values: %j", (value) => {
+    extract.mockReturnValue(spanContext);
+    expect(
+      extractor.extract(event({ "topic-0": [{ headers: [{ invalid: value }, ...record(datadogHeaders).headers] }] })),
+    ).toBe(spanContext);
+    expect(extract).toHaveBeenCalledWith(datadogHeaders);
+  });
+
+  it.each([
+    undefined,
+    null,
+    "invalid",
+    {},
+    { "topic-0": [] },
+    { "topic-0": null },
+    { "topic-0": "invalid" },
+    { "topic-0": [null, {}, { headers: null }, { headers: "invalid" }, { headers: [null, 42, []] }] },
+  ])("returns null for empty or malformed records: %j", (records) => {
+    expect(extractor.extract(event(records))).toBeNull();
+    expect(extract).not.toHaveBeenCalled();
+  });
+
+  it("returns null after an extraction error", () => {
+    extract.mockImplementationOnce(() => {
+      throw new Error("invalid carrier");
+    });
+    expect(extractor.extract(event({ "topic-0": [record(w3cHeaders), record(datadogHeaders)] }))).toBeNull();
+    expect(extract).toHaveBeenCalledTimes(1);
+  });
+
+  it("returns null when the tracer cannot extract context", () => {
+    expect(extractor.extract(event({ "topic-0": [record(datadogHeaders)] }))).toBeNull();
+  });
+
+  it("does not mutate the event", () => {
+    const payload = event({ "topic-0": [record(datadogHeaders)] });
+    const original = JSON.stringify(payload);
+    extractor.extract(payload);
+    expect(JSON.stringify(payload)).toBe(original);
+  });
+});

src/trace/context/extractors/msk.ts

@@ -0,0 +1,56 @@
+import { MSKEvent, MSKRecord } from "aws-lambda";
+import { logDebug } from "../../../utils";
+import { EventTraceExtractor } from "../extractor";
+import { handleExtractionError } from "../extractor-utils";
+import { TracerWrapper } from "../../tracer-wrapper";
+import { SpanContextWrapper } from "../../span-context-wrapper";
+
+export class MSKEventTraceExtractor implements EventTraceExtractor {
+  constructor(private tracerWrapper: TracerWrapper) {}
+
+  extract(event: MSKEvent): SpanContextWrapper | null {
+    if (!event?.records) {
+      logDebug("Failed to extract trace context from MSK event");
+      return null;
+    }
+
+    try {
+      // A Lambda span can have only one parent. Use the first record with valid
+      // trace context, without combining headers from different records.
+      for (const records of Object.values(event.records)) {
+        if (!Array.isArray(records)) continue;
+        for (const record of records) {
+          const headers = this.getParsedRecordHeaders(record);
+          if (!headers) continue;
+          const traceContext = this.tracerWrapper.extract(headers);
+          if (traceContext) {
+            logDebug("Extracted trace context from MSK event");
+            return traceContext;
+          }
+        }
+      }
+    } catch (error) {
+      handleExtractionError(error, "MSK");
+    }
+
+    logDebug("Failed to extract trace context from MSK event");
+    return null;
+  }
+
+  private getParsedRecordHeaders(record: MSKRecord): Record<string, string> | null {
+    if (!Array.isArray(record?.headers)) return null;
+
+    let headers: Record<string, string> | null = null;
+    for (const entry of record.headers) {
+      if (!entry || typeof entry !== "object" || Array.isArray(entry)) continue;
+      for (const [name, value] of Object.entries(entry)) {
+        // MSK serializes Kafka header bytes as integer arrays, not base64.
+        if (Array.isArray(value)) {
+          headers ??= {};
+          headers[name.toLowerCase()] = Buffer.from(value).toString("utf8");
+        }
+      }
+    }
+    return headers;
+  }
+}

src/utils/event-validator.ts

@@ -4,6 +4,7 @@ import {
   AppSyncResolverEvent,
   EventBridgeEvent,
   KinesisStreamEvent,
+  MSKEvent,
   SNSEvent,
   SNSMessage,
   SQSEvent,
@@ -75,6 +76,15 @@ export class EventValidator {
     return Array.isArray(event.Records) && event.Records.length > 0 && event.Records[0].kinesis !== undefined;
   }
 
+  static isMSKEvent(event: any): event is MSKEvent {
+    return (
+      event?.eventSource === "aws:kafka" &&
+      event.records !== null &&
+      typeof event.records === "object" &&
+      !Array.isArray(event.records)
+    );
+  }
+
   static isDurableExecutionEvent(event: any): boolean {
     return typeof event.DurableExecutionArn === "string";
   }