Inject Kinesis DSM context into all batch records#8640
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 30ad5edac9
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| foreach (var requestRecord in request.Records) | ||
| { | ||
| InjectTraceIntoData(tracer, record, scope, streamName); | ||
| if (requestRecord.DuckCast<IContainsData>() is { } record) | ||
| { | ||
| InjectTraceIntoData(tracer, record, scope, streamName); |
There was a problem hiding this comment.
For PutRecords batches whose original payload is close to the service request-size limit, this loop now adds a _datadog object to every JSON record without tracking the aggregate request size, so a call that previously only grew the first record can now exceed Kinesis' whole-request limit and be rejected. The per-record MaxKinesisDataSize check inside InjectTraceIntoData doesn't protect the batch total; please account for the total batch size, including partition keys/added context, before mutating every record.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 30ad5edac9
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| foreach (var requestRecord in request.Records) | ||
| { | ||
| InjectTraceIntoData(tracer, record, scope, streamName); | ||
| if (requestRecord.DuckCast<IContainsData>() is { } record) | ||
| { | ||
| InjectTraceIntoData(tracer, record, scope, streamName); |
There was a problem hiding this comment.
For PutRecords batches whose original payload is close to the service request-size limit, this loop now adds a _datadog object to every JSON record without tracking the aggregate request size, so a call that previously only grew the first record can now exceed Kinesis' whole-request limit and be rejected. The per-record MaxKinesisDataSize check inside InjectTraceIntoData doesn't protect the batch total; please account for the total batch size, including partition keys/added context, before mutating every record.
Useful? React with 👍 / 👎.
1 participant
Summary of changes
Updates the AWS Kinesis PutRecords instrumentation so trace propagation and DSM injection are attempted for every record in the batch, instead of only Records[0].
Reason for change
Batch propagation for Kinesis was inconsistent with the other AWS messaging integrations. SQS and SNS already inject Datadog context into every message in a batch, while Kinesis only injected into the first record.
Bringing Kinesis in line with the other integrations makes batch behavior more predictable and avoids partial propagation when a single PutRecords request contains multiple independently consumed records.
Implementation details
ContextPropagation.InjectTraceIntoRecords() now iterates through request.Records and calls InjectTraceIntoData() for each duck-typed record entry.
The underlying per-record behavior is unchanged:
Injection still only happens when the payload can be parsed as JSON.
DSM injection still depends on the existing payload/header-size checks.
Injection is still skipped if the final payload would exceed the Kinesis 1 MB limit.
PutRecord behavior is unchanged.
One important tradeoff is that this introduces additional per-batch work for Kinesis. Previously we only parsed and attempted injection for the first record; now we do that work for every record in the batch. That means CPU/allocation overhead will scale with batch size, especially for larger JSON payloads. This is a potential performance cost, but it also aligns Kinesis with the existing SQS/SNS model, where Datadog context is added to every message in a batch.
If there are concerns around performance, we could possibly hide this behind a feature flag.
dd-trace-jsdoes this.Test coverage
Updated the managed Kinesis propagation tests to verify that all records in a PutRecords batch receive _datadog context, not just the first record.
Focused test run:
Datadog.Trace.ClrProfiler.Managed.Tests.AutoInstrumentation.AWS.Kinesis.ContextPropagationTests
11 passed, 0 failed on net10.0
Other details