Skip to content

feat: lambda support for DSM #672

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 21 commits into
base: main
Choose a base branch
from
Draft

feat: lambda support for DSM #672

Show file tree
Hide file tree
Changes from 15 commits
Commits
Show all changes
21 commits
Select commit Hold shift + click to select a range
d0764da
set checkpoint for all supported dsm queue types
michael-zhao459 Jul 24, 2025
5d4da9f
removed not needed checkpoint set
michael-zhao459 Jul 28, 2025
4ea0470
clean
michael-zhao459 Jul 28, 2025
c2cefc2
clearer return type
michael-zhao459 Jul 28, 2025
2626672
add tests
michael-zhao459 Jul 28, 2025
ca9edb9
add the manual checkpoint parameter
michael-zhao459 Jul 28, 2025
d0e71e6
Correct approach to ensure public api is used
michael-zhao459 Jul 28, 2025
197abd0
remove unused variable
michael-zhao459 Jul 28, 2025
c9d6c23
match logging patterns of repo
michael-zhao459 Jul 28, 2025
75e976f
lint
michael-zhao459 Jul 28, 2025
a12d0ac
lint
michael-zhao459 Jul 28, 2025
99f1c67
add False manual_checkpoint false
michael-zhao459 Jul 29, 2025
7b4c4aa
fix formatting
michael-zhao459 Jul 29, 2025
5059c5c
optimization
michael-zhao459 Jul 29, 2025
7d3fdba
clean
michael-zhao459 Jul 29, 2025
c49e96f
add prettier ignore
michael-zhao459 Jul 29, 2025
e422375
add optional chaining
michael-zhao459 Jul 29, 2025
893d391
add log debug for no arn case
michael-zhao459 Jul 29, 2025
e44b6a7
Merge remote-tracking branch 'origin/main' into michael.zhao/lambda-dsm
michael-zhao459 Aug 13, 2025
bc8cb73
merge commit
michael-zhao459 Aug 13, 2025
70c5824
add checkpoints for headers
michael-zhao459 Aug 13, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 44 additions & 6 deletions src/trace/context/extractors/kinesis.spec.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,15 @@ import { TracerWrapper } from "../../tracer-wrapper";
import { KinesisEventTraceExtractor } from "./kinesis";

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
Expand All @@ -12,6 +21,7 @@ jest.mock("dd-trace", () => {
...ddTrace,
_tracer: { _service: {} },
extract: (_carrier: any, _headers: any) => mockSpanContext,
dataStreamsCheckpointer: mockDataStreamsCheckpointer,
};
});
const spyTracerWrapper = jest.spyOn(TracerWrapper.prototype, "extract");
Expand All @@ -20,10 +30,13 @@ describe("KinesisEventTraceExtractor", () => {
describe("extract", () => {
beforeEach(() => {
mockSpanContext = null;
mockDataStreamsCheckpointer.setConsumeCheckpoint.mockClear();
process.env["DD_DATA_STREAMS_ENABLED"] = "true";
});

afterEach(() => {
jest.resetModules();
delete process.env["DD_DATA_STREAMS_ENABLED"];
});

it("extracts trace context with valid payload", () => {
Expand All @@ -43,7 +56,7 @@ describe("KinesisEventTraceExtractor", () => {
kinesisSchemaVersion: "1.0",
partitionKey: "cdbfd750-cec0-4f0f-a4b0-82ae6152c7fb",
sequenceNumber: "49625698045709644136382874226371117765484751339579768834",
data: "eyJJJ20gbWFkZSBvZiB3YXgsIExhcnJ5IjoiV2hhdCBhcmUgeW91IG1hZGUgb2Y/IiwiX2RhdGFkb2ciOnsieC1kYXRhZG9nLXRyYWNlLWlkIjoiNjY3MzA5NTE0MjIxMDM1NTM4IiwieC1kYXRhZG9nLXBhcmVudC1pZCI6IjEzNTA3MzUwMzU0OTc4MTE4MjgiLCJ4LWRhdGFkb2ctc2FtcGxlZCI6IjEiLCJ4LWRhdGFkb2ctc2FtcGxpbmctcHJpb3JpdHkiOiIxIn19",
data: "eyJJJ20gbWFkZSBvZiB3YXgsIExhcnJ5IjoiV2hhdCBhcmUgeW91IG1hZGUgb2Y/IiwiX2RhdGFkb2ciOnsieC1kYXRhZG9nLXRyYWNlLWlkIjoiNjY3MzA5NTE0MjIxMDM1NTM4IiwieC1kYXRhZG9nLXBhcmVudC1pZCI6IjEzNTA3MzUwMzU0OTc4MTE4MjgiLCJ4LWRhdGFkb2ctc2FtcGxlZCI6IjEiLCJ4LWRhdGFkb2ctc2FtcGxpbmctcHJpb3JpdHkiOiIxIiwiZGQtcGF0aHdheS1jdHgtYmFzZTY0Ijoic29tZS1iYXNlNjQtZW5jb2RlZC1jb250ZXh0In19Cg==",
approximateArrivalTimestamp: 1642518727.248,
},
eventSource: "aws:kinesis",
Expand All @@ -67,25 +80,50 @@ describe("KinesisEventTraceExtractor", () => {
"x-datadog-sampled": "1",
"x-datadog-sampling-priority": "1",
"x-datadog-trace-id": "667309514221035538",
"dd-pathway-ctx-base64": "some-base64-encoded-context",
});

expect(traceContext?.toTraceId()).toBe("667309514221035538");
expect(traceContext?.toSpanId()).toBe("1350735035497811828");
expect(traceContext?.sampleMode()).toBe("1");
expect(traceContext?.source).toBe("event");

expect(mockDataStreamsCheckpointer.setConsumeCheckpoint).toHaveBeenCalledWith(
"kinesis",
"arn:aws:kinesis:EXAMPLE",
{
"x-datadog-parent-id": "1350735035497811828",
"x-datadog-sampled": "1",
"x-datadog-sampling-priority": "1",
"x-datadog-trace-id": "667309514221035538",
"dd-pathway-ctx-base64": "some-base64-encoded-context",
},
false,
);
});

it.each([
["Records", {}],
["Records first entry", { Records: [] }],
["valid data in kinesis", { Records: [{ kinesis: { data: "{" } }] }], // JSON.parse should fail
["_datadog in data", { Records: [{ kinesis: { data: "e30=" } }] }],
])("returns null and skips extracting when payload is missing '%s'", (_, payload) => {
["Records", {}, 0],
["Records first entry", { Records: [] }, 0],
["valid data in kinesis", { Records: [{ kinesis: { data: "{" }, eventSourceARN: "arn:aws:kinesis:test" }] }, 1], // JSON.parse should fail
["_datadog in data", { Records: [{ kinesis: { data: "e30=" }, eventSourceARN: "arn:aws:kinesis:test" }] }, 1],
])("returns null and skips extracting when payload is missing '%s'", (_, payload, dsmCalls) => {
const tracerWrapper = new TracerWrapper();
const extractor = new KinesisEventTraceExtractor(tracerWrapper);

const traceContext = extractor.extract(payload as any);
expect(traceContext).toBeNull();

expect(mockDataStreamsCheckpointer.setConsumeCheckpoint).toHaveBeenCalledTimes(dsmCalls);

if (dsmCalls > 0) {
expect(mockDataStreamsCheckpointer.setConsumeCheckpoint).toHaveBeenCalledWith(
"kinesis",
"arn:aws:kinesis:test",
null,
false,
);
}
});

it("returns null when extracted span context by tracer is null", () => {
Expand Down
6 changes: 6 additions & 0 deletions src/trace/context/extractors/kinesis.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,15 +8,19 @@ export class KinesisEventTraceExtractor implements EventTraceExtractor {
constructor(private tracerWrapper: TracerWrapper) {}

extract(event: KinesisStreamEvent): SpanContextWrapper | null {
let sourceARN = "";
const kinesisData = event?.Records?.[0]?.kinesis.data;
if (kinesisData === undefined) return null;

sourceARN = event.Records[0].eventSourceARN;
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what happens if event is null or Records is null or it has no records? i would use optional chaining like above.

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We are using records[0] (so the first message). Can a lambda be called with many messages? What do we want to do in that case? Set many checkpoints or just one for the first checkpoint ? What do we do in the Python lambda?


try {
const decodedData = Buffer.from(kinesisData, "base64").toString("ascii");
const parsedBody = JSON.parse(decodedData);
const headers = parsedBody?._datadog;
if (headers) {
const traceContext = this.tracerWrapper.extract(headers);
this.tracerWrapper.setConsumeCheckpoint(headers, "kinesis", sourceARN);
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

if we fail to capture the sourceARN, we probably want a debug log & not set a checkpoint

if (traceContext === null) return null;

logDebug(`Extracted trace context from Kinesis event`, { traceContext, headers });
Expand All @@ -28,6 +32,8 @@ export class KinesisEventTraceExtractor implements EventTraceExtractor {
}
}

// Still want to set a DSM checkpoint even if DSM context not propagated
this.tracerWrapper.setConsumeCheckpoint(null, "kinesis", sourceARN);
return null;
}
}
71 changes: 61 additions & 10 deletions src/trace/context/extractors/sns-sqs.spec.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,15 @@ import { TracerWrapper } from "../../tracer-wrapper";
import { SNSSQSEventTraceExtractor } from "./sns-sqs";

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
Expand All @@ -12,6 +21,7 @@ jest.mock("dd-trace", () => {
...ddTrace,
_tracer: { _service: {} },
extract: (_carrier: any, _headers: any) => mockSpanContext,
dataStreamsCheckpointer: mockDataStreamsCheckpointer,
};
});
const spyTracerWrapper = jest.spyOn(TracerWrapper.prototype, "extract");
Expand All @@ -21,10 +31,13 @@ describe("SNSSQSEventTraceExtractor", () => {
beforeEach(() => {
mockSpanContext = null;
spyTracerWrapper.mockClear();
mockDataStreamsCheckpointer.setConsumeCheckpoint.mockClear();
process.env["DD_DATA_STREAMS_ENABLED"] = "true";
});

afterEach(() => {
jest.resetModules();
delete process.env["DD_DATA_STREAMS_ENABLED"];
});

it("extracts trace context with valid payload with String Value", () => {
Expand All @@ -43,7 +56,7 @@ describe("SNSSQSEventTraceExtractor", () => {
messageId: "64812b68-4d9b-4dca-b3fb-9b18f255ee51",
receiptHandle:
"AQEBER6aRkfG8092GvkL7FRwCwbQ7LLDW9Tlk/CembqHe+suS2kfFxXiukomvaIN61QoyQMoRgWuV52SDkiQno2u+5hP64BDbmw+e/KR9ayvIfHJ3M6RfyQLaWNWm3hDFBCKTnBMVIxtdx0N9epZZewyokjKcrNYtmCghFgTCvZzsQkowi5rnoHAVHJ3je1c3bDnQ1KLrZFgajDnootYXDwEPuMq5FIxrf4EzTe0S7S+rnRm+GaQfeBLBVAY6dASL9usV3/AFRqDtaI7GKI+0F2NCgLlqj49VlPRz4ldhkGknYlKTZTluAqALWLJS62/J1GQo53Cs3nneJcmu5ajB2zzmhhRXoXINEkLhCD5ujZfcsw9H4xqW69Or4ECvlqx14bUU2rtMIW0QM2p7pEeXnyocymQv6m1te113eYWTVmaJ4I=",
body: '{\n "Type" : "Notification",\n "MessageId" : "0a0ab23e-4861-5447-82b7-e8094ff3e332",\n "TopicArn" : "arn:aws:sns:eu-west-1:601427279990:js-library-test-dev-demoTopic-15WGUVRCBMPAA",\n "Message" : "{\\"hello\\":\\"harv\\",\\"nice of you to join us\\":\\"david\\",\\"anotherThing\\":{\\"foo\\":\\"bar\\",\\"blah\\":null,\\"harv\\":123},\\"vals\\":[{\\"thingOne\\":1},{\\"thingTwo\\":2}],\\"ajTimestamp\\":1639777617957}",\n "Timestamp" : "2021-12-17T21:46:58.040Z",\n "SignatureVersion" : "1",\n "Signature" : "FR35/7E8C3LHEVk/rC4XxXlXwV/5mNkFNPgDhHSnJ2I6hIoSrTROAm7h5xm1PuBkAeFDvq0zofw91ouk9zZyvhdrMLFIIgrjEyNayRmEffmoEAkzLFUsgtQX7MmTl644r4NuWiM0Oiz7jueRvIcKXcZr7Nc6GJcWV1ymec8oOmuHNMisnPMxI07LIQVYSyAfv6P9r2jEWMVIukRoCzwTnRk4bUUYhPSGHI7OC3AsxxXBbv8snqTrLM/4z2rXCf6jHCKNxWeLlm9/45PphCkEyx5BWS4/71KaoMWUWy8+6CCsy+uF3XTCVmvSEYLyEwTSzOY+vCUjazrRW93498i70g==",\n "SigningCertUrl" : "https://sns.eu-west-1.amazonaws.com/SimpleNotificationService-************************33ab7e69.pem",\n "UnsubscribeUrl" : "https://sns.eu-west-1.amazonaws.com/?Action=Unsubscribe&SubscriptionArn=arn:aws:sns:eu-west-1:601427279990:js-library-test-dev-demoTopic-15WGUVRCBMPAA:1290f550-9a8a-4e8f-a900-8f5f96dcddda",\n "MessageAttributes" : {\n "_datadog" : {"Type":"String","Value":"{\\"x-datadog-trace-id\\":\\"2776434475358637757\\",\\"x-datadog-parent-id\\":\\"4493917105238181843\\",\\"x-datadog-sampled\\":\\"1\\",\\"x-datadog-sampling-priority\\":\\"1\\"}"}\n }\n}',
body: '{\n "Type" : "Notification",\n "MessageId" : "0a0ab23e-4861-5447-82b7-e8094ff3e332",\n "TopicArn" : "arn:aws:sns:eu-west-1:601427279990:js-library-test-dev-demoTopic-15WGUVRCBMPAA",\n "Message" : "{\\"hello\\":\\"harv\\",\\"nice of you to join us\\":\\"david\\",\\"anotherThing\\":{\\"foo\\":\\"bar\\",\\"blah\\":null,\\"harv\\":123},\\"vals\\":[{\\"thingOne\\":1},{\\"thingTwo\\":2}],\\"ajTimestamp\\":1639777617957}",\n "Timestamp" : "2021-12-17T21:46:58.040Z",\n "SignatureVersion" : "1",\n "Signature" : "FR35/7E8C3LHEVk/rC4XxXlXwV/5mNkFNPgDhHSnJ2I6hIoSrTROAm7h5xm1PuBkAeFDvq0zofw91ouk9zZyvhdrMLFIIgrjEyNayRmEffmoEAkzLFUsgtQX7MmTl644r4NuWiM0Oiz7jueRvIcKXcZr7Nc6GJcWV1ymec8oOmuHNMisnPMxI07LIQVYSyAfv6P9r2jEWMVIukRoCzwTnRk4bUUYhPSGHI7OC3AsxxXBbv8snqTrLM/4z2rXCf6jHCKNxWeLlm9/45PphCkEyx5BWS4/71KaoMWUWy8+6CCsy+uF3XTCVmvSEYLyEwTSzOY+vCUjazrRW93498i70g==",\n "SigningCertUrl" : "https://sns.eu-west-1.amazonaws.com/SimpleNotificationService-************************33ab7e69.pem",\n "UnsubscribeUrl" : "https://sns.eu-west-1.amazonaws.com/?Action=Unsubscribe&SubscriptionArn=arn:aws:sns:eu-west-1:601427279990:js-library-test-dev-demoTopic-15WGUVRCBMPAA:1290f550-9a8a-4e8f-a900-8f5f96dcddda",\n "MessageAttributes" : {\n "_datadog" : {"Type":"String","Value":"{\\"x-datadog-trace-id\\":\\"2776434475358637757\\",\\"x-datadog-parent-id\\":\\"4493917105238181843\\",\\"x-datadog-sampled\\":\\"1\\",\\"x-datadog-sampling-priority\\":\\"1\\",\\"dd-pathway-ctx-base64\\":\\"some-base64-encoded-context\\"}"}\n }\n}',
attributes: {
ApproximateReceiveCount: "1",
SentTimestamp: "1639777618130",
Expand All @@ -69,12 +82,26 @@ describe("SNSSQSEventTraceExtractor", () => {
"x-datadog-sampled": "1",
"x-datadog-sampling-priority": "1",
"x-datadog-trace-id": "2776434475358637757",
"dd-pathway-ctx-base64": "some-base64-encoded-context",
});

expect(traceContext?.toTraceId()).toBe("2776434475358637757");
expect(traceContext?.toSpanId()).toBe("4493917105238181843");
expect(traceContext?.sampleMode()).toBe("1");
expect(traceContext?.source).toBe("event");

expect(mockDataStreamsCheckpointer.setConsumeCheckpoint).toHaveBeenCalledWith(
"sqs",
"arn:aws:sqs:eu-west-1:601427279990:aj-js-library-test-dev-demo-queue",
{
"x-datadog-parent-id": "4493917105238181843",
"x-datadog-sampled": "1",
"x-datadog-sampling-priority": "1",
"x-datadog-trace-id": "2776434475358637757",
"dd-pathway-ctx-base64": "some-base64-encoded-context",
},
false,
);
});

it("extracts trace context with valid payload with Binary Value", () => {
Expand All @@ -93,7 +120,7 @@ describe("SNSSQSEventTraceExtractor", () => {
messageId: "64812b68-4d9b-4dca-b3fb-9b18f255ee51",
receiptHandle:
"AQEBER6aRkfG8092GvkL7FRwCwbQ7LLDW9Tlk/CembqHe+suS2kfFxXiukomvaIN61QoyQMoRgWuV52SDkiQno2u+5hP64BDbmw+e/KR9ayvIfHJ3M6RfyQLaWNWm3hDFBCKTnBMVIxtdx0N9epZZewyokjKcrNYtmCghFgTCvZzsQkowi5rnoHAVHJ3je1c3bDnQ1KLrZFgajDnootYXDwEPuMq5FIxrf4EzTe0S7S+rnRm+GaQfeBLBVAY6dASL9usV3/AFRqDtaI7GKI+0F2NCgLlqj49VlPRz4ldhkGknYlKTZTluAqALWLJS62/J1GQo53Cs3nneJcmu5ajB2zzmhhRXoXINEkLhCD5ujZfcsw9H4xqW69Or4ECvlqx14bUU2rtMIW0QM2p7pEeXnyocymQv6m1te113eYWTVmaJ4I=",
body: '{\n "Type" : "Notification",\n "MessageId" : "0a0ab23e-4861-5447-82b7-e8094ff3e332",\n "TopicArn" : "arn:aws:sns:eu-west-1:601427279990:js-library-test-dev-demoTopic-15WGUVRCBMPAA",\n "Message" : "{\\"hello\\":\\"harv\\",\\"nice of you to join us\\":\\"david\\",\\"anotherThing\\":{\\"foo\\":\\"bar\\",\\"blah\\":null,\\"harv\\":123},\\"vals\\":[{\\"thingOne\\":1},{\\"thingTwo\\":2}],\\"ajTimestamp\\":1639777617957}",\n "Timestamp" : "2021-12-17T21:46:58.040Z",\n "SignatureVersion" : "1",\n "Signature" : "FR35/7E8C3LHEVk/rC4XxXlXwV/5mNkFNPgDhHSnJ2I6hIoSrTROAm7h5xm1PuBkAeFDvq0zofw91ouk9zZyvhdrMLFIIgrjEyNayRmEffmoEAkzLFUsgtQX7MmTl644r4NuWiM0Oiz7jueRvIcKXcZr7Nc6GJcWV1ymec8oOmuHNMisnPMxI07LIQVYSyAfv6P9r2jEWMVIukRoCzwTnRk4bUUYhPSGHI7OC3AsxxXBbv8snqTrLM/4z2rXCf6jHCKNxWeLlm9/45PphCkEyx5BWS4/71KaoMWUWy8+6CCsy+uF3XTCVmvSEYLyEwTSzOY+vCUjazrRW93498i70g==",\n "SigningCertUrl" : "https://sns.eu-west-1.amazonaws.com/SimpleNotificationService-7ff5318490ec183fbaddaa2a969abfda.pem",\n "UnsubscribeUrl" : "https://sns.eu-west-1.amazonaws.com/?Action=Unsubscribe&SubscriptionArn=arn:aws:sns:eu-west-1:601427279990:js-library-test-dev-demoTopic-15WGUVRCBMPAA:1290f550-9a8a-4e8f-a900-8f5f96dcddda",\n "MessageAttributes" : {\n "_datadog" : {"Type":"Binary","Value":"eyJ4LWRhdGFkb2ctdHJhY2UtaWQiOiI3MTAyMjkxNjI4NDQzMTM0OTE5IiwieC1kYXRhZG9nLXBhcmVudC1pZCI6IjQyNDc1NTAxMDE2NDg2MTg2MTgiLCJ4LWRhdGFkb2ctc2FtcGxpbmctcHJpb3JpdHkiOiIxIn0="}\n }\n}',
body: '{\n "Type" : "Notification",\n "MessageId" : "0a0ab23e-4861-5447-82b7-e8094ff3e332",\n "TopicArn" : "arn:aws:sns:eu-west-1:601427279990:js-library-test-dev-demoTopic-15WGUVRCBMPAA",\n "Message" : "{\\"hello\\":\\"harv\\",\\"nice of you to join us\\":\\"david\\",\\"anotherThing\\":{\\"foo\\":\\"bar\\",\\"blah\\":null,\\"harv\\":123},\\"vals\\":[{\\"thingOne\\":1},{\\"thingTwo\\":2}],\\"ajTimestamp\\":1639777617957}",\n "Timestamp" : "2021-12-17T21:46:58.040Z",\n "SignatureVersion" : "1",\n "Signature" : "FR35/7E8C3LHEVk/rC4XxXlXwV/5mNkFNPgDhHSnJ2I6hIoSrTROAm7h5xm1PuBkAeFDvq0zofw91ouk9zZyvhdrMLFIIgrjEyNayRmEffmoEAkzLFUsgtQX7MmTl644r4NuWiM0Oiz7jueRvIcKXcZr7Nc6GJcWV1ymec8oOmuHNMisnPMxI07LIQVYSyAfv6P9r2jEWMVIukRoCzwTnRk4bUUYhPSGHI7OC3AsxxXBbv8snqTrLM/4z2rXCf6jHCKNxWeLlm9/45PphCkEyx5BWS4/71KaoMWUWy8+6CCsy+uF3XTCVmvSEYLyEwTSzOY+vCUjazrRW93498i70g==",\n "SigningCertUrl" : "https://sns.eu-west-1.amazonaws.com/SimpleNotificationService-7ff5318490ec183fbaddaa2a969abfda.pem",\n "UnsubscribeUrl" : "https://sns.eu-west-1.amazonaws.com/?Action=Unsubscribe&SubscriptionArn=arn:aws:sns:eu-west-1:601427279990:js-library-test-dev-demoTopic-15WGUVRCBMPAA:1290f550-9a8a-4e8f-a900-8f5f96dcddda",\n "MessageAttributes" : {\n "_datadog" : {"Type":"Binary","Value":"eyJ4LWRhdGFkb2ctdHJhY2UtaWQiOiI3MTAyMjkxNjI4NDQzMTM0OTE5IiwieC1kYXRhZG9nLXBhcmVudC1pZCI6IjQyNDc1NTAxMDE2NDg2MTg2MTgiLCJ4LWRhdGFkb2ctc2FtcGxpbmctcHJpb3JpdHkiOiIxIiwiZGQtcGF0aHdheS1jdHgtYmFzZTY0Ijoic29tZS1iYXNlNjQtZW5jb2RlZC1jb250ZXh0In0="}\n }\n}',
attributes: {
ApproximateReceiveCount: "1",
SentTimestamp: "1639777618130",
Expand All @@ -118,28 +145,52 @@ describe("SNSSQSEventTraceExtractor", () => {
"x-datadog-parent-id": "4247550101648618618",
"x-datadog-sampling-priority": "1",
"x-datadog-trace-id": "7102291628443134919",
"dd-pathway-ctx-base64": "some-base64-encoded-context",
});

expect(traceContext?.toTraceId()).toBe("7102291628443134919");
expect(traceContext?.toSpanId()).toBe("4247550101648618618");
expect(traceContext?.sampleMode()).toBe("1");
expect(traceContext?.source).toBe("event");

expect(mockDataStreamsCheckpointer.setConsumeCheckpoint).toHaveBeenCalledWith(
"sqs",
"arn:aws:sqs:eu-west-1:601427279990:aj-js-library-test-dev-demo-queue",
{
"x-datadog-parent-id": "4247550101648618618",
"x-datadog-sampling-priority": "1",
"x-datadog-trace-id": "7102291628443134919",
"dd-pathway-ctx-base64": "some-base64-encoded-context",
},
false,
);
});

it.each([
["Records", {}],
["Records first entry", { Records: [] }],
["Records first entry body", { Records: [{}] }],
["valid data in body", { Records: [{ body: "{" }] }], // JSON.parse should fail
["MessageAttributes in body", { Records: [{ body: "{}" }] }],
["_datadog in MessageAttributes", { Records: [{ body: '{"MessageAttributes":{"text":"Hello, world!"}}' }] }],
["Value in _datadog", { Records: [{ body: '{"MessageAttributes":{"_datadog":{}}}' }] }],
])("returns null and skips extracting when payload is missing '%s'", (_, payload) => {
["Records", {}, 0],
["Records first entry", { Records: [] }, 0],
["Records first entry body", { Records: [{}] }, 0],
["valid data in body", { Records: [{ body: "{", eventSourceARN: "arn:aws:sqs:us-east-1:test" }] }, 1], // JSON.parse should fail
["MessageAttributes in body", { Records: [{ body: "{}", eventSourceARN: "arn:aws:sqs:us-east-1:test" }] }, 1],
["_datadog in MessageAttributes", { Records: [{ body: '{"MessageAttributes":{"text":"Hello, world!"}}', eventSourceARN: "arn:aws:sqs:us-east-1:test" }] }, 1],
["Value in _datadog", { Records: [{ body: '{"MessageAttributes":{"_datadog":{}}}', eventSourceARN: "arn:aws:sqs:us-east-1:test" }] }, 1],
])("returns null and skips extracting when payload is missing '%s'", (_, payload, dsmCalls) => {
const tracerWrapper = new TracerWrapper();
const extractor = new SNSSQSEventTraceExtractor(tracerWrapper);

const traceContext = extractor.extract(payload as any);
expect(traceContext).toBeNull();

expect(mockDataStreamsCheckpointer.setConsumeCheckpoint).toHaveBeenCalledTimes(dsmCalls);

if (dsmCalls > 0) {
expect(mockDataStreamsCheckpointer.setConsumeCheckpoint).toHaveBeenCalledWith(
"sqs",
"arn:aws:sqs:us-east-1:test",
null,
false,
);
}
});

it("returns null when extracted span context by tracer is null", () => {
Expand Down
7 changes: 6 additions & 1 deletion src/trace/context/extractors/sns-sqs.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,12 @@ export class SNSSQSEventTraceExtractor implements EventTraceExtractor {
constructor(private tracerWrapper: TracerWrapper) {}

extract(event: SQSEvent): SpanContextWrapper | null {
let sourceARN = "";

try {
// First try to extract trace context from message attributes
if (event?.Records?.[0]?.body) {
sourceARN = event.Records[0].eventSourceARN;
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same comments as for Kinesis

const parsedBody = JSON.parse(event?.Records?.[0]?.body) as SNSMessage;
const messageAttribute = parsedBody?.MessageAttributes?._datadog;
if (messageAttribute?.Value) {
Expand All @@ -24,6 +27,7 @@ export class SNSSQSEventTraceExtractor implements EventTraceExtractor {
}

const traceContext = this.tracerWrapper.extract(headers);
this.tracerWrapper.setConsumeCheckpoint(headers, "sqs", sourceARN);
if (traceContext) {
logDebug("Extracted trace context from SNS-SQS event");
return traceContext;
Expand All @@ -48,7 +52,8 @@ export class SNSSQSEventTraceExtractor implements EventTraceExtractor {
logDebug("Unable to extract trace context from SNS-SQS event", error);
}
}

// Still want to set a DSM checkpoint even if DSM context not propagated
this.tracerWrapper.setConsumeCheckpoint(null, "sqs", sourceARN);
return null;
}
}
Loading
Loading