feat: lambda support for DSM #672

michael-zhao459 · 2025-07-28T21:44:20Z

What does this PR do?

This PR adds lambda support for Data Streams Monitoring (DSM), sets a checkpoint after extracting context from trace propagation headers, DSM context gets packaged in with the trace prop headers. Queues supported by DSM include SQS, SNS, SNS->SQS, and Kinesis

Motivation

Consume calls of lambdas are set as triggers to queue and do not have any explicit calls the tracers can hook into to set a checkpoint

Testing Guidelines

Tried all configurations with sandbox AWS account. Wrote unit tests ensuring all cases were covered.

Additional Notes

Types of Changes

Bug fix
New feature
Breaking change
Misc (docs, refactoring, dependency upgrade, etc.)

Check all that apply

This PR's description is comprehensive
This PR contains breaking changes that are documented in the description
This PR introduces new APIs or parameters that are documented and unlikely to change in the foreseeable future
This PR impacts documentation, and it has been updated (or a ticket has been logged)
This PR's changes are covered by the automated tests
This PR collects user input/sensitive content into Datadog
This PR passes the integration tests (ask a Datadog member to run the tests)

src/trace/tracer-wrapper.ts

src/trace/context/extractors/kinesis.ts

src/trace/context/extractors/sns-sqs.ts

src/trace/context/extractors/sns.ts

src/trace/tracer-wrapper.ts

src/trace/context/extractors/kinesis.ts

src/trace/tracer-wrapper.ts

src/trace/context/extractors/sns-sqs.ts

…t of the loop of records earlier when data streams is not enabled

purple4reina · 2025-08-22T16:39:28Z

src/index.ts

+ */
+export function getDataStreamsEnabled(): boolean {
+  return getConfig().dataStreamsEnabled;
+}


This is not what we want. Calling getConfig will read each and every environment variable each and every time it is called.

Instead, you will want to use the global config that is created just once at start up.

Pass this.config to the constructor of the trace extractor

datadog-lambda-js/src/trace/context/extractor.ts

Line 89 in 1619bce

if (EventValidator.isSNSEvent(event)) return new SNSEventTraceExtractor(this.tracerWrapper);

then you can use it during extraction without having to read environment variables each time.

thanks for clarifying. i am not super-familiar with this repo. made this change as well

purple4reina · 2025-08-22T16:43:12Z

src/trace/context/extractors/kinesis.ts

+        // If we already have a context and dsm is not enabled, we can break out of the loop early
+        if (!getDataStreamsEnabled() && context) {
+          break;
+        }


This changes the behavior of the extractor and in the case where there is no trace propagation context, will loop through each and every record attempting to look for one. That will potentially add significant overhead to the function execution time.

Instead, after the context has been extracted (or not), we should break if dsm is not enabled, whether the context exists or not.

makes sense. i changed the extractors to not change anything about the context fetching from main, but add the loop on records only when dsm is enabled

purple4reina

Two requested changes:

Read from global config object
Always break from loop when dsm not enabled

…oints with the context stuff completely

purple4reina · 2025-08-22T19:11:35Z

src/trace/tracer-wrapper.ts

+
+    if (!this.config.dataStreamsEnabled) {
+      return;
+    }


Can we put this before the arn check, that way when dsm is disabled, users will never see dsm related messages in their logs.

purple4reina · 2025-08-22T19:15:56Z

src/index.ts

    };
-    const tracerWrapper = new TracerWrapper();
+    const config = getConfig();
+    const tracerWrapper = new TracerWrapper(config);


I don't understand why is this change necessary? This method is only called when there is an error attempting to initialize the handler.

purple4reina · 2025-08-22T19:17:10Z

src/trace/tracer-wrapper.ts

+      coldStartTraceSkipLib: "",
+      addSpanPointers: true,
+      dataStreamsEnabled: false,
+    };


I'd prefer not to have all of this testing data in the code. Is there not a way we can move it to the tests? We don't want to have to maintain two separate locations for managing default configuration.

In fact, we don't need to add the config to this class in the first place. We should just rely on the config on the extractor alone.

👍 true. I think I was being too deferential to the change Michael made where there was a check in the tracer wrapper because set consume checkpoint was always run. But with the optimizations where we are only setting checkpoints when dsm is enabled, so there is no need for the second check

purple4reina · 2025-08-22T19:19:43Z

src/trace/context/extractors/sqs.ts

+            }
+          }
+
+          headers = headers ? JSON.parse(headers) : null;


When dsm is enabled, we'll now json parse the headers twice. I feel like that can and should be avoided. This method is in the hotpath of people's invocations and needs to be as performant as possible.

… the extractors

purple4reina · 2025-08-25T16:42:40Z

src/trace/tracer-wrapper.ts

+    }
+
+    try {
+      this.tracer.dataStreamsCheckpointer.setConsumeCheckpoint(eventType, arn, contextJson, false);


Question: is there a particular version of the tracer that this requires?

meaning the 4th argument (whether its a manually instrumented or not)? that is on 5.62.0, https://github.com/DataDog/dd-trace-js/releases/tag/v5.62.0

In that case, you will need to bump the dependency version of ddtrace.

datadog-lambda-js/package.json

Line 32 in 1619bce

"dd-trace": "^5.58.0",

purple4reina · 2025-08-25T16:55:16Z

src/trace/context/extractors/sqs.ts

-        // Then try to get from binary value. This happens when SNS->SQS, but SNS has raw message delivery enabled.
-        // In this case, SNS maps any messageAttributes to the SQS messageAttributes.
-        // We can at least get trace context from SQS, but we won't be able to create the SNS inferred span.


Can we keep this comment please?

…the case where DSM is instrumented

purple4reina · 2025-08-25T17:08:01Z

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

      expect(traceContext).toBeNull();
    });
+
+    it("extracts trace context from multiple records when DSM is disabled but does not call setConsumeCheckpoint", () => {


This seems to be wrong now, we only extract from the first record, yeah?

true. the test (trace context still works when no dsm) is still relevant but i updated the description to be accurate

michael-zhao459 added 11 commits July 24, 2025 19:51

set checkpoint for all supported dsm queue types

d0764da

removed not needed checkpoint set

5d4da9f

clean

4ea0470

clearer return type

c2cefc2

add tests

2626672

add the manual checkpoint parameter

ca9edb9

Correct approach to ensure public api is used

d0e71e6

remove unused variable

197abd0

match logging patterns of repo

c9d6c23

lint

75e976f

lint

a12d0ac

michael-zhao459 requested a review from piochelepiotr July 28, 2025 21:56

michael-zhao459 added 4 commits July 28, 2025 23:56

add False manual_checkpoint false

99f1c67

fix formatting

7b4c4aa

optimization

5059c5c

clean

7d3fdba

michael-zhao459 commented Jul 29, 2025

View reviewed changes

src/trace/tracer-wrapper.ts Show resolved Hide resolved

add prettier ignore

c49e96f

piochelepiotr reviewed Jul 29, 2025

View reviewed changes

michael-zhao459 added 2 commits July 29, 2025 13:55

add optional chaining

e422375

add log debug for no arn case

893d391

michael-zhao459 requested a review from piochelepiotr July 29, 2025 18:23

michael-zhao459 and others added 8 commits August 13, 2025 12:44

Merge remote-tracking branch 'origin/main' into michael.zhao/lambda-dsm

e44b6a7

merge commit

bc8cb73

add checkpoints for headers

70c5824

Update SQS to work with batch processing

1675ddb

Update Kinesis to work with batch processing

c2b457f

Update SNS-SQS to work with batch processing

5974b32

Prettier changes

c39db0a

Fix test for set sqs-sns

5d06798

purple4reina reviewed Aug 21, 2025

View reviewed changes

src/trace/context/extractors/kinesis.ts Outdated Show resolved Hide resolved

purple4reina reviewed Aug 21, 2025

View reviewed changes

src/trace/tracer-wrapper.ts Outdated Show resolved Hide resolved

purple4reina reviewed Aug 21, 2025

View reviewed changes

src/trace/context/extractors/sns-sqs.ts Outdated Show resolved Hide resolved

Move the env var check on data streams enabled to config and break ou…

2a1b75d

…t of the loop of records earlier when data streams is not enabled

ericfirth requested a review from purple4reina August 21, 2025 16:36

test because i think the integration test failed for flake reasons

21780b6

purple4reina reviewed Aug 22, 2025

View reviewed changes

purple4reina requested changes Aug 22, 2025

View reviewed changes

Move from getConfig() to passed in config for dataStreamsEnabled

632126b

ericfirth force-pushed the michael.zhao/lambda-dsm branch from c94ecbb to adce22d Compare August 22, 2025 18:12

Update the extraction functions to separate the setting of dsm checkp…

e00c8e1

…oints with the context stuff completely

ericfirth force-pushed the michael.zhao/lambda-dsm branch from adce22d to e00c8e1 Compare August 22, 2025 18:23

purple4reina reviewed Aug 22, 2025

View reviewed changes

Move the config of dataStreamsEnabled out of the trace-wrapper but in…

49da2bb

… the extractors

purple4reina reviewed Aug 25, 2025

View reviewed changes

Optimize such that you only parse the first messages headers once in …

bc3648b

…the case where DSM is instrumented

ericfirth force-pushed the michael.zhao/lambda-dsm branch from 72680f4 to bc3648b Compare August 25, 2025 17:02

purple4reina reviewed Aug 25, 2025

View reviewed changes

purple4reina approved these changes Aug 25, 2025

View reviewed changes

ericfirth added 2 commits August 26, 2025 09:05

Update the dd-trace version

598ed8b

Fix test description to be accurate for DSM is disabled

e861214

ericfirth force-pushed the michael.zhao/lambda-dsm branch from e46b50e to e861214 Compare August 26, 2025 13:06

empty

20e7eec

ericfirth merged commit 5d3da3d into main Aug 26, 2025
26 checks passed

ericfirth deleted the michael.zhao/lambda-dsm branch August 26, 2025 15:22

feat: lambda support for DSM #672

feat: lambda support for DSM #672

Uh oh!

Conversation

michael-zhao459 commented Jul 28, 2025 • edited by ericfirth Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Motivation

Testing Guidelines

Additional Notes

Types of Changes

Check all that apply

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

purple4reina left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

michael-zhao459 commented Jul 28, 2025 •

edited by ericfirth

Loading