diff --git a/src/trace/context/extractor.spec.ts b/src/trace/context/extractor.spec.ts index f41a645e..1bdb4430 100644 --- a/src/trace/context/extractor.spec.ts +++ b/src/trace/context/extractor.spec.ts @@ -14,6 +14,7 @@ import { EventBridgeEventTraceExtractor, EventBridgeSQSEventTraceExtractor, HTTPEventTraceExtractor, + KafkaEventTraceExtractor, KinesisEventTraceExtractor, SNSEventTraceExtractor, SNSSQSEventTraceExtractor, @@ -906,6 +907,24 @@ describe("getTraceEventExtractor", () => { ], }, ], + [ + "KafkaEventTraceExtractor", + "Kafka event", + KafkaEventTraceExtractor, + { + eventSource: "aws:kafka", + records: {}, + }, + ], + [ + "KafkaEventTraceExtractor", + "self-managed Kafka event", + KafkaEventTraceExtractor, + { + eventSource: "aws:SelfManagedKafka", + records: {}, + }, + ], [ "EventBridgeEventTraceExtractor", "EventBridge event", diff --git a/src/trace/context/extractor.ts b/src/trace/context/extractor.ts index 145ae0c9..e551ff6f 100644 --- a/src/trace/context/extractor.ts +++ b/src/trace/context/extractor.ts @@ -9,6 +9,7 @@ import { EventBridgeEventTraceExtractor, EventBridgeSQSEventTraceExtractor, HTTPEventTraceExtractor, + KafkaEventTraceExtractor, KinesisEventTraceExtractor, LambdaContextTraceExtractor, SNSEventTraceExtractor, @@ -98,6 +99,7 @@ export class TraceContextExtractor { if (EventValidator.isSQSEvent(event)) return new SQSEventTraceExtractor(this.tracerWrapper, this.config); if (EventValidator.isKinesisStreamEvent(event)) return new KinesisEventTraceExtractor(this.tracerWrapper, this.config); + if (EventValidator.isKafkaEvent(event)) return new KafkaEventTraceExtractor(this.tracerWrapper, this.config); if (EventValidator.isEventBridgeEvent(event)) return new EventBridgeEventTraceExtractor(this.tracerWrapper, this.config); diff --git a/src/trace/context/extractors/index.ts b/src/trace/context/extractors/index.ts index cd5c757d..2a5b6ebc 100644 --- a/src/trace/context/extractors/index.ts +++ b/src/trace/context/extractors/index.ts @@ -1,6 +1,7 @@ export { AppSyncEventTraceExtractor } from "./app-sync"; export { EventBridgeEventTraceExtractor } from "./event-bridge"; export { EventBridgeSQSEventTraceExtractor } from "./event-bridge-sqs"; +export { KafkaEventTraceExtractor } from "./kafka"; export { KinesisEventTraceExtractor } from "./kinesis"; export { HTTPEventTraceExtractor } from "./http"; export { SQSEventTraceExtractor } from "./sqs"; diff --git a/src/trace/context/extractors/kafka.spec.ts b/src/trace/context/extractors/kafka.spec.ts new file mode 100644 index 00000000..400b9a1f --- /dev/null +++ b/src/trace/context/extractors/kafka.spec.ts @@ -0,0 +1,237 @@ +import { TracerWrapper } from "../../tracer-wrapper"; +import { KafkaEventTraceExtractor } from "./kafka"; + +let mockSpanContext: any = null; +let mockDataStreamsCheckpointer: any = { + setConsumeCheckpoint: jest.fn(), +}; + +jest.mock("dd-trace/packages/dd-trace/src/datastreams/checkpointer", () => { + return { + DataStreamsCheckpointer: jest.fn().mockImplementation(() => mockDataStreamsCheckpointer), + }; +}); + +// Mocking extract is needed, due to dd-trace being a No-op +// if the detected environment is testing. This is expected, since +// we don't want to test dd-trace extraction, but our components. +jest.mock("dd-trace", () => { + return { + ...jest.requireActual("dd-trace"), + _tracer: { _service: {} }, + extract: (_carrier: any, _headers: any) => mockSpanContext, + dataStreamsCheckpointer: mockDataStreamsCheckpointer, + }; +}); +const spyTracerWrapper = jest.spyOn(TracerWrapper.prototype, "extract"); + +/** Encodes a header value the way Kafka/Lambda delivers it: a byte array. */ +const toBytes = (value: string): number[] => Array.from(Buffer.from(value, "utf8")); + +const PATHWAY_CTX = "4eia3s7L38Gs9qeLlGis9qeLlGg="; + +const buildRecord = (overrides: Partial = {}): any => ({ + topic: "demo-topic", + partition: 0, + offset: "5", + timestamp: 1789402611155, + timestampType: "CREATE_TIME", + key: Buffer.from("some-key").toString("base64"), + value: Buffer.from(JSON.stringify({ id: "abc" })).toString("base64"), + headers: [ + { "content-type": toBytes("application/json") }, + { "x-datadog-trace-id": toBytes("6043338675393224912") }, + { "x-datadog-parent-id": toBytes("2162012466251825722") }, + { "x-datadog-sampling-priority": toBytes("1") }, + { "dd-pathway-ctx-base64": toBytes(PATHWAY_CTX) }, + ], + ...overrides, +}); + +const buildEvent = (records: Record): any => ({ + eventSource: "aws:kafka", + eventSourceArn: "arn:aws:kafka:us-east-2:123456789012:cluster/demo-cluster/39027fc5-2658-4c2d-a102-4dcd3a10301b-s2", + bootstrapServers: "boot-abc123.c2.kafka-serverless.us-east-2.amazonaws.com:9098", + records, +}); + +describe("KafkaEventTraceExtractor", () => { + const mockConfig = { + autoPatchHTTP: true, + captureLambdaPayload: false, + captureLambdaPayloadMaxDepth: 10, + createInferredSpan: true, + encodeAuthorizerContext: true, + decodeAuthorizerContext: true, + mergeDatadogXrayTraces: false, + injectLogContext: false, + minColdStartTraceDuration: 3, + coldStartTraceSkipLib: "", + addSpanPointers: true, + dataStreamsEnabled: true, + appsecEnabled: false, + }; + + describe("extract", () => { + beforeEach(() => { + mockSpanContext = null; + mockDataStreamsCheckpointer.setConsumeCheckpoint.mockClear(); + }); + + it("extracts trace context with valid payload", () => { + mockSpanContext = { + toTraceId: () => "6043338675393224912", + toSpanId: () => "2162012466251825722", + _sampling: { priority: "1" }, + }; + const tracerWrapper = new TracerWrapper(); + + const payload = buildEvent({ "demo-topic-0": [buildRecord()] }); + + const extractor = new KafkaEventTraceExtractor(tracerWrapper, mockConfig); + const traceContext = extractor.extract(payload); + expect(traceContext).not.toBeNull(); + + // Headers are decoded from byte arrays into a flat string carrier. + expect(spyTracerWrapper).toHaveBeenCalledWith({ + "content-type": "application/json", + "x-datadog-trace-id": "6043338675393224912", + "x-datadog-parent-id": "2162012466251825722", + "x-datadog-sampling-priority": "1", + "dd-pathway-ctx-base64": PATHWAY_CTX, + }); + + expect(traceContext?.toTraceId()).toBe("6043338675393224912"); + expect(traceContext?.toSpanId()).toBe("2162012466251825722"); + expect(traceContext?.sampleMode()).toBe("1"); + expect(traceContext?.source).toBe("event"); + }); + + it("sets the DSM consume checkpoint with the topic name, not the cluster ARN", () => { + const tracerWrapper = new TracerWrapper(); + const payload = buildEvent({ "demo-topic-0": [buildRecord()] }); + + const extractor = new KafkaEventTraceExtractor(tracerWrapper, mockConfig); + extractor.extract(payload); + + // The DSM target must match the `topic:` tag dd-trace's kafkajs + // producer plugin emits. Passing the ARN would break the pathway. + expect(mockDataStreamsCheckpointer.setConsumeCheckpoint).toHaveBeenCalledWith( + "kafka", + "demo-topic", + expect.objectContaining({ "dd-pathway-ctx-base64": PATHWAY_CTX }), + false, + ); + }); + + it("sets one checkpoint per record across multiple topic-partitions", () => { + const tracerWrapper = new TracerWrapper(); + const payload = buildEvent({ + "demo-topic-0": [buildRecord({ offset: "1" }), buildRecord({ offset: "2" })], + "other-topic-3": [buildRecord({ topic: "other-topic", partition: 3, offset: "9" })], + }); + + const extractor = new KafkaEventTraceExtractor(tracerWrapper, mockConfig); + extractor.extract(payload); + + expect(mockDataStreamsCheckpointer.setConsumeCheckpoint).toHaveBeenCalledTimes(3); + const targets = mockDataStreamsCheckpointer.setConsumeCheckpoint.mock.calls.map((c: any[]) => c[1]); + expect(targets).toEqual(["demo-topic", "demo-topic", "other-topic"]); + }); + + it("does not set checkpoints when Data Streams is disabled", () => { + const tracerWrapper = new TracerWrapper(); + const payload = buildEvent({ "demo-topic-0": [buildRecord()] }); + + const extractor = new KafkaEventTraceExtractor(tracerWrapper, { + ...mockConfig, + dataStreamsEnabled: false, + }); + extractor.extract(payload); + + expect(mockDataStreamsCheckpointer.setConsumeCheckpoint).not.toHaveBeenCalled(); + }); + + it("recovers the topic from the map key when the record omits it", () => { + const tracerWrapper = new TracerWrapper(); + const record = buildRecord(); + delete record.topic; + // Topic names may themselves contain hyphens; only the trailing + // partition number is trimmed. + const payload = buildEvent({ "orders-eu-west-12": [record] }); + + const extractor = new KafkaEventTraceExtractor(tracerWrapper, mockConfig); + extractor.extract(payload); + + expect(mockDataStreamsCheckpointer.setConsumeCheckpoint).toHaveBeenCalledWith( + "kafka", + "orders-eu-west", + expect.any(Object), + false, + ); + }); + + it("handles byte arrays delivered as numeric strings", () => { + mockSpanContext = { + toTraceId: () => "6043338675393224912", + toSpanId: () => "2162012466251825722", + _sampling: { priority: "1" }, + }; + const tracerWrapper = new TracerWrapper(); + + const record = buildRecord({ + headers: [{ "dd-pathway-ctx-base64": toBytes(PATHWAY_CTX).map(String) }], + }); + const payload = buildEvent({ "demo-topic-0": [record] }); + + const extractor = new KafkaEventTraceExtractor(tracerWrapper, mockConfig); + extractor.extract(payload); + + expect(mockDataStreamsCheckpointer.setConsumeCheckpoint).toHaveBeenCalledWith( + "kafka", + "demo-topic", + { "dd-pathway-ctx-base64": PATHWAY_CTX }, + false, + ); + }); + + it("returns null when payload is empty", () => { + const tracerWrapper = new TracerWrapper(); + const payload = buildEvent({}); + + const extractor = new KafkaEventTraceExtractor(tracerWrapper, mockConfig); + const traceContext = extractor.extract(payload); + + expect(traceContext).toBeNull(); + expect(mockDataStreamsCheckpointer.setConsumeCheckpoint).not.toHaveBeenCalled(); + }); + + it("returns null but still checkpoints when a record carries no headers", () => { + const tracerWrapper = new TracerWrapper(); + const payload = buildEvent({ "demo-topic-0": [buildRecord({ headers: [] })] }); + + const extractor = new KafkaEventTraceExtractor(tracerWrapper, mockConfig); + const traceContext = extractor.extract(payload); + + expect(traceContext).toBeNull(); + + // Records from uninstrumented producers have no context to extract, but + // they must still produce a DSM consume node as a pathway root, matching + // the SQS/Kinesis extractors' null-carrier behaviour. + expect(mockDataStreamsCheckpointer.setConsumeCheckpoint).toHaveBeenCalledWith("kafka", "demo-topic", null, false); + }); + + it("returns null when headers contain no trace context", () => { + const tracerWrapper = new TracerWrapper(); + const payload = buildEvent({ + "demo-topic-0": [buildRecord({ headers: [{ "content-type": toBytes("application/json") }] })], + }); + + const extractor = new KafkaEventTraceExtractor(tracerWrapper, mockConfig); + const traceContext = extractor.extract(payload); + + // mockSpanContext is null -> extract yields nothing. + expect(traceContext).toBeNull(); + }); + }); +}); diff --git a/src/trace/context/extractors/kafka.ts b/src/trace/context/extractors/kafka.ts new file mode 100644 index 00000000..3608064a --- /dev/null +++ b/src/trace/context/extractors/kafka.ts @@ -0,0 +1,154 @@ +import { MSKEvent, MSKRecord, SelfManagedKafkaEvent, SelfManagedKafkaRecord } from "aws-lambda"; +import { logDebug } from "../../../utils"; +import { EventTraceExtractor } from "../extractor"; +import { TracerWrapper } from "../../tracer-wrapper"; +import { SpanContextWrapper } from "../../span-context-wrapper"; +import { TraceConfig } from "../../listener"; + +/** + * Extracts trace context and sets DSM consume checkpoints for Kafka events + * delivered by a Lambda event source mapping (Amazon MSK, MSK Serverless, and + * self-managed Kafka). + * + * Kafka differs from the SQS/SNS/Kinesis extractors in three ways: + * + * 1. Records arrive in `event.records`, an object keyed by "-" + * whose values are arrays of records -- not a flat `event.Records` array. + * 2. Context lives in real Kafka message headers (`record.headers`), which are + * an array of single-key objects whose values are byte arrays. They must be + * decoded to UTF-8 to form a carrier. There is no `_datadog` payload + * envelope, so the payload is never parsed. + * 3. The DSM target must be the **topic name**, not the cluster ARN. + * `dd-trace`'s kafkajs producer plugin tags produce edges with + * `topic:`; `setConsumeCheckpoint` turns its `source` argument into + * `topic:`. Passing the ARN (as the Kinesis/SQS extractors do) would + * tag the consume edge with the ARN and the pathway would not connect to the + * producer's edge in the DSM graph. + */ +type KafkaEvent = MSKEvent | SelfManagedKafkaEvent; +type KafkaRecord = MSKRecord | SelfManagedKafkaRecord; + +export class KafkaEventTraceExtractor implements EventTraceExtractor { + constructor(private tracerWrapper: TracerWrapper, private config: TraceConfig) {} + + extract(event: KafkaEvent): SpanContextWrapper | null { + // Set DSM consume checkpoints if enabled, and capture the first record's + // headers for trace context extraction. + let firstRecordHeaders: Record | null = null; + const topicPartitions = Object.keys(event?.records ?? {}); + + if (this.config.dataStreamsEnabled) { + let isFirstRecord = true; + + for (const topicPartition of topicPartitions) { + for (const record of event.records[topicPartition] ?? []) { + try { + const headers = this.getParsedRecordHeaders(record); + + if (isFirstRecord) { + firstRecordHeaders = headers; + isFirstRecord = false; + } + + // A batch can span topics, so the target is resolved per record + // rather than once for the whole event. + const topic = this.getTopicName(record, topicPartition); + if (topic) { + this.tracerWrapper.setConsumeCheckpoint(headers, "kafka", topic); + } else { + logDebug("Skipping DSM checkpoint for Kafka record: unable to resolve topic name", { + topicPartition, + }); + } + } catch (error) { + if (error instanceof Error) { + logDebug("Unable to set DSM checkpoint for Kafka event", error); + } + } + } + } + } + + const firstRecord = this.getFirstRecord(event); + if (firstRecord === undefined) return null; + + try { + // Reuse the headers already decoded for DSM when available. + if (!firstRecordHeaders) { + firstRecordHeaders = this.getParsedRecordHeaders(firstRecord); + } + + if (firstRecordHeaders) { + const traceContext = this.tracerWrapper.extract(firstRecordHeaders); + if (traceContext === null) return null; + + logDebug(`Extracted trace context from Kafka event`, { + traceContext, + headers: firstRecordHeaders, + }); + return traceContext; + } + } catch (error) { + if (error instanceof Error) { + logDebug("Unable to extract trace context from Kafka event", error); + } + } + + return null; + } + + private getFirstRecord(event: KafkaEvent): KafkaRecord | undefined { + for (const topicPartition of Object.keys(event?.records ?? {})) { + const records = event.records[topicPartition]; + if (Array.isArray(records) && records.length > 0) { + return records[0]; + } + } + return undefined; + } + + /** + * Prefers the record's own `topic` field. Falls back to parsing the + * "-" map key, trimming only the trailing partition number + * so topics containing hyphens survive. + */ + private getTopicName(record: KafkaRecord | undefined, topicPartition: string): string | undefined { + if (record?.topic) return record.topic; + const match = /^(.*)-\d+$/.exec(topicPartition); + return match?.[1] ?? undefined; + } + + /** + * Kafka headers arrive as an array of single-key objects whose values are + * byte arrays, e.g. `[{ "dd-pathway-ctx-base64": [52, 101, ...] }]`. + * Flattened to a `Record` carrier that both + * `tracerWrapper.extract` and the DSM codec can read. Header names are left + * as-is: dd-trace reads the lowercase keys Kafka producers already emit. + */ + private getParsedRecordHeaders(record: KafkaRecord | undefined): Record | null { + if (!record?.headers) { + return null; + } + + try { + const headers: Record = {}; + + for (const header of record.headers) { + if (header === null || typeof header !== "object") continue; + + for (const [name, value] of Object.entries(header)) { + if (value === null || value === undefined) continue; + + // Byte arrays may arrive as numbers or numeric strings depending on + // how the event was serialized in transit. + const bytes = Array.isArray(value) ? value.map((byte) => Number(byte)) : value; + headers[name] = Buffer.from(bytes as any).toString("utf8"); + } + } + + return Object.keys(headers).length > 0 ? headers : null; + } catch (error) { + return null; + } + } +} diff --git a/src/utils/event-validator.ts b/src/utils/event-validator.ts index b2ba6ea9..fda6daa9 100644 --- a/src/utils/event-validator.ts +++ b/src/utils/event-validator.ts @@ -4,6 +4,8 @@ import { AppSyncResolverEvent, EventBridgeEvent, KinesisStreamEvent, + MSKEvent, + SelfManagedKafkaEvent, SNSEvent, SNSMessage, SQSEvent, @@ -75,6 +77,14 @@ export class EventValidator { return Array.isArray(event.Records) && event.Records.length > 0 && event.Records[0].kinesis !== undefined; } + static isKafkaEvent(event: any): event is MSKEvent | SelfManagedKafkaEvent { + return ( + (event.eventSource === "aws:kafka" || event.eventSource === "aws:SelfManagedKafka") && + event.records !== null && + typeof event.records === "object" + ); + } + static isDurableExecutionEvent(event: any): boolean { return typeof event.DurableExecutionArn === "string"; }