feat(dsm): Add support for automatic DSM context extraction inside the extension#1265
feat(dsm): Add support for automatic DSM context extraction inside the extension#1265jeastham1993 wants to merge 26 commits into
Conversation
|
There was a problem hiding this comment.
Pull request overview
This PR extends Bottlecap’s universal instrumentation trace context propagation to automatically extract and emit Data Streams Monitoring (DSM) consume-side checkpoints within the extension for SQS, SNS, Kinesis, and EventBridge. It adds a DSM aggregation/serialization pipeline that drains into the existing proxy flush path and introduces configuration knobs to enable DSM consume extraction and provide an EventBridge exchange fallback.
Changes:
- Add a new
traces::data_streamsmodule implementing DSM context decode, pathway hashing, checkpoint computation, aggregation, and msgpack+gzip payload generation. - Wire an optional
DsmProcessorinto invocation start processing (universal instrumentation) and into the flushing pipeline so pipeline-stats are shipped via the proxy flusher. - Add trigger-specific DSM edge tags for SQS/SNS/Kinesis/EventBridge and introduce
DD_DSM_CONSUME_ENABLED/DD_DSM_EXCHANGE_NAMEconfig support.
Reviewed changes
Copilot reviewed 25 out of 26 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
| bottlecap/src/traces/mod.rs | Exposes the new data_streams module from traces. |
| bottlecap/src/traces/data_streams/mod.rs | DSM module entrypoint; re-exports key DSM types/functions. |
| bottlecap/src/traces/data_streams/context.rs | Implements inbound DSM pathway context decoding (base64 + zigzag varints). |
| bottlecap/src/traces/data_streams/pathway.rs | Implements dd-trace-js-compatible pathway hash computation. |
| bottlecap/src/traces/data_streams/propagation_hash.rs | Implements optional FNV-1 propagation hash. |
| bottlecap/src/traces/data_streams/checkpoint.rs | Computes consume-side checkpoints from extracted context + tags. |
| bottlecap/src/traces/data_streams/sketch.rs | Implements tracer-compatible DDSketch + protobuf bytes serialization. |
| bottlecap/src/traces/data_streams/aggregator.rs | Aggregates checkpoints into 10s buckets and serializes msgpack payloads. |
| bottlecap/src/traces/data_streams/processor.rs | Bridges checkpoint aggregation to proxy flush by gzipping/enqueueing proxy requests. |
| bottlecap/src/traces/data_streams/fixtures/sketch_golden.json | Golden vectors to validate DDSketch compatibility. |
| bottlecap/src/tags/lambda/tags.rs | Updates hardcoded extension version string used in tags/logging. |
| bottlecap/src/proxy/interceptor.rs | Enables universal instrumentation processing when experimental proxy env var is set. |
| bottlecap/src/lifecycle/invocation/triggers/mod.rs | Adds Trigger::get_dsm_edge_tags() default hook. |
| bottlecap/src/lifecycle/invocation/triggers/sqs_event.rs | Adds SQS DSM consume edge tag derivation. |
| bottlecap/src/lifecycle/invocation/triggers/sns_event.rs | Adds SNS DSM consume edge tag derivation. |
| bottlecap/src/lifecycle/invocation/triggers/kinesis_event.rs | Adds Kinesis DSM consume edge tag derivation. |
| bottlecap/src/lifecycle/invocation/triggers/event_bridge_event.rs | Adds EventBridge DSM consume tags + best-effort bus name extraction + tests. |
| bottlecap/src/lifecycle/invocation/processor.rs | Hooks DSM consume recording into universal instrumentation start; adds EventBridge exchange fallback. |
| bottlecap/src/lifecycle/invocation/processor_service.rs | Threads optional DsmProcessor into Processor initialization. |
| bottlecap/src/flushing/service.rs | Drains DSM pipeline-stats into proxy aggregator immediately before proxy flush. |
| bottlecap/src/config/yaml.rs | Adds YAML config fields for DSM consume enabling and exchange fallback. |
| bottlecap/src/config/mod.rs | Adds Config fields: dsm_consume_enabled, dsm_exchange_name. |
| bottlecap/src/config/env.rs | Adds env config parsing/merging for DD_DSM_CONSUME_ENABLED and DD_DSM_EXCHANGE_NAME. |
| bottlecap/src/bin/bottlecap/main.rs | Constructs shared proxy aggregator, conditionally instantiates DSM processor, wires it into services. |
| bottlecap/Cargo.toml | Promotes msgpack/gzip dependencies to main deps and adds serde_bytes. |
| bottlecap/Cargo.lock | Adds serde_bytes to resolved dependency set. |
![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: 579f6e6d9b
ℹ️ 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".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 947380af7b
ℹ️ 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".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ca5b7da286
ℹ️ 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".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9b923a55cf
ℹ️ 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".
| // functions whose tracer does not drive the invocation lifecycle. | ||
| let experimental_proxy_enabled = std::env::var("DD_EXPERIMENTAL_ENABLE_PROXY") | ||
| .is_ok_and(|v| v.eq_ignore_ascii_case("true")); | ||
| if aws_config.aws_lwa_proxy_lambda_runtime_api.is_some() || experimental_proxy_enabled { |
There was a problem hiding this comment.
When the runtime API proxy is running for the Datadog wrapper because AppSec is enabled, and DD_DSM_CONSUME_ENABLED=true but no tracer calls /lambda/start-invocation, this condition stays false unless LWA or DD_EXPERIMENTAL_ENABLE_PROXY is also set. In that no-tracer/AppSec setup the event body is available here but lwa::process_invocation_next is skipped, so process_on_universal_instrumentation_start never records any DSM consume checkpoints; include the DSM flag, or all proxy-active wrapper cases, in this gate.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c5dc4159e2
ℹ️ 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".
…pport # Conflicts: # bottlecap/src/config/env.rs # bottlecap/src/config/mod.rs # bottlecap/src/config/yaml.rs
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7d710d7418
ℹ️ 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".
| bottlecap::traces::data_streams::DsmProcessor::new( | ||
| service, | ||
| env, | ||
| env!("CARGO_PKG_VERSION").to_string(), |
There was a problem hiding this comment.
When DD_DSM_CONSUME_ENABLED is set, every extension-generated pipeline_stats payload reports TracerVersion from Cargo's package version, which is currently the crate's fixed 0.1.0 rather than the released Lambda extension version exposed via EXTENSION_VERSION. That makes DSM intake data from different extension releases indistinguishable for rollout/debugging of this new feature; pass the extension version constant instead of CARGO_PKG_VERSION.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c27c642406
ℹ️ 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".
| } | ||
|
|
||
| #[derive(Serialize)] | ||
| struct StatsPayloadSer { |
There was a problem hiding this comment.
When DD_VERSION or DD_TAGS are configured, extension-produced DSM points will not carry those unified service/custom tags because this payload shape has no Version or Tags fields, while the proxy flusher only adds _dd.origin and functionname headers. As a result data_streams.latency emitted by DD_DSM_CONSUME_ENABLED cannot be filtered or grouped by deployment version, team, or other global tags even though the rest of the extension telemetry can; pass config.version and config.tags into the DSM aggregator and serialize them with the payload.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f99e4b3c27
ℹ️ 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".
duncanista
left a comment
duncanista
left a comment
There was a problem hiding this comment.
a few things worth digging into before merge, mostly architecture and correctness plus one cold-start nit. nothing blocking that can't be sorted in-thread, inline below.
|
@duncanista think that's everything addressed 👍 |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f7f80283aa
ℹ️ 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".
| let dsm_consume_enabled = std::env::var("DD_DATA_STREAMS_ENABLED") | ||
| .is_ok_and(|v| v.eq_ignore_ascii_case("true")); |
There was a problem hiding this comment.
When Data Streams is enabled through the normal config path, e.g. datadog.yaml with data_streams_enabled: true, main creates the DSM processor from config.ext.dsm_consume_enabled, but this proxy gate re-reads only the DD_DATA_STREAMS_ENABLED environment variable. In an AppSec/datadog-wrapper setup with no tracer /lambda/start-invocation, the intercepted payload is available but this condition remains false, so the DSM hook never records consume checkpoints despite the feature being enabled; pass the resolved config flag into the interceptor instead of checking the environment here.
Useful? React with 👍 / 👎.
| let carrier = record | ||
| .get("headers") | ||
| .map_or_else(HashMap::new, headers_to_string_map); |
There was a problem hiding this comment.
For MSK/Kafka records that still carry the legacy dd-pathway-ctx header, the Lambda event exposes the header value as bytes, and those bytes are intentionally not UTF-8-safe because the first 8 bytes are an arbitrary pathway hash. This added DSM path converts the record headers through headers_to_string_map, which only keeps values that pass String::from_utf8, so legacy raw pathway contexts are dropped before extract_pathway_context can decode them with from_bytes, breaking continuity for those Kafka consumers; keep raw bytes for dd-pathway-ctx or encode that header losslessly.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 745c6f5274
ℹ️ 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".
| .is_ok_and(|v| v.eq_ignore_ascii_case("true")); | ||
| if aws_config.aws_lwa_proxy_lambda_runtime_api.is_some() | ||
| || experimental_proxy_enabled | ||
| || dsm_consume_enabled |
There was a problem hiding this comment.
When the Runtime API proxy is active for AppSec and DD_DATA_STREAMS_ENABLED=true in a runtime that still calls /lambda/start-invocation (Java/Go/.NET), this new branch runs lwa::process_invocation_next, which unconditionally queues reparenting with a freshly generated invocation span id, while the later tracer start replaces the invocation span id from x-datadog-span-id. The trace flusher then reparents child spans to the synthetic id that is never emitted, so traces with inbound context become orphaned; drive only the DSM extraction here or avoid adding LWA reparenting when a tracer start is expected.
Useful? React with 👍 / 👎.
duncanista
left a comment
duncanista
left a comment
There was a problem hiding this comment.
nice work, this is a clean feature and the compatibility-sensitive parts are well isolated. i checked the core against dd-trace-js and the fixtures, and verified the utf8 lossy round-trip matches node over a 50k fuzz, so the pathway hashes line up. none of the notes below are blockers. mostly one flag-parsing consistency thing, a couple of confirmations on the sns/eventbridge consume tags, and small nits.
| // conflict with the tracer-provided invocation span id. | ||
| let experimental_proxy_enabled = std::env::var("DD_EXPERIMENTAL_ENABLE_PROXY") | ||
| .is_ok_and(|v| v.eq_ignore_ascii_case("true")); | ||
| let dsm_consume_enabled = std::env::var("DD_DATA_STREAMS_ENABLED") |
There was a problem hiding this comment.
blocker. two things here. this reads the flag straight from env with a strict "true" check, but the config path uses the lenient bool parser (accepts 1, yes, TRUE). so DD_DATA_STREAMS_ENABLED=1 builds the processor and runs the flusher, but silently skips this proxy branch, so dsm goes missing on that path with no error. it also does an env::var read on every invocation on the hot path. can we thread config.ext.dsm_consume_enabled into the interceptor state and use that? that fixes the mismatch and the per-invocation read in one go.
There was a problem hiding this comment.
quick context on why this is env and not config: the interceptor state only carries AwsConfig, and Config (which holds ext.dsm_consume_enabled) is never threaded into interceptor::start, so the handler has to fall back to env::var. the neighbouring DD_EXPERIMENTAL_ENABLE_PROXY read is env-only by design, but dsm already has a parsed config field, so re-reading env here parses it with a stricter rule, which is the =1 mismatch. start_api_runtime_proxy already has config in scope, so threading dsm_consume_enabled (or Arc) through into the state and reading it from there would drop the per-invocation env::var and make the gating consistent in one change.
| }; | ||
|
|
||
| let payload_value = serde_json::from_slice(body).unwrap_or_else(|e| { | ||
| error!("PROXY | invocation_next_proxy | error parsing DSM payload as JSON: {e}"); |
There was a problem hiding this comment.
nit, not a blocker. non-json payloads are fairly common, and this logs at error for each one. since it falls back to {} and just skips, debug feels more appropriate here.
| Some(vec
Please include Jira ticket in title.
Overview
Update the trace context propagation support in universal instrumentation to include the automatic extraction of Data Streams Monitoring context for SQS, SNS, Kinesis and EventBridge.
Testing
Added unit tests. Manually tested the functionality for Java, Go and .NET. Java and Go work correctly. The .NET tracer needs updating to support the 2.x version of
Amazon.Lambda.RuntimeSupportso can't test that manually yet.