From 3714e34ce669de60c0ec50b6a2b2afb0d7910dda Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 26 Aug 2026 16:00:16 -0400 Subject: [PATCH 01/23] refactor(bottlecap): extract trace-agent startup into the library crate Moves start_trace_agent out of the Lambda binary into bottlecap::startup so both [[bin]] targets can share it, and splits it into: - build_trace_agent, which returns an unspawned TraceAgent plus a TraceAgentPipeline handle struct - start_trace_agent, a thin wrapper that spawns the agent bottlecap-test-mode needs the unspawned agent so it can attach its /flush RouterExtension before spawning; the Lambda binary keeps calling start_trace_agent and its call site is unchanged. Placed at the crate root rather than under traces/ because it wires trace, stats, proxy, lifecycle, tags, appsec, and flushing together. --- bottlecap/src/bin/bottlecap/main.rs | 147 +----------------- bottlecap/src/lib.rs | 1 + bottlecap/src/startup.rs | 226 ++++++++++++++++++++++++++++ bottlecap/src/traces/trace_agent.rs | 9 +- 4 files changed, 237 insertions(+), 146 deletions(-) create mode 100644 bottlecap/src/startup.rs diff --git a/bottlecap/src/bin/bottlecap/main.rs b/bottlecap/src/bin/bottlecap/main.rs index d3cf4dac2..9834610fe 100644 --- a/bottlecap/src/bin/bottlecap/main.rs +++ b/bottlecap/src/bin/bottlecap/main.rs @@ -61,21 +61,11 @@ use bottlecap::{ provider::Provider as TagProvider, }, traces::{ - http_client as trace_http_client, propagation::DatadogCompositePropagator, proxy_aggregator, - proxy_flusher::Flusher as ProxyFlusher, - span_dedup_service, - stats_aggregator::StatsAggregator, - stats_concentrator_service::{StatsConcentratorHandle, StatsConcentratorService}, - stats_flusher, + stats_concentrator_service::StatsConcentratorHandle, stats_generator::StatsGenerator, - stats_processor, trace_agent, trace_aggregator::SendDataBuilderInfo, - trace_aggregator_service::{ - AggregatorHandle as TraceAggregatorHandle, AggregatorService as TraceAggregatorService, - }, - trace_flusher, trace_processor::{self, SendingTraceProcessor}, }, }; @@ -95,7 +85,6 @@ use dogstatsd::{ flusher::{Flusher as MetricsFlusher, FlusherConfig as MetricsFlusherConfig}, metric::{EMPTY_TAGS, SortedTags}, }; -use libdd_trace_obfuscation::obfuscation_config; use reqwest::Client; use std::{collections::hash_map, env, path::Path, str::FromStr, sync::Arc}; use tokio::time::Instant; @@ -414,16 +403,16 @@ async fn extension_loop_active( } }; - let ( - trace_agent_channel, + let bottlecap::startup::TraceAgentPipeline { + trace_tx: trace_agent_channel, trace_flusher, trace_processor, stats_flusher, proxy_flusher, - trace_agent_shutdown_token, + shutdown_token: trace_agent_shutdown_token, stats_concentrator, trace_aggregator_handle, - ) = start_trace_agent( + } = bottlecap::startup::start_trace_agent( config, &api_key_factory, &tags_provider, @@ -1158,132 +1147,6 @@ fn start_logs_agent( ) } -#[allow(clippy::type_complexity)] -fn start_trace_agent( - config: &Arc, - api_key_factory: &Arc, - tags_provider: &Arc, - invocation_processor_handle: InvocationProcessorHandle, - appsec_processor: Option>>, - client: &Client, - proxy_aggregator: Arc>, -) -> ( - Sender, - Arc, - Arc, - Arc, - Arc, - tokio_util::sync::CancellationToken, - StatsConcentratorHandle, - TraceAggregatorHandle, -) { - // Build one shared hyper-based HTTP client for trace and stats flushing. - // This client type is required by libdd_trace_utils for SendData::send(). - let trace_http_client = trace_http_client::create_client( - config.proxy_https.as_ref(), - config.tls_cert_file.as_ref(), - config.skip_ssl_validation, - ) - .expect("Failed to create trace HTTP client"); - - // Stats - let (stats_concentrator_service, stats_concentrator_handle) = - StatsConcentratorService::new(Arc::clone(config)); - tokio::spawn(stats_concentrator_service.run()); - let stats_aggregator: Arc> = Arc::new(TokioMutex::new( - StatsAggregator::new_with_concentrator(stats_concentrator_handle.clone()), - )); - let stats_flusher = Arc::new(stats_flusher::StatsFlusher::new( - api_key_factory.clone(), - stats_aggregator.clone(), - Arc::clone(config), - trace_http_client.clone(), - libdd_trace_utils::config_utils::trace_stats_url(&config.site), - )); - - let stats_processor = Arc::new(stats_processor::ServerlessStatsProcessor {}); - - // Traces - let (trace_aggregator_service, trace_aggregator_handle) = TraceAggregatorService::default(); - tokio::spawn(trace_aggregator_service.run()); - - let trace_flusher = Arc::new(trace_flusher::TraceFlusher::new( - trace_aggregator_handle.clone(), - config.clone(), - api_key_factory.clone(), - trace_http_client, - )); - - let obfuscation_config = obfuscation_config::ObfuscationConfig { - tag_replace_rules: config.apm_replace_tags.clone(), - http: obfuscation_config::HttpConfig { - remove_paths_with_digits: config.apm_config_obfuscation_http_remove_paths_with_digits, - remove_query_string: config.apm_config_obfuscation_http_remove_query_string, - }, - ..Default::default() - }; - - // The Agent's error sampler knob has no effect here: the extension's - // sampler is a plain on/off switch, not a TPS budget. - if env::var("DD_APM_ERROR_TPS").is_ok_and(|v| !v.trim().is_empty()) { - warn!( - "DD_APM_ERROR_TPS is not supported by the Lambda extension; error trace rescue is an on/off switch controlled by DD_SERVERLESS_ERROR_SAMPLER_ENABLED" - ); - } - - let trace_processor = Arc::new(trace_processor::ServerlessTraceProcessor { - obfuscation_config: Arc::new(obfuscation_config), - error_sampler: trace_processor::new_error_sampler( - config.ext.serverless_error_sampler_enabled, - ), - }); - - let (span_dedup_service, span_dedup_handle) = span_dedup_service::DedupService::new(); - tokio::spawn(span_dedup_service.run()); - - // Proxy - let proxy_flusher = Arc::new(ProxyFlusher::new( - api_key_factory.clone(), - Arc::clone(&proxy_aggregator), - Arc::clone(tags_provider), - Arc::clone(config), - client.clone(), - )); - - let trace_agent = trace_agent::TraceAgent::new( - Arc::clone(config), - trace_aggregator_handle.clone(), - trace_processor.clone(), - stats_aggregator, - stats_processor, - proxy_aggregator, - invocation_processor_handle, - appsec_processor, - Arc::clone(tags_provider), - stats_concentrator_handle.clone(), - span_dedup_handle, - ); - let trace_agent_channel = trace_agent.get_sender_copy(); - let shutdown_token = trace_agent.shutdown_token(); - - tokio::spawn(async move { - if let Err(e) = trace_agent.start().await { - error!("Error starting trace agent: {e:?}"); - } - }); - - ( - trace_agent_channel, - trace_flusher, - trace_processor, - stats_flusher, - proxy_flusher, - shutdown_token, - stats_concentrator_handle, - trace_aggregator_handle, - ) -} - async fn start_dogstatsd( tags_provider: Arc, api_key_factory: Arc, diff --git a/bottlecap/src/lib.rs b/bottlecap/src/lib.rs index df94fd246..97784956c 100644 --- a/bottlecap/src/lib.rs +++ b/bottlecap/src/lib.rs @@ -35,6 +35,7 @@ pub mod otlp; pub mod proc; pub mod proxy; pub mod secrets; +pub mod startup; pub mod tags; pub mod traces; diff --git a/bottlecap/src/startup.rs b/bottlecap/src/startup.rs new file mode 100644 index 000000000..71e826c96 --- /dev/null +++ b/bottlecap/src/startup.rs @@ -0,0 +1,226 @@ +// Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +//! Cross-cutting startup assembly for the trace-processing pipeline. +//! +//! Lives in the library crate rather than under `src/bin/` because it is +//! shared by both `[[bin]]` targets: the Lambda extension (`bottlecap`) and +//! `bottlecap-test-mode`, which builds the same pipeline but attaches its +//! own router extension before spawning the agent. +//! +//! Placed at the crate root rather than under `traces/` because it wires +//! together trace, stats, proxy, lifecycle, tags, appsec, and flushing +//! pieces: it is orchestration, not a trace-domain API. + +use std::{env, sync::Arc}; + +use dogstatsd::api_key::ApiKeyFactory; +use libdd_trace_obfuscation::obfuscation_config; +use tokio::sync::{Mutex as TokioMutex, mpsc::Sender}; +use tokio_util::sync::CancellationToken; +use tracing::{error, warn}; + +use crate::appsec::processor::Processor as AppSecProcessor; +use crate::config::Config; +use crate::lifecycle::invocation::processor_service::InvocationProcessorHandle; +use crate::tags::provider::Provider as TagProvider; +use crate::traces::{ + http_client as trace_http_client, proxy_aggregator, + proxy_flusher::Flusher as ProxyFlusher, + span_dedup_service, + stats_aggregator::StatsAggregator, + stats_concentrator_service::{self, StatsConcentratorHandle}, + stats_flusher, stats_processor, trace_agent, + trace_aggregator::SendDataBuilderInfo, + trace_aggregator_service::{self, AggregatorHandle as TraceAggregatorHandle}, + trace_flusher, trace_processor, +}; + +/// Handles produced by [`build_trace_agent`] / [`start_trace_agent`]. Holds +/// the trace-channel sender, the per-domain flushers, the shutdown token, and +/// the aggregator/concentrator handles the caller needs to drive flushes and +/// shut the pipeline down. +pub struct TraceAgentPipeline { + pub trace_tx: Sender, + pub trace_flusher: Arc, + pub trace_processor: Arc, + pub stats_flusher: Arc, + pub proxy_flusher: Arc, + pub shutdown_token: CancellationToken, + pub stats_concentrator: StatsConcentratorHandle, + pub trace_aggregator_handle: TraceAggregatorHandle, +} + +/// Builds the full trace-processing pipeline (trace + stats + proxy +/// aggregators, services, flushers) and the [`trace_agent::TraceAgent`] that +/// owns the HTTP listener. Spawns the aggregator/concentrator/dedup services +/// onto the current tokio runtime; `TraceAgent::new` additionally spawns a +/// trace-payload drain task. Does **not** spawn the `TraceAgent` itself. +/// The caller owns `trace_agent` and is responsible for spawning +/// `trace_agent.start()`, optionally after further configuring it (for +/// example, via [`trace_agent::TraceAgent::with_router_extension`]). +/// +/// Note: the four background tasks started during this call (aggregator, +/// concentrator, dedup, and the trace-payload drain task inside +/// `TraceAgent::new`) have no external shutdown signal; they run until +/// their command channels are dropped. Callers that abandon the returned +/// `TraceAgent` without either spawning it or dropping the pipeline handles +/// will leak those background tasks for the lifetime of the process. +/// +/// Most callers want [`start_trace_agent`] instead, which handles the spawn. +pub fn build_trace_agent( + config: &Arc, + api_key_factory: &Arc, + tags_provider: &Arc, + invocation_processor_handle: InvocationProcessorHandle, + appsec_processor: Option>>, + client: &reqwest::Client, + proxy_aggregator: Arc>, +) -> (trace_agent::TraceAgent, TraceAgentPipeline) { + // Build one shared hyper-based HTTP client for trace and stats flushing. + // This client type is required by libdd_trace_utils for SendData::send(). + let trace_http_client = trace_http_client::create_client( + config.proxy_https.as_ref(), + config.tls_cert_file.as_ref(), + config.skip_ssl_validation, + ) + .expect("Failed to create trace HTTP client"); + + // Stats + let (stats_concentrator_service, stats_concentrator_handle) = + stats_concentrator_service::StatsConcentratorService::new(Arc::clone(config)); + tokio::spawn(stats_concentrator_service.run()); + let stats_aggregator: Arc> = Arc::new(TokioMutex::new( + StatsAggregator::new_with_concentrator(stats_concentrator_handle.clone()), + )); + let stats_flusher = Arc::new(stats_flusher::StatsFlusher::new( + api_key_factory.clone(), + stats_aggregator.clone(), + Arc::clone(config), + trace_http_client.clone(), + libdd_trace_utils::config_utils::trace_stats_url(&config.site), + )); + + let stats_processor = Arc::new(stats_processor::ServerlessStatsProcessor {}); + + // Traces + let (trace_aggregator_service, trace_aggregator_handle) = + trace_aggregator_service::AggregatorService::default(); + tokio::spawn(trace_aggregator_service.run()); + + let trace_flusher = Arc::new(trace_flusher::TraceFlusher::new( + trace_aggregator_handle.clone(), + config.clone(), + api_key_factory.clone(), + trace_http_client, + )); + + let obfuscation_config = obfuscation_config::ObfuscationConfig { + tag_replace_rules: config.apm_replace_tags.clone(), + http: obfuscation_config::HttpConfig { + remove_paths_with_digits: config.apm_config_obfuscation_http_remove_paths_with_digits, + remove_query_string: config.apm_config_obfuscation_http_remove_query_string, + }, + ..Default::default() + }; + + // The Agent's error sampler knob has no effect here: the extension's + // sampler is a plain on/off switch, not a TPS budget. + if env::var("DD_APM_ERROR_TPS").is_ok_and(|v| !v.trim().is_empty()) { + warn!( + "DD_APM_ERROR_TPS is not supported by the Lambda extension; error trace rescue is an on/off switch controlled by DD_SERVERLESS_ERROR_SAMPLER_ENABLED" + ); + } + + let trace_processor = Arc::new(trace_processor::ServerlessTraceProcessor { + obfuscation_config: Arc::new(obfuscation_config), + error_sampler: trace_processor::new_error_sampler( + config.ext.serverless_error_sampler_enabled, + ), + }); + + let (span_dedup_service, span_dedup_handle) = span_dedup_service::DedupService::new(); + tokio::spawn(span_dedup_service.run()); + + // Proxy + let proxy_flusher = Arc::new(ProxyFlusher::new( + api_key_factory.clone(), + Arc::clone(&proxy_aggregator), + Arc::clone(tags_provider), + Arc::clone(config), + client.clone(), + )); + + let trace_agent = trace_agent::TraceAgent::new( + Arc::clone(config), + trace_aggregator_handle.clone(), + trace_processor.clone(), + stats_aggregator, + stats_processor, + proxy_aggregator, + invocation_processor_handle, + appsec_processor, + Arc::clone(tags_provider), + stats_concentrator_handle.clone(), + span_dedup_handle, + ); + let pipeline = TraceAgentPipeline { + trace_tx: trace_agent.get_sender_copy(), + trace_flusher, + trace_processor, + stats_flusher, + proxy_flusher, + shutdown_token: trace_agent.shutdown_token(), + stats_concentrator: stats_concentrator_handle, + trace_aggregator_handle, + }; + + (trace_agent, pipeline) +} + +/// Builds the trace-processing pipeline with [`build_trace_agent`] and spawns +/// the [`trace_agent::TraceAgent`] HTTP listener onto the current tokio +/// runtime. Convenience entry point for callers that do not need to +/// further configure the `TraceAgent` before spawning it. +/// +/// Errors from `TraceAgent::start` (TCP bind failures, router-extension +/// failures, axum serve errors) are logged and discarded; this preserves +/// the pre-extraction behavior from `main.rs` and means the surrounding +/// pipeline keeps running with a dead trace channel. Callers that need to +/// react to startup errors should use [`build_trace_agent`] and spawn the +/// agent themselves. +/// +/// Callers that need to customize the `TraceAgent` (for example via +/// [`trace_agent::TraceAgent::with_router_extension`]) should use +/// [`build_trace_agent`] directly and spawn the returned `TraceAgent` +/// themselves after applying the extra configuration. +pub fn start_trace_agent( + config: &Arc, + api_key_factory: &Arc, + tags_provider: &Arc, + invocation_processor_handle: InvocationProcessorHandle, + appsec_processor: Option>>, + client: &reqwest::Client, + proxy_aggregator: Arc>, +) -> TraceAgentPipeline { + let (trace_agent, pipeline) = build_trace_agent( + config, + api_key_factory, + tags_provider, + invocation_processor_handle, + appsec_processor, + client, + proxy_aggregator, + ); + + // Log-only error handling preserved from the pre-extraction code in + // main.rs. See the doc comment above for callers that need reactive + // error handling. + tokio::spawn(async move { + if let Err(e) = trace_agent.start().await { + error!("Error starting trace agent: {e:?}"); + } + }); + + pipeline +} diff --git a/bottlecap/src/traces/trace_agent.rs b/bottlecap/src/traces/trace_agent.rs index aea1d084d..8500f6522 100644 --- a/bottlecap/src/traces/trace_agent.rs +++ b/bottlecap/src/traces/trace_agent.rs @@ -112,10 +112,11 @@ pub struct ProxyState { /// route to [`TraceAgent`]. /// /// Returning `Err` propagates out of [`TraceAgent::start`], aborting the -/// HTTP listener task. Note that the Lambda binary's `start_trace_agent` -/// helper spawns `start` and only logs its error; the surrounding pipeline -/// does not observe the failure. Callers that need to react to startup -/// errors must spawn the agent themselves. +/// HTTP listener task. Note that the production convenience entry point +/// [`crate::startup::start_trace_agent`] spawns `start` and only logs its +/// error; the surrounding pipeline does not observe the failure. Callers +/// that need to react to startup errors must use +/// [`crate::startup::build_trace_agent`] and spawn the agent themselves. /// /// Note that a path collision is not an `Err`: `Router::merge` panics when /// both routers define the same path, so a colliding extension aborts the From 3968c3f50b1e38c9e2c7b58ea5a6cf7f58162fa1 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 29 Apr 2026 19:31:58 -0400 Subject: [PATCH 02/23] feat(bottlecap): add bottlecap-test-mode binary MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A second [[bin]] target that runs the APM trace-processing surface as a long-lived HTTP server with no Lambda lifecycle. Listens on 127.0.0.1:8126 and exposes the standard tracer endpoints (/v0.4/traces, /v0.5/traces, /v0.6/stats, /info) plus POST /flush for deterministic harness-driven flushing. Configured by the same DD_* env vars the Lambda binary reads. Optional periodic flushing via DD_SERVERLESS_FLUSH_STRATEGY (decoupled from managed-instance mode). Gated behind the `test-mode` cargo feature (required-features), so it is not built in default or fips builds. Build with `cargo build --bin bottlecap-test-mode --features test-mode`. Intended for the cross-agent parity harness (APMSVLS-496) and for local dev workflows that need a tracer endpoint without standing up a Lambda. APMSVLS-501 🤖 Co-Authored-By: Claude Code --- bottlecap/Cargo.toml | 6 +- bottlecap/src/bin/bottlecap-test-mode/main.rs | 240 ++++++++++++++++++ 2 files changed, 245 insertions(+), 1 deletion(-) create mode 100644 bottlecap/src/bin/bottlecap-test-mode/main.rs diff --git a/bottlecap/Cargo.toml b/bottlecap/Cargo.toml index f14e48730..2e20bfb24 100644 --- a/bottlecap/Cargo.toml +++ b/bottlecap/Cargo.toml @@ -35,7 +35,7 @@ flate2 = { version = "1.1", default-features = false, features = ["rust_backend" thiserror = { version = "1.0", default-features = false } # Transitive dependency (pulled in via cookie). Pinned to >=0.3.47 so cargo audit / CI passes (RUSTSEC-2026-0009). time = { version = "0.3.47", default-features = false } -tokio = { version = "1.47", default-features = false, features = ["macros", "rt-multi-thread", "time"] } +tokio = { version = "1.47", default-features = false, features = ["macros", "rt-multi-thread", "signal", "time"] } tokio-util = { version = "0.7", default-features = false } tracing = { version = "0.1", default-features = false } tracing-core = { version = "0.1", default-features = false } @@ -117,6 +117,10 @@ flate2 = { version = "1.1", default-features = false, features = ["rust_backend" [[bin]] name = "bottlecap" +[[bin]] +name = "bottlecap-test-mode" +required-features = ["test-mode"] + [profile.dev] debug = true # same as debuginfo=2 and no stripping strip = false diff --git a/bottlecap/src/bin/bottlecap-test-mode/main.rs b/bottlecap/src/bin/bottlecap-test-mode/main.rs new file mode 100644 index 000000000..0e48bfa5f --- /dev/null +++ b/bottlecap/src/bin/bottlecap-test-mode/main.rs @@ -0,0 +1,240 @@ +//! Test-mode entry point for the bottlecap APM trace processor. +//! +//! Runs the trace-processing surface (accept -> aggregate -> flush) as a +//! long-lived HTTP server with no AWS Lambda Extension lifecycle. Intended for +//! the cross-agent parity harness ([APMSVLS-496]) and for local developer +//! workflows that need to point a tracer at bottlecap without standing up a +//! Lambda. +//! +//! Endpoints exposed on `127.0.0.1:8126`: +//! +//! | Path | Method | Source | +//! |----------------|-----------|------------------------------------------| +//! | `/v0.4/traces` | POST, PUT | trace agent | +//! | `/v0.5/traces` | POST, PUT | trace agent | +//! | `/v0.6/stats` | POST, PUT | trace agent | +//! | `/info` | GET | trace agent | +//! | `/flush` | POST | this binary's `FlushRouterExtension` | +//! +//! Environment variables this binary reads: +//! +//! | Variable | Purpose | +//! |---------------------------------|-------------------------------------------------------------------------| +//! | `DD_APM_DD_URL` | Override trace intake URL (parity harness points this at fake-intake) | +//! | `DD_SITE` | Derive stats intake URL when `DD_APM_DD_URL` is unset | +//! | `DD_SERVERLESS_FLUSH_STRATEGY` | Enable periodic flushing (e.g. `periodically,5000`); default = manual | +//! | `DD_TESTMODE_FUNCTION_ARN` | Override stub function ARN for tag generation | +//! | `DD_LOG_LEVEL` | Logging verbosity, parsed by [`bottlecap::config::log_level::LogLevel`] | +//! +//! [APMSVLS-496]: https://datadoghq.atlassian.net/browse/APMSVLS-496 + +#![deny(clippy::all)] +#![deny(clippy::pedantic)] +#![deny(clippy::unwrap_used)] +#![deny(unused_extern_crates)] +#![deny(unused_allocation)] +#![deny(unused_assignments)] +#![deny(unused_comparisons)] +#![deny(unreachable_pub)] +#![deny(missing_copy_implementations)] +#![deny(missing_debug_implementations)] + +#[cfg(not(target_env = "msvc"))] +use tikv_jemallocator::Jemalloc; + +#[cfg(not(target_env = "msvc"))] +#[global_allocator] +static GLOBAL: Jemalloc = Jemalloc; + +use std::{collections::HashMap, env, path::Path, str::FromStr, sync::Arc}; + +use axum::{Router, http::StatusCode, routing::post}; +use bottlecap::{ + LAMBDA_RUNTIME_SLUG, + config::{self, flush_strategy::FlushStrategy, log_level::LogLevel}, + flushing::FlushingService, + lifecycle::{ + flush_control::FlushControl, invocation::processor_service::InvocationProcessorHandle, + }, + logger, + logs::{aggregator_service::AggregatorService as LogsAggregatorService, flusher::LogsFlusher}, + startup::build_trace_agent, + tags::{lambda::tags::FUNCTION_ARN_KEY, provider::Provider as TagProvider}, + traces::{proxy_aggregator, trace_agent::RouterExtension}, +}; +use dogstatsd::{ + aggregator::AggregatorService as MetricsAggregatorService, api_key::ApiKeyFactory, + constants::CONTEXTS, flusher::Flusher as MetricsFlusher, metric::EMPTY_TAGS, +}; +use tokio::signal; +use tracing::error; +use tracing_subscriber::EnvFilter; +use ustr::Ustr; + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + init_ustr(); + enable_logging_subsystem(); + + // Outside Lambda every `AwsConfig` field falls back via `unwrap_or_default()`, + // so loading from env is safe with no AWS env vars set. The struct is only + // read by the secrets resolver, which test-mode bypasses. + let config = Arc::new(config::get_config(Path::new(""))); + let shared_client = bottlecap::http::get_client(&config); + + // Hardcoded literal API key. The parity harness points at a fake-intake + // that ignores auth; local dev with a real intake requires a code change. + let api_key_factory = Arc::new(ApiKeyFactory::new("stub-key")); + + let function_arn = env::var("DD_TESTMODE_FUNCTION_ARN") + .unwrap_or_else(|_| "arn:aws:lambda:us-east-1:000000000000:function:testmode".to_string()); + let metadata = HashMap::from([(FUNCTION_ARN_KEY.to_string(), function_arn)]); + let tags_provider = Arc::new(TagProvider::new( + Arc::clone(&config), + LAMBDA_RUNTIME_SLUG.to_string(), + &metadata, + )); + + let invocation_processor_handle = InvocationProcessorHandle::noop(); + + // Shared proxy aggregator backing the trace agent's proxy endpoints. + let proxy_aggregator = Arc::new(tokio::sync::Mutex::new( + proxy_aggregator::Aggregator::default(), + )); + + // Build the trace pipeline unspawned so we can attach our /flush extension + // before starting the listener. + let (trace_agent, pipeline) = build_trace_agent( + &config, + &api_key_factory, + &tags_provider, + invocation_processor_handle, + None, + &shared_client, + Arc::clone(&proxy_aggregator), + ); + let trace_flusher = Arc::clone(&pipeline.trace_flusher); + let stats_flusher = Arc::clone(&pipeline.stats_flusher); + let proxy_flusher = Arc::clone(&pipeline.proxy_flusher); + let shutdown_token = pipeline.shutdown_token.clone(); + + // FlushingService::new takes six non-optional owned values. Test-mode only + // exercises the trace/stats/proxy flushers; the logs and metrics stubs + // below stand up real services with empty queues so flushes are no-ops. + let (logs_aggregator_service, logs_aggregator_handle) = LogsAggregatorService::default(); + tokio::spawn(async move { logs_aggregator_service.run().await }); + let logs_flusher = LogsFlusher::new( + Arc::clone(&api_key_factory), + logs_aggregator_handle, + Arc::clone(&config), + shared_client.clone(), + ); + + let (metrics_aggregator_service, metrics_aggregator_handle) = + MetricsAggregatorService::new(EMPTY_TAGS, CONTEXTS).expect("metrics aggregator"); + tokio::spawn(async move { metrics_aggregator_service.run().await }); + let metrics_flushers: Arc> = Arc::new(Vec::new()); + + let flushing_service = Arc::new(FlushingService::new( + logs_flusher, + trace_flusher, + stats_flusher, + proxy_flusher, + metrics_flushers, + metrics_aggregator_handle, + None, + )); + + let flush_extension = Arc::new(FlushRouterExtension { + flushing_service: Arc::clone(&flushing_service), + }); + let trace_agent = trace_agent.with_router_extension(flush_extension); + tokio::spawn(async move { + if let Err(e) = trace_agent.start().await { + error!("Error starting trace agent: {e:?}"); + } + }); + + // Periodic flush driver. Decoupled from managed-instance mode: any non-Default + // strategy enables it. Manual flushing via POST /flush always works regardless. + if config.ext.serverless_flush_strategy != FlushStrategy::Default { + let mut interval = + FlushControl::new(config.ext.serverless_flush_strategy, config.flush_timeout) + .get_flush_interval(); + let fs = Arc::clone(&flushing_service); + let token = shutdown_token.clone(); + tokio::spawn(async move { + loop { + tokio::select! { + biased; + () = token.cancelled() => break, + _ = interval.tick() => fs.flush_blocking().await, + } + } + }); + } + + signal::ctrl_c().await?; + // Cancel before the final drain so axum's graceful shutdown drives any + // in-flight /v0.4/traces requests through the aggregator before the flush + // reads from it. + shutdown_token.cancel(); + flushing_service.flush_blocking_final().await; + Ok(()) +} + +#[derive(Debug)] +struct FlushRouterExtension { + flushing_service: Arc, +} + +impl RouterExtension for FlushRouterExtension { + fn extend(&self, router: Router) -> Result> { + let fs = Arc::clone(&self.flushing_service); + Ok(router.route( + "/flush", + post(move || { + let fs = Arc::clone(&fs); + async move { + fs.flush_blocking_final().await; + StatusCode::NO_CONTENT + } + }), + )) + } +} + +// Warm the ustr pool early so the first SortedTags::parse call (inside +// build_trace_agent and downstream) doesn't pay the 10+ ms init cost. +fn init_ustr() { + tokio::spawn(async { + Ustr::from(""); + }); +} + +fn enable_logging_subsystem() { + let log_level = LogLevel::from_str( + std::env::var("DD_LOG_LEVEL") + .unwrap_or("info".to_string()) + .as_str(), + ) + .unwrap_or(LogLevel::Info); + + let env_filter = format!( + "h2=off,hyper=off,reqwest=off,rustls=off,datadog-trace-mini-agent=off,{log_level:?}", + ); + let subscriber = tracing_subscriber::fmt::Subscriber::builder() + .with_env_filter( + EnvFilter::try_new(env_filter).expect("could not parse log level in configuration"), + ) + .with_level(true) + .with_thread_names(false) + .with_thread_ids(false) + .with_line_number(false) + .with_file(false) + .with_target(false) + .without_time() + .event_format(logger::Formatter) + .finish(); + tracing::subscriber::set_global_default(subscriber).expect("setting default subscriber failed"); +} From a0966e2806eea5db2bf94c61d6fd2342e02892d3 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 26 Aug 2026 18:03:39 -0400 Subject: [PATCH 03/23] fix(test-mode): bound and isolate the POST /flush handler flush_blocking_final expects on the metrics aggregator handle, so a dead aggregator task panicked the connection task and the harness saw a dropped connection instead of a status it could act on. The five flushers also bound only their individual HTTP calls, so stacked retries could leave a request outstanding far longer than a harness should wait. Runs the flush in a spawned task and caps it at 30s: 204 on success, 500 if the task panics, 504 after aborting a timed-out flush. Restores the hardening that previously lived in the trace agent's hardcoded /flush handler, now on the consumer side where the route lives. --- bottlecap/src/bin/bottlecap-test-mode/main.rs | 31 +++++++++++++++++-- 1 file changed, 28 insertions(+), 3 deletions(-) diff --git a/bottlecap/src/bin/bottlecap-test-mode/main.rs b/bottlecap/src/bin/bottlecap-test-mode/main.rs index 0e48bfa5f..758688216 100644 --- a/bottlecap/src/bin/bottlecap-test-mode/main.rs +++ b/bottlecap/src/bin/bottlecap-test-mode/main.rs @@ -46,7 +46,7 @@ use tikv_jemallocator::Jemalloc; #[global_allocator] static GLOBAL: Jemalloc = Jemalloc; -use std::{collections::HashMap, env, path::Path, str::FromStr, sync::Arc}; +use std::{collections::HashMap, env, path::Path, str::FromStr, sync::Arc, time::Duration}; use axum::{Router, http::StatusCode, routing::post}; use bottlecap::{ @@ -188,6 +188,11 @@ struct FlushRouterExtension { flushing_service: Arc, } +/// Upper bound on a single `POST /flush`. The flushers already bound their own +/// HTTP calls via `flush_timeout`, but retries across the five flushers can +/// stack, so this caps total wall-clock time for the harness. +const FLUSH_REQUEST_TIMEOUT: Duration = Duration::from_secs(30); + impl RouterExtension for FlushRouterExtension { fn extend(&self, router: Router) -> Result> { let fs = Arc::clone(&self.flushing_service); @@ -196,8 +201,28 @@ impl RouterExtension for FlushRouterExtension { post(move || { let fs = Arc::clone(&fs); async move { - fs.flush_blocking_final().await; - StatusCode::NO_CONTENT + // Isolate panics and bound execution time. flush_blocking_final + // expects on the metrics aggregator handle, so a dead aggregator + // task would otherwise panic the connection task instead of + // returning a status the harness can act on. + let mut task = tokio::task::spawn(async move { + fs.flush_blocking_final().await; + }); + match tokio::time::timeout(FLUSH_REQUEST_TIMEOUT, &mut task).await { + Ok(Ok(())) => StatusCode::NO_CONTENT, + Ok(Err(e)) => { + error!("Flush task failed: {e:?}"); + StatusCode::INTERNAL_SERVER_ERROR + } + Err(_) => { + task.abort(); + error!( + "Flush timed out after {}s, aborting", + FLUSH_REQUEST_TIMEOUT.as_secs() + ); + StatusCode::GATEWAY_TIMEOUT + } + } } }), )) From a6b6a3ba37fbf42d280c9503569b71302422bd0b Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 2 Sep 2026 16:56:25 -0400 Subject: [PATCH 04/23] ci: lint the test-mode feature The bottlecap-test-mode binary is gated behind required-features, so no existing CI job compiles it and a change to library code could break it without any job failing. Add a clippy pass with the test-mode feature enabled to both the GitHub Actions and GitLab pipelines. --- .github/workflows/rs_ci.yml | 4 ++-- .gitlab/templates/pipeline.yaml.tpl | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/rs_ci.yml b/.github/workflows/rs_ci.yml index 0518386dc..a5572db18 100644 --- a/.github/workflows/rs_ci.yml +++ b/.github/workflows/rs_ci.yml @@ -82,8 +82,8 @@ jobs: run: cargo clippy --workspace --all-targets --features default - working-directory: bottlecap run: cargo clippy --workspace --all-targets --no-default-features --features fips - # No other job compiles the test-mode feature: it gates test-only - # constructors that are absent from default and fips builds. + # The test-mode feature gates the bottlecap-test-mode binary via + # required-features, so no other job compiles it. - working-directory: bottlecap run: cargo clippy --workspace --all-targets --features default,test-mode diff --git a/.gitlab/templates/pipeline.yaml.tpl b/.gitlab/templates/pipeline.yaml.tpl index d0223a677..4f49c591c 100644 --- a/.gitlab/templates/pipeline.yaml.tpl +++ b/.gitlab/templates/pipeline.yaml.tpl @@ -63,8 +63,8 @@ cargo clippy: # We need to do these separately because the fips feature is incompatible with the default feature. - cargo clippy --workspace --features default - cargo clippy --workspace --no-default-features --features fips - # No other job compiles the test-mode feature: it gates test-only - # constructors that are absent from default and fips builds. + # The test-mode feature gates the bottlecap-test-mode binary via + # required-features, so no other job compiles it. - cargo clippy --workspace --features default,test-mode {{ range $flavor := (ds "flavors").flavors }} From 2355edf11ce69e4e73f98c15a08097ee4dd7741e Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Thu, 10 Sep 2026 15:32:42 -0400 Subject: [PATCH 05/23] fix(test-mode): drain queued payloads before flushing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Trace and stats requests only await the hand-off into an intermediate channel, so a request can be answered while its payload is still queued ahead of the aggregator. A flush issued right after an accepted request could therefore miss it, leaving the payload buffered until the next flush and making the parity harness nondeterministic. Add an ingest barrier over both forwarder tasks and await it in POST /flush and in the shutdown drain. 🤖 --- bottlecap/src/bin/bottlecap-test-mode/main.rs | 18 +- bottlecap/src/startup.rs | 9 +- bottlecap/src/traces/trace_agent.rs | 167 +++++++++++++++--- 3 files changed, 165 insertions(+), 29 deletions(-) diff --git a/bottlecap/src/bin/bottlecap-test-mode/main.rs b/bottlecap/src/bin/bottlecap-test-mode/main.rs index 758688216..b0c17ef5d 100644 --- a/bottlecap/src/bin/bottlecap-test-mode/main.rs +++ b/bottlecap/src/bin/bottlecap-test-mode/main.rs @@ -60,7 +60,10 @@ use bottlecap::{ logs::{aggregator_service::AggregatorService as LogsAggregatorService, flusher::LogsFlusher}, startup::build_trace_agent, tags::{lambda::tags::FUNCTION_ARN_KEY, provider::Provider as TagProvider}, - traces::{proxy_aggregator, trace_agent::RouterExtension}, + traces::{ + proxy_aggregator, + trace_agent::{IngestBarrier, RouterExtension}, + }, }; use dogstatsd::{ aggregator::AggregatorService as MetricsAggregatorService, api_key::ApiKeyFactory, @@ -145,8 +148,10 @@ async fn main() -> anyhow::Result<()> { None, )); + let ingest_barrier = trace_agent.ingest_barrier(); let flush_extension = Arc::new(FlushRouterExtension { flushing_service: Arc::clone(&flushing_service), + ingest_barrier: ingest_barrier.clone(), }); let trace_agent = trace_agent.with_router_extension(flush_extension); tokio::spawn(async move { @@ -179,6 +184,9 @@ async fn main() -> anyhow::Result<()> { // in-flight /v0.4/traces requests through the aggregator before the flush // reads from it. shutdown_token.cancel(); + // Graceful shutdown only guarantees the handlers returned; the accepted + // payloads may still be queued ahead of the aggregators. + ingest_barrier.wait().await; flushing_service.flush_blocking_final().await; Ok(()) } @@ -186,6 +194,7 @@ async fn main() -> anyhow::Result<()> { #[derive(Debug)] struct FlushRouterExtension { flushing_service: Arc, + ingest_barrier: IngestBarrier, } /// Upper bound on a single `POST /flush`. The flushers already bound their own @@ -196,16 +205,23 @@ const FLUSH_REQUEST_TIMEOUT: Duration = Duration::from_secs(30); impl RouterExtension for FlushRouterExtension { fn extend(&self, router: Router) -> Result> { let fs = Arc::clone(&self.flushing_service); + let barrier = self.ingest_barrier.clone(); Ok(router.route( "/flush", post(move || { let fs = Arc::clone(&fs); + let barrier = barrier.clone(); async move { // Isolate panics and bound execution time. flush_blocking_final // expects on the metrics aggregator handle, so a dead aggregator // task would otherwise panic the connection task instead of // returning a status the harness can act on. let mut task = tokio::task::spawn(async move { + // A payload can be accepted, and its request answered, + // while it is still queued ahead of the aggregators. + // Drain those queues first so this flush is + // deterministic from the caller's point of view. + barrier.wait().await; fs.flush_blocking_final().await; }); match tokio::time::timeout(FLUSH_REQUEST_TIMEOUT, &mut task).await { diff --git a/bottlecap/src/startup.rs b/bottlecap/src/startup.rs index 71e826c96..dee04dccc 100644 --- a/bottlecap/src/startup.rs +++ b/bottlecap/src/startup.rs @@ -54,14 +54,15 @@ pub struct TraceAgentPipeline { /// Builds the full trace-processing pipeline (trace + stats + proxy /// aggregators, services, flushers) and the [`trace_agent::TraceAgent`] that /// owns the HTTP listener. Spawns the aggregator/concentrator/dedup services -/// onto the current tokio runtime; `TraceAgent::new` additionally spawns a -/// trace-payload drain task. Does **not** spawn the `TraceAgent` itself. +/// onto the current tokio runtime; `TraceAgent::new` additionally spawns the +/// trace- and stats-payload forwarder tasks. Does **not** spawn the +/// `TraceAgent` itself. /// The caller owns `trace_agent` and is responsible for spawning /// `trace_agent.start()`, optionally after further configuring it (for /// example, via [`trace_agent::TraceAgent::with_router_extension`]). /// -/// Note: the four background tasks started during this call (aggregator, -/// concentrator, dedup, and the trace-payload drain task inside +/// Note: the background tasks started during this call (aggregator, +/// concentrator, dedup, and the two payload forwarder tasks inside /// `TraceAgent::new`) have no external shutdown signal; they run until /// their command channels are dropped. Callers that abandon the returned /// `TraceAgent` without either spawning it or dropping the pipeline handles diff --git a/bottlecap/src/traces/trace_agent.rs b/bottlecap/src/traces/trace_agent.rs index 8500f6522..3d9c2024b 100644 --- a/bottlecap/src/traces/trace_agent.rs +++ b/bottlecap/src/traces/trace_agent.rs @@ -16,6 +16,7 @@ use std::time::Instant; use tokio::sync::{ Mutex, mpsc::{self, Receiver, Sender}, + oneshot, }; use tokio_util::sync::CancellationToken; use tower_http::limit::RequestBodyLimitLayer; @@ -79,6 +80,7 @@ const INSTRUMENTATION_INTAKE_PATH: &str = "/api/v2/apmtelemetry"; const TRACER_PAYLOAD_CHANNEL_BUFFER_SIZE: usize = 10; const STATS_PAYLOAD_CHANNEL_BUFFER_SIZE: usize = 10; +const BARRIER_CHANNEL_BUFFER_SIZE: usize = 1; pub const TRACE_REQUEST_BODY_LIMIT: usize = 50 * 1024 * 1024; pub const DEFAULT_REQUEST_BODY_LIMIT: usize = 2 * 1024 * 1024; pub const MAX_CONTENT_LENGTH: usize = 50 * 1024 * 1024; @@ -135,6 +137,40 @@ pub trait RouterExtension: Send + Sync { fn extend(&self, router: Router) -> Result>; } +/// Barrier over the forwarder tasks that move accepted payloads from the +/// request handlers into the trace and stats aggregators. +/// +/// The handlers only `await` the hand-off into an intermediate channel, so a +/// request can return `200` while its payload is still queued. [`wait`] +/// returns once every payload enqueued before the call has reached its +/// aggregator, which makes a flush issued afterwards observe it. +/// +/// [`wait`]: IngestBarrier::wait +#[derive(Clone, Debug)] +pub struct IngestBarrier { + trace_tx: Sender>, + stats_tx: Sender>, +} + +impl IngestBarrier { + pub async fn wait(&self) { + Self::wait_for(&self.trace_tx).await; + Self::wait_for(&self.stats_tx).await; + } + + /// Each forwarder task selects on its payload channel before its barrier + /// channel, so an acknowledgement can only be sent on an iteration where + /// the payload channel was empty. + async fn wait_for(tx: &Sender>) { + let (ack_tx, ack_rx) = oneshot::channel(); + // A closed channel or a dropped acknowledgement means the forwarder + // task is gone, so there is nothing left to drain. + if tx.send(ack_tx).await.is_ok() { + let _ = ack_rx.await; + } + } +} + pub struct TraceAgent { pub config: Arc, pub trace_processor: Arc, @@ -146,6 +182,8 @@ pub struct TraceAgent { appsec_processor: Option>>, shutdown_token: CancellationToken, tx: Sender, + stats_tx: Sender, + ingest_barrier: IngestBarrier, stats_concentrator: StatsConcentratorHandle, span_deduper: DedupHandle, /// `None` when the caller wants no extra routes. See @@ -180,12 +218,51 @@ impl TraceAgent { // processed trace payloads to our trace aggregator. let (trace_tx, mut trace_rx): (Sender, Receiver) = mpsc::channel(TRACER_PAYLOAD_CHANNEL_BUFFER_SIZE); + let (trace_barrier_tx, mut trace_barrier_rx) = + mpsc::channel::>(BARRIER_CHANNEL_BUFFER_SIZE); // Start the trace aggregator, which receives and buffers trace payloads to be consumed by the trace flusher. tokio::spawn(async move { - while let Some(tracer_payload_info) = trace_rx.recv().await { - if let Err(e) = aggregator_handle.insert_payload(tracer_payload_info) { - error!("TRACE_AGENT | Failed to insert payload into aggregator: {e}"); + loop { + tokio::select! { + biased; + tracer_payload_info = trace_rx.recv() => { + let Some(tracer_payload_info) = tracer_payload_info else { break }; + if let Err(e) = aggregator_handle.insert_payload(tracer_payload_info) { + error!("TRACE_AGENT | Failed to insert payload into aggregator: {e}"); + } + } + // Reached only on an iteration where the payload channel is + // empty, so everything queued before the barrier request is + // already in the aggregator. See [`IngestBarrier`]. + Some(ack) = trace_barrier_rx.recv() => { + let _ = ack.send(()); + } + } + } + }); + + // Set up a channel to send processed stats to our stats aggregator. + let (stats_tx, mut stats_rx): ( + Sender, + Receiver, + ) = mpsc::channel(STATS_PAYLOAD_CHANNEL_BUFFER_SIZE); + let (stats_barrier_tx, mut stats_barrier_rx) = + mpsc::channel::>(BARRIER_CHANNEL_BUFFER_SIZE); + + // Start the stats aggregator, which receives and buffers stats payloads to be consumed by the stats flusher. + let stats_aggregator_task = stats_aggregator.clone(); + tokio::spawn(async move { + loop { + tokio::select! { + biased; + stats_payload = stats_rx.recv() => { + let Some(stats_payload) = stats_payload else { break }; + stats_aggregator_task.lock().await.add(stats_payload); + } + Some(ack) = stats_barrier_rx.recv() => { + let _ = ack.send(()); + } } } }); @@ -200,6 +277,11 @@ impl TraceAgent { appsec_processor, tags_provider, tx: trace_tx, + stats_tx, + ingest_barrier: IngestBarrier { + trace_tx: trace_barrier_tx, + stats_tx: stats_barrier_tx, + }, shutdown_token: CancellationToken::new(), stats_concentrator, span_deduper, @@ -221,22 +303,7 @@ impl TraceAgent { pub async fn start(&self) -> Result<(), Box> { let now = Instant::now(); - // Set up a channel to send processed stats to our stats aggregator. - let (stats_tx, mut stats_rx): ( - Sender, - Receiver, - ) = mpsc::channel(STATS_PAYLOAD_CHANNEL_BUFFER_SIZE); - - // Start the stats aggregator, which receives and buffers stats payloads to be consumed by the stats flusher. - let stats_aggregator = self.stats_aggregator.clone(); - tokio::spawn(async move { - while let Some(stats_payload) = stats_rx.recv().await { - let mut aggregator = stats_aggregator.lock().await; - aggregator.add(stats_payload); - } - }); - - let router = self.make_router(stats_tx)?; + let router = self.make_router(self.stats_tx.clone())?; let port = u16::try_from(TRACE_AGENT_PORT).expect("TRACE_AGENT_PORT is too large"); let socket = SocketAddr::from(([127, 0, 0, 1], port)); @@ -785,6 +852,13 @@ impl TraceAgent { pub fn shutdown_token(&self) -> CancellationToken { self.shutdown_token.clone() } + + /// Barrier over the trace and stats forwarder tasks. Await it before a + /// flush to make that flush observe every payload accepted so far. + #[must_use] + pub fn ingest_barrier(&self) -> IngestBarrier { + self.ingest_barrier.clone() + } } fn handle_reparenting(reparenting_info: &mut VecDeque, span: &mut pb::Span) { @@ -854,13 +928,17 @@ mod tests { LAMBDA_RUNTIME_SLUG, config, traces::{ span_dedup_service::DedupService, stats_concentrator_service::StatsConcentratorService, - trace_aggregator_service::AggregatorService, + trace_aggregator::OwnedTracerHeaderTags, trace_aggregator_service::AggregatorService, }, }; use axum::body::Body; use axum::http::{HeaderMap, HeaderName, HeaderValue, Request}; + use libdd_common::Endpoint; use libdd_trace_obfuscation::obfuscation_config::ObfuscationConfig; - use libdd_trace_utils::trace_utils::TracerHeaderTags; + use libdd_trace_utils::{ + send_data::SendDataBuilder, trace_utils::TracerHeaderTags, + tracer_payload::TracerPayloadCollection, + }; use std::collections::HashMap; use std::sync::atomic::{AtomicUsize, Ordering}; use tower::ServiceExt; @@ -982,6 +1060,10 @@ mod tests { } fn build_test_agent() -> TraceAgent { + build_test_agent_with_aggregator().0 + } + + fn build_test_agent_with_aggregator() -> (TraceAgent, AggregatorHandle) { let config = Arc::new(config::Config::default()); let (concentrator_svc, concentrator) = StatsConcentratorService::new(Arc::clone(&config)); tokio::spawn(concentrator_svc.run()); @@ -1003,9 +1085,9 @@ mod tests { &HashMap::from([("function_arn".to_string(), "test-arn".to_string())]), )); - TraceAgent::new( + let agent = TraceAgent::new( config, - aggregator_handle, + aggregator_handle.clone(), trace_processor, stats_aggregator, Arc::new(stats_processor::ServerlessStatsProcessor {}), @@ -1015,7 +1097,44 @@ mod tests { tags_provider, concentrator, dedup_handle, - ) + ); + + (agent, aggregator_handle) + } + + fn stub_payload() -> SendDataBuilderInfo { + let header_tags = TracerHeaderTags::default(); + let owned_tags = OwnedTracerHeaderTags::from(header_tags.clone()); + let size = 1; + let builder = SendDataBuilder::new( + size, + TracerPayloadCollection::V07(Vec::new()), + header_tags, + &Endpoint::from_slice("localhost"), + ); + + SendDataBuilderInfo::new(builder, size, owned_tags) + } + + /// A payload is only handed to the aggregator by a background forwarder + /// task, so accepting it does not by itself make it visible to a flush. + /// `ingest_barrier` closes that gap: on this single-threaded runtime the + /// forwarder cannot have run before the barrier is awaited. + #[tokio::test] + async fn ingest_barrier_waits_for_queued_payloads_to_reach_the_aggregator() { + let (agent, aggregator_handle) = build_test_agent_with_aggregator(); + let trace_tx = agent.get_sender_copy(); + let barrier = agent.ingest_barrier(); + + for _ in 0..3 { + trace_tx.send(stub_payload()).await.expect("send payload"); + } + + barrier.wait().await; + + let batches = aggregator_handle.get_batches().await.expect("get_batches"); + let payloads: usize = batches.iter().map(Vec::len).sum(); + assert_eq!(payloads, 3); } #[tokio::test] From 3b5b717b403a8987955726d3c8e74f7d29fb15a6 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Thu, 10 Sep 2026 15:37:59 -0400 Subject: [PATCH 06/23] fix(test-mode): send stats to the overridden trace intake MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DD_APM_DD_URL only moved the trace intake, so a harness pointing traces at a local fake-intake still sent stats toward Datadog with the stub API key, leaving the binary's advertised /v0.6/stats path unexercised. Derive the stats endpoint from the resolved trace intake in test mode. With DD_APM_DD_URL unset this reproduces the site-derived default, and production routing is unchanged. 🤖 --- bottlecap/src/bin/bottlecap-test-mode/main.rs | 46 ++++++++++++++++++- bottlecap/src/startup.rs | 12 ++++- 2 files changed, 55 insertions(+), 3 deletions(-) diff --git a/bottlecap/src/bin/bottlecap-test-mode/main.rs b/bottlecap/src/bin/bottlecap-test-mode/main.rs index b0c17ef5d..f5d20f87f 100644 --- a/bottlecap/src/bin/bottlecap-test-mode/main.rs +++ b/bottlecap/src/bin/bottlecap-test-mode/main.rs @@ -20,8 +20,8 @@ //! //! | Variable | Purpose | //! |---------------------------------|-------------------------------------------------------------------------| -//! | `DD_APM_DD_URL` | Override trace intake URL (parity harness points this at fake-intake) | -//! | `DD_SITE` | Derive stats intake URL when `DD_APM_DD_URL` is unset | +//! | `DD_APM_DD_URL` | Override trace intake URL; stats follow it (harness points at fake-intake) | +//! | `DD_SITE` | Derive trace and stats intake URLs when `DD_APM_DD_URL` is unset | //! | `DD_SERVERLESS_FLUSH_STRATEGY` | Enable periodic flushing (e.g. `periodically,5000`); default = manual | //! | `DD_TESTMODE_FUNCTION_ARN` | Override stub function ARN for tag generation | //! | `DD_LOG_LEVEL` | Logging verbosity, parsed by [`bottlecap::config::log_level::LogLevel`] | @@ -115,6 +115,7 @@ async fn main() -> anyhow::Result<()> { None, &shared_client, Arc::clone(&proxy_aggregator), + Some(stats_url_from_trace_intake(&config.apm_dd_url)), ); let trace_flusher = Arc::clone(&pipeline.trace_flusher); let stats_flusher = Arc::clone(&pipeline.stats_flusher); @@ -191,6 +192,25 @@ async fn main() -> anyhow::Result<()> { Ok(()) } +/// Path the config crate appends to `DD_APM_DD_URL` to build `apm_dd_url`. +const TRACE_INTAKE_ROUTE: &str = "/api/v0.2/traces"; + +/// Point stats at the same host as traces. +/// +/// `DD_APM_DD_URL` only moves the trace intake; stats would otherwise be +/// derived from `DD_SITE` and leave the harness's fake-intake, so the +/// binary's `/v0.6/stats` path could not be exercised locally. `apm_dd_url` +/// is already a fully-resolved trace endpoint, so strip the trace route +/// before appending the stats one. With `DD_APM_DD_URL` unset this +/// reproduces the site-derived default. +fn stats_url_from_trace_intake(apm_dd_url: &str) -> String { + libdd_trace_utils::config_utils::trace_stats_url_prefixed( + apm_dd_url + .trim_end_matches('/') + .trim_end_matches(TRACE_INTAKE_ROUTE), + ) +} + #[derive(Debug)] struct FlushRouterExtension { flushing_service: Arc, @@ -279,3 +299,25 @@ fn enable_logging_subsystem() { .finish(); tracing::subscriber::set_global_default(subscriber).expect("setting default subscriber failed"); } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn stats_url_follows_the_overridden_trace_intake() { + assert_eq!( + stats_url_from_trace_intake("http://127.0.0.1:8080/api/v0.2/traces"), + "http://127.0.0.1:8080/api/v0.2/stats" + ); + } + + #[test] + fn stats_url_matches_the_site_default_when_not_overridden() { + let site = "datadoghq.com"; + assert_eq!( + stats_url_from_trace_intake(&libdd_trace_utils::config_utils::trace_intake_url(site)), + libdd_trace_utils::config_utils::trace_stats_url(site) + ); + } +} diff --git a/bottlecap/src/startup.rs b/bottlecap/src/startup.rs index dee04dccc..b1ff66d16 100644 --- a/bottlecap/src/startup.rs +++ b/bottlecap/src/startup.rs @@ -68,7 +68,14 @@ pub struct TraceAgentPipeline { /// `TraceAgent` without either spawning it or dropping the pipeline handles /// will leak those background tasks for the lifetime of the process. /// +/// `stats_url_override` replaces the stats intake that would otherwise be +/// derived from `DD_SITE`. `DD_APM_DD_URL` only moves the trace intake, so a +/// caller that redirects traces to a local intake must pass the matching +/// stats endpoint here or its stats will still be sent to Datadog. `None` +/// keeps the site-derived default. +/// /// Most callers want [`start_trace_agent`] instead, which handles the spawn. +#[allow(clippy::too_many_arguments)] pub fn build_trace_agent( config: &Arc, api_key_factory: &Arc, @@ -77,6 +84,7 @@ pub fn build_trace_agent( appsec_processor: Option>>, client: &reqwest::Client, proxy_aggregator: Arc>, + stats_url_override: Option, ) -> (trace_agent::TraceAgent, TraceAgentPipeline) { // Build one shared hyper-based HTTP client for trace and stats flushing. // This client type is required by libdd_trace_utils for SendData::send(). @@ -99,7 +107,8 @@ pub fn build_trace_agent( stats_aggregator.clone(), Arc::clone(config), trace_http_client.clone(), - libdd_trace_utils::config_utils::trace_stats_url(&config.site), + stats_url_override + .unwrap_or_else(|| libdd_trace_utils::config_utils::trace_stats_url(&config.site)), )); let stats_processor = Arc::new(stats_processor::ServerlessStatsProcessor {}); @@ -212,6 +221,7 @@ pub fn start_trace_agent( appsec_processor, client, proxy_aggregator, + None, ); // Log-only error handling preserved from the pre-extraction code in From 1d9de67b022e52f0ee3a55b85d8ee5892d416b51 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Thu, 10 Sep 2026 15:48:37 -0400 Subject: [PATCH 07/23] fix(test-mode): fail POST /flush when delivery fails MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The blocking flush path discarded every flusher result, so payloads that could not be delivered were dropped while POST /flush still answered 204. The harness read that as a successful drain. Report undelivered payloads from the blocking flush, log which domains dropped data, and answer 502 instead of 204 when any did. 🤖 --- bottlecap/src/bin/bottlecap-test-mode/main.rs | 15 ++++-- bottlecap/src/flushing/service.rs | 53 ++++++++++++++++--- 2 files changed, 59 insertions(+), 9 deletions(-) diff --git a/bottlecap/src/bin/bottlecap-test-mode/main.rs b/bottlecap/src/bin/bottlecap-test-mode/main.rs index f5d20f87f..d5b90722f 100644 --- a/bottlecap/src/bin/bottlecap-test-mode/main.rs +++ b/bottlecap/src/bin/bottlecap-test-mode/main.rs @@ -174,7 +174,9 @@ async fn main() -> anyhow::Result<()> { tokio::select! { biased; () = token.cancelled() => break, - _ = interval.tick() => fs.flush_blocking().await, + // The periodic driver has no caller to report to; the + // flushing service already logs what it dropped. + _ = interval.tick() => { fs.flush_blocking().await; }, } } }); @@ -242,10 +244,17 @@ impl RouterExtension for FlushRouterExtension { // Drain those queues first so this flush is // deterministic from the caller's point of view. barrier.wait().await; - fs.flush_blocking_final().await; + fs.flush_blocking_final().await }); match tokio::time::timeout(FLUSH_REQUEST_TIMEOUT, &mut task).await { - Ok(Ok(())) => StatusCode::NO_CONTENT, + Ok(Ok(false)) => StatusCode::NO_CONTENT, + // The flush ran, but a flusher gave up on payloads it + // could not deliver and they were dropped. Reporting 204 + // here would tell the harness the drain succeeded. + Ok(Ok(true)) => { + error!("Flush completed with undelivered payloads"); + StatusCode::BAD_GATEWAY + } Ok(Err(e)) => { error!("Flush task failed: {e:?}"); StatusCode::INTERNAL_SERVER_ERROR diff --git a/bottlecap/src/flushing/service.rs b/bottlecap/src/flushing/service.rs index a8181d7e9..edd4ac3bf 100644 --- a/bottlecap/src/flushing/service.rs +++ b/bottlecap/src/flushing/service.rs @@ -300,8 +300,11 @@ impl FlushingService { /// /// The stats flusher respects its normal timing constraints (time-based bucketing), /// which may result in some stats being held back until the next flush cycle. - pub async fn flush_blocking(&self) { - self.flush_blocking_inner(false).await; + /// + /// Returns `true` when at least one flusher finished with payloads it could + /// not deliver. See [`flush_blocking_inner`](Self::flush_blocking_inner). + pub async fn flush_blocking(&self) -> bool { + self.flush_blocking_inner(false).await } /// Performs a final blocking flush of all telemetry data before shutdown. @@ -311,14 +314,23 @@ impl FlushingService { /// flush immediately regardless of its normal timing constraints. /// /// Use this during shutdown when this is the last opportunity to send data. - pub async fn flush_blocking_final(&self) { - self.flush_blocking_inner(true).await; + /// + /// Returns `true` when at least one flusher finished with payloads it could + /// not deliver. See [`flush_blocking_inner`](Self::flush_blocking_inner). + pub async fn flush_blocking_final(&self) -> bool { + self.flush_blocking_inner(true).await } /// Internal implementation for blocking flush operations. /// /// Fetches metrics from the aggregator and flushes all data types in parallel. - async fn flush_blocking_inner(&self, force_stats: bool) { + /// + /// Returns `true` when at least one flusher handed back payloads it could + /// not deliver after exhausting its own retries. Unlike the + /// [`spawn_non_blocking`](Self::spawn_non_blocking) path, this one does not + /// redrive them, so the return value is the only signal that data was + /// drained without reaching the intake. + async fn flush_blocking_inner(&self, force_stats: bool) -> bool { let flush_response = self .metrics_aggr_handle .flush() @@ -341,13 +353,42 @@ impl FlushingService { dsm.drain_into_proxy().await; } - tokio::join!( + let (logs, metrics, traces, stats, proxy) = tokio::join!( self.logs_flusher.flush(None), futures::future::join_all(metrics_futures), self.trace_flusher.flush(None), self.stats_flusher.flush(force_stats, None), self.proxy_flusher.flush(None), ); + + let undelivered = [ + ("logs", !logs.is_empty()), + ( + "metrics", + metrics.iter().any(|retry| { + retry.as_ref().is_some_and(|(series, sketches)| { + !series.is_empty() || !sketches.is_empty() + }) + }), + ), + ("traces", traces.is_some_and(|t| !t.is_empty())), + ("stats", stats.is_some_and(|s| !s.is_empty())), + ("proxy", proxy.is_some_and(|p| !p.is_empty())), + ]; + + let failed: Vec<&str> = undelivered + .iter() + .filter_map(|(name, failed)| failed.then_some(*name)) + .collect(); + + if !failed.is_empty() { + error!( + "FLUSHING_SERVICE | Dropping undelivered payloads after blocking flush: {}", + failed.join(", ") + ); + } + + !failed.is_empty() } } From 9a0d99fa7d6c2b79b9791062da4c07e6234ceefb Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Thu, 10 Sep 2026 15:50:30 -0400 Subject: [PATCH 08/23] fix(test-mode): await the listener before the final drain MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cancelling the shutdown token only signals graceful shutdown, so the final drain could run while a request handler was still processing and its payload would be dropped when the process exited. Await the listener task, bounded so a lingering connection cannot wedge shutdown. 🤖 --- bottlecap/src/bin/bottlecap-test-mode/main.rs | 22 ++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/bottlecap/src/bin/bottlecap-test-mode/main.rs b/bottlecap/src/bin/bottlecap-test-mode/main.rs index d5b90722f..8580c6bb7 100644 --- a/bottlecap/src/bin/bottlecap-test-mode/main.rs +++ b/bottlecap/src/bin/bottlecap-test-mode/main.rs @@ -155,7 +155,7 @@ async fn main() -> anyhow::Result<()> { ingest_barrier: ingest_barrier.clone(), }); let trace_agent = trace_agent.with_router_extension(flush_extension); - tokio::spawn(async move { + let listener_task = tokio::spawn(async move { if let Err(e) = trace_agent.start().await { error!("Error starting trace agent: {e:?}"); } @@ -187,8 +187,20 @@ async fn main() -> anyhow::Result<()> { // in-flight /v0.4/traces requests through the aggregator before the flush // reads from it. shutdown_token.cancel(); - // Graceful shutdown only guarantees the handlers returned; the accepted - // payloads may still be queued ahead of the aggregators. + // Cancelling only signals the shutdown; awaiting the listener is what + // guarantees in-flight handlers have finished. Bounded so a lingering + // connection cannot wedge shutdown: draining late data beats hanging. + if tokio::time::timeout(SHUTDOWN_TIMEOUT, listener_task) + .await + .is_err() + { + error!( + "Trace agent did not shut down within {}s, draining anyway", + SHUTDOWN_TIMEOUT.as_secs() + ); + } + // Handlers returning does not mean their payloads reached the + // aggregators; they may still be queued ahead of them. ingest_barrier.wait().await; flushing_service.flush_blocking_final().await; Ok(()) @@ -224,6 +236,10 @@ struct FlushRouterExtension { /// stack, so this caps total wall-clock time for the harness. const FLUSH_REQUEST_TIMEOUT: Duration = Duration::from_secs(30); +/// Upper bound on waiting for the HTTP listener to finish its graceful +/// shutdown before the final drain runs. +const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(10); + impl RouterExtension for FlushRouterExtension { fn extend(&self, router: Router) -> Result> { let fs = Arc::clone(&self.flushing_service); From d069f384c23bea50ac141f38c8642b6dd2b9d845 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Thu, 10 Sep 2026 15:54:19 -0400 Subject: [PATCH 09/23] fix(test-mode): exit when the listener fails to start MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A failure to bind port 8126 was only logged, leaving a live process with no listener while main waited on ctrl-c. The harness saw a healthy process, or connected to whatever already held the port. Propagate the listener error and race it against ctrl-c so startup failures end the process with a non-zero exit. 🤖 --- bottlecap/src/bin/bottlecap-test-mode/main.rs | 35 +++++++++++++------ 1 file changed, 24 insertions(+), 11 deletions(-) diff --git a/bottlecap/src/bin/bottlecap-test-mode/main.rs b/bottlecap/src/bin/bottlecap-test-mode/main.rs index 8580c6bb7..9fff6fdfb 100644 --- a/bottlecap/src/bin/bottlecap-test-mode/main.rs +++ b/bottlecap/src/bin/bottlecap-test-mode/main.rs @@ -155,10 +155,14 @@ async fn main() -> anyhow::Result<()> { ingest_barrier: ingest_barrier.clone(), }); let trace_agent = trace_agent.with_router_extension(flush_extension); - let listener_task = tokio::spawn(async move { - if let Err(e) = trace_agent.start().await { - error!("Error starting trace agent: {e:?}"); - } + // Errors are returned rather than logged so that a startup failure (port + // 8126 already bound, for instance) ends the process instead of leaving a + // live one with no listener for the harness to connect to. + let mut listener_task = tokio::spawn(async move { + trace_agent + .start() + .await + .map_err(|e| anyhow::anyhow!("trace agent failed: {e}")) }); // Periodic flush driver. Decoupled from managed-instance mode: any non-Default @@ -182,7 +186,16 @@ async fn main() -> anyhow::Result<()> { }); } - signal::ctrl_c().await?; + // The listener finishing first means it never came up, or stopped serving + // without being asked to. Either way there is nothing left to drain. + tokio::select! { + result = signal::ctrl_c() => result?, + result = &mut listener_task => match result? { + Ok(()) => anyhow::bail!("trace agent listener stopped unexpectedly"), + Err(e) => return Err(e), + }, + } + // Cancel before the final drain so axum's graceful shutdown drives any // in-flight /v0.4/traces requests through the aggregator before the flush // reads from it. @@ -190,14 +203,14 @@ async fn main() -> anyhow::Result<()> { // Cancelling only signals the shutdown; awaiting the listener is what // guarantees in-flight handlers have finished. Bounded so a lingering // connection cannot wedge shutdown: draining late data beats hanging. - if tokio::time::timeout(SHUTDOWN_TIMEOUT, listener_task) - .await - .is_err() - { - error!( + match tokio::time::timeout(SHUTDOWN_TIMEOUT, listener_task).await { + Ok(Ok(Ok(()))) => {} + Ok(Ok(Err(e))) => error!("Trace agent shut down with an error: {e:?}"), + Ok(Err(e)) => error!("Trace agent task failed: {e:?}"), + Err(_) => error!( "Trace agent did not shut down within {}s, draining anyway", SHUTDOWN_TIMEOUT.as_secs() - ); + ), } // Handlers returning does not mean their payloads reached the // aggregators; they may still be queued ahead of them. From bb6e9200cd21986cc0dbad591a2e44f5fdce953b Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Thu, 10 Sep 2026 17:17:38 -0400 Subject: [PATCH 10/23] fix(test-mode): serialize overlapping flushes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Each flusher drains its aggregator before awaiting network delivery, so the periodic driver, POST /flush, and the shutdown drain could overlap and let one report success on empty queues while another's send was still in flight. All three now run the barrier-plus-flush sequence under a shared lock, which also makes the final drain wait for an in-flight periodic flush. 🤖 --- bottlecap/src/bin/bottlecap-test-mode/main.rs | 114 +++++++++++++----- 1 file changed, 84 insertions(+), 30 deletions(-) diff --git a/bottlecap/src/bin/bottlecap-test-mode/main.rs b/bottlecap/src/bin/bottlecap-test-mode/main.rs index 9fff6fdfb..76437bcaa 100644 --- a/bottlecap/src/bin/bottlecap-test-mode/main.rs +++ b/bottlecap/src/bin/bottlecap-test-mode/main.rs @@ -70,6 +70,7 @@ use dogstatsd::{ constants::CONTEXTS, flusher::Flusher as MetricsFlusher, metric::EMPTY_TAGS, }; use tokio::signal; +use tokio_util::sync::CancellationToken; use tracing::error; use tracing_subscriber::EnvFilter; use ustr::Ustr; @@ -150,9 +151,12 @@ async fn main() -> anyhow::Result<()> { )); let ingest_barrier = trace_agent.ingest_barrier(); + // Serializes every flush path. See [`drain_and_flush`]. + let flush_lock = Arc::new(tokio::sync::Mutex::new(())); let flush_extension = Arc::new(FlushRouterExtension { flushing_service: Arc::clone(&flushing_service), ingest_barrier: ingest_barrier.clone(), + flush_lock: Arc::clone(&flush_lock), }); let trace_agent = trace_agent.with_router_extension(flush_extension); // Errors are returned rather than logged so that a startup failure (port @@ -165,26 +169,13 @@ async fn main() -> anyhow::Result<()> { .map_err(|e| anyhow::anyhow!("trace agent failed: {e}")) }); - // Periodic flush driver. Decoupled from managed-instance mode: any non-Default - // strategy enables it. Manual flushing via POST /flush always works regardless. - if config.ext.serverless_flush_strategy != FlushStrategy::Default { - let mut interval = - FlushControl::new(config.ext.serverless_flush_strategy, config.flush_timeout) - .get_flush_interval(); - let fs = Arc::clone(&flushing_service); - let token = shutdown_token.clone(); - tokio::spawn(async move { - loop { - tokio::select! { - biased; - () = token.cancelled() => break, - // The periodic driver has no caller to report to; the - // flushing service already logs what it dropped. - _ = interval.tick() => { fs.flush_blocking().await; }, - } - } - }); - } + spawn_periodic_flush( + &config, + &flushing_service, + &ingest_barrier, + &flush_lock, + &shutdown_token, + ); // The listener finishing first means it never came up, or stopped serving // without being asked to. Either way there is nothing left to drain. @@ -212,13 +203,78 @@ async fn main() -> anyhow::Result<()> { SHUTDOWN_TIMEOUT.as_secs() ), } - // Handlers returning does not mean their payloads reached the - // aggregators; they may still be queued ahead of them. - ingest_barrier.wait().await; - flushing_service.flush_blocking_final().await; + // Taking the lock here is what makes an in-flight periodic flush finish + // before this one starts, so its queued payloads are not counted as + // already drained. + drain_and_flush(&flush_lock, &ingest_barrier, &flushing_service, true).await; Ok(()) } +/// Spawns the periodic flush driver. Decoupled from managed-instance mode: any +/// non-Default strategy enables it. Manual flushing via `POST /flush` always +/// works regardless. +fn spawn_periodic_flush( + config: &config::Config, + flushing_service: &Arc, + ingest_barrier: &IngestBarrier, + flush_lock: &Arc>, + shutdown_token: &CancellationToken, +) { + if config.ext.serverless_flush_strategy == FlushStrategy::Default { + return; + } + + let mut interval = + FlushControl::new(config.ext.serverless_flush_strategy, config.flush_timeout) + .get_flush_interval(); + let fs = Arc::clone(flushing_service); + let barrier = ingest_barrier.clone(); + let lock = Arc::clone(flush_lock); + let token = shutdown_token.clone(); + tokio::spawn(async move { + loop { + tokio::select! { + biased; + () = token.cancelled() => break, + // The periodic driver has no caller to report to; the + // flushing service already logs what it dropped. + _ = interval.tick() => { + drain_and_flush(&lock, &barrier, &fs, false).await; + }, + } + } + }); +} + +/// Runs the barrier-plus-flush sequence under `lock`. +/// +/// Every flusher drains its aggregator before awaiting network delivery, so +/// two overlapping flushes let one observe empty queues and report success +/// while the other's send is still in flight and may still fail. Holding the +/// lock across both steps keeps the periodic driver, `POST /flush`, and the +/// shutdown drain from overlapping. +/// +/// Returns `true` when a flusher dropped payloads it could not deliver. +async fn drain_and_flush( + lock: &tokio::sync::Mutex<()>, + barrier: &IngestBarrier, + flushing_service: &FlushingService, + is_final: bool, +) -> bool { + let _guard = lock.lock().await; + // A payload can be accepted, and its request answered, while it is still + // queued ahead of the aggregators. Drain those queues first so the flush + // is deterministic from the caller's point of view. Handlers returning + // during shutdown does not mean their payloads reached the aggregators + // either. + barrier.wait().await; + if is_final { + flushing_service.flush_blocking_final().await + } else { + flushing_service.flush_blocking().await + } +} + /// Path the config crate appends to `DD_APM_DD_URL` to build `apm_dd_url`. const TRACE_INTAKE_ROUTE: &str = "/api/v0.2/traces"; @@ -242,6 +298,7 @@ fn stats_url_from_trace_intake(apm_dd_url: &str) -> String { struct FlushRouterExtension { flushing_service: Arc, ingest_barrier: IngestBarrier, + flush_lock: Arc>, } /// Upper bound on a single `POST /flush`. The flushers already bound their own @@ -257,23 +314,20 @@ impl RouterExtension for FlushRouterExtension { fn extend(&self, router: Router) -> Result> { let fs = Arc::clone(&self.flushing_service); let barrier = self.ingest_barrier.clone(); + let lock = Arc::clone(&self.flush_lock); Ok(router.route( "/flush", post(move || { let fs = Arc::clone(&fs); let barrier = barrier.clone(); + let lock = Arc::clone(&lock); async move { // Isolate panics and bound execution time. flush_blocking_final // expects on the metrics aggregator handle, so a dead aggregator // task would otherwise panic the connection task instead of // returning a status the harness can act on. let mut task = tokio::task::spawn(async move { - // A payload can be accepted, and its request answered, - // while it is still queued ahead of the aggregators. - // Drain those queues first so this flush is - // deterministic from the caller's point of view. - barrier.wait().await; - fs.flush_blocking_final().await + drain_and_flush(&lock, &barrier, &fs, true).await }); match tokio::time::timeout(FLUSH_REQUEST_TIMEOUT, &mut task).await { Ok(Ok(false)) => StatusCode::NO_CONTENT, From 19cde53da0619e9f8e07471ffcd33b15bdc02730 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Thu, 10 Sep 2026 17:21:00 -0400 Subject: [PATCH 11/23] fix(test-mode): report a failed ingest barrier MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A closed barrier channel or a dropped acknowledgement means the forwarder task is gone and the payloads queued ahead of it were lost, not drained. The barrier now returns an error instead of treating that as success, and POST /flush reports 502 rather than 204. 🤖 --- bottlecap/src/bin/bottlecap-test-mode/main.rs | 22 ++++++++++------ bottlecap/src/traces/trace_agent.rs | 26 ++++++++++++------- 2 files changed, 30 insertions(+), 18 deletions(-) diff --git a/bottlecap/src/bin/bottlecap-test-mode/main.rs b/bottlecap/src/bin/bottlecap-test-mode/main.rs index 76437bcaa..cd7ba4820 100644 --- a/bottlecap/src/bin/bottlecap-test-mode/main.rs +++ b/bottlecap/src/bin/bottlecap-test-mode/main.rs @@ -254,7 +254,8 @@ fn spawn_periodic_flush( /// lock across both steps keeps the periodic driver, `POST /flush`, and the /// shutdown drain from overlapping. /// -/// Returns `true` when a flusher dropped payloads it could not deliver. +/// Returns `true` when payloads were lost: either the barrier reported a dead +/// forwarder, or a flusher dropped payloads it could not deliver. async fn drain_and_flush( lock: &tokio::sync::Mutex<()>, barrier: &IngestBarrier, @@ -267,12 +268,17 @@ async fn drain_and_flush( // is deterministic from the caller's point of view. Handlers returning // during shutdown does not mean their payloads reached the aggregators // either. - barrier.wait().await; - if is_final { + let barrier_failed = barrier.wait().await.is_err(); + if barrier_failed { + error!("Ingest barrier reported a dead forwarder; accepted payloads were lost"); + } + // Flush regardless: whatever did reach the aggregators should still go out. + let undelivered = if is_final { flushing_service.flush_blocking_final().await } else { flushing_service.flush_blocking().await - } + }; + barrier_failed || undelivered } /// Path the config crate appends to `DD_APM_DD_URL` to build `apm_dd_url`. @@ -331,11 +337,11 @@ impl RouterExtension for FlushRouterExtension { }); match tokio::time::timeout(FLUSH_REQUEST_TIMEOUT, &mut task).await { Ok(Ok(false)) => StatusCode::NO_CONTENT, - // The flush ran, but a flusher gave up on payloads it - // could not deliver and they were dropped. Reporting 204 - // here would tell the harness the drain succeeded. + // The flush ran, but payloads were lost on the way to + // the aggregators or to the intake. Reporting 204 here + // would tell the harness the drain succeeded. Ok(Ok(true)) => { - error!("Flush completed with undelivered payloads"); + error!("Flush completed with lost payloads"); StatusCode::BAD_GATEWAY } Ok(Err(e)) => { diff --git a/bottlecap/src/traces/trace_agent.rs b/bottlecap/src/traces/trace_agent.rs index 3d9c2024b..555f89a99 100644 --- a/bottlecap/src/traces/trace_agent.rs +++ b/bottlecap/src/traces/trace_agent.rs @@ -152,22 +152,28 @@ pub struct IngestBarrier { stats_tx: Sender>, } +/// A forwarder task ended without acknowledging the barrier, so payloads its +/// handlers had already accepted never reached the aggregator. +#[derive(Debug, Clone, Copy, thiserror::Error)] +#[error("ingest forwarder stopped before draining accepted payloads")] +pub struct IngestBarrierError; + impl IngestBarrier { - pub async fn wait(&self) { - Self::wait_for(&self.trace_tx).await; - Self::wait_for(&self.stats_tx).await; + pub async fn wait(&self) -> Result<(), IngestBarrierError> { + Self::wait_for(&self.trace_tx).await?; + Self::wait_for(&self.stats_tx).await } /// Each forwarder task selects on its payload channel before its barrier /// channel, so an acknowledgement can only be sent on an iteration where /// the payload channel was empty. - async fn wait_for(tx: &Sender>) { + async fn wait_for(tx: &Sender>) -> Result<(), IngestBarrierError> { let (ack_tx, ack_rx) = oneshot::channel(); - // A closed channel or a dropped acknowledgement means the forwarder - // task is gone, so there is nothing left to drain. - if tx.send(ack_tx).await.is_ok() { - let _ = ack_rx.await; - } + // Both a closed channel and a dropped acknowledgement mean the + // forwarder task is gone. Anything still queued ahead of it is lost, + // not drained, so this is a failure rather than an early success. + tx.send(ack_tx).await.map_err(|_| IngestBarrierError)?; + ack_rx.await.map_err(|_| IngestBarrierError) } } @@ -1130,7 +1136,7 @@ mod tests { trace_tx.send(stub_payload()).await.expect("send payload"); } - barrier.wait().await; + barrier.wait().await.expect("barrier"); let batches = aggregator_handle.get_batches().await.expect("get_batches"); let payloads: usize = batches.iter().map(Vec::len).sum(); From ab0075ca3e5f96fb97d15a371e056161cb5c581d Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Thu, 10 Sep 2026 17:23:50 -0400 Subject: [PATCH 12/23] docs(test-mode): note the inherited proxy routes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The endpoint table listed only the five core paths, but the binary also serves every proxy route the trace agent's router registers. 🤖 --- bottlecap/src/bin/bottlecap-test-mode/main.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/bottlecap/src/bin/bottlecap-test-mode/main.rs b/bottlecap/src/bin/bottlecap-test-mode/main.rs index cd7ba4820..68bbc5578 100644 --- a/bottlecap/src/bin/bottlecap-test-mode/main.rs +++ b/bottlecap/src/bin/bottlecap-test-mode/main.rs @@ -6,7 +6,7 @@ //! workflows that need to point a tracer at bottlecap without standing up a //! Lambda. //! -//! Endpoints exposed on `127.0.0.1:8126`: +//! Core endpoints on `127.0.0.1:8126`: //! //! | Path | Method | Source | //! |----------------|-----------|------------------------------------------| @@ -16,6 +16,11 @@ //! | `/info` | GET | trace agent | //! | `/flush` | POST | this binary's `FlushRouterExtension` | //! +//! The inherited `TraceAgent` router also serves its proxy routes (DSM, +//! profiling, LLM observability, debugger, diagnostics, and instrumentation +//! telemetry). They are not what this binary exists to exercise, but they +//! answer on the same port. `TraceAgent::make_router` has the full set. +//! //! Environment variables this binary reads: //! //! | Variable | Purpose | From 3d56a18a8225074cb14f6ce5bb52467ff3211428 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Thu, 10 Sep 2026 17:33:17 -0400 Subject: [PATCH 13/23] test(test-mode): cover the POST /flush status contract MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The handler now takes the flush work as an injectable operation, so the 204, 502, 500 and 504 branches can be driven directly. Without this a regression in the status the harness reads would pass CI. 🤖 --- bottlecap/src/bin/bottlecap-test-mode/main.rs | 99 +++++++++++++++---- 1 file changed, 82 insertions(+), 17 deletions(-) diff --git a/bottlecap/src/bin/bottlecap-test-mode/main.rs b/bottlecap/src/bin/bottlecap-test-mode/main.rs index 68bbc5578..612480c20 100644 --- a/bottlecap/src/bin/bottlecap-test-mode/main.rs +++ b/bottlecap/src/bin/bottlecap-test-mode/main.rs @@ -51,7 +51,7 @@ use tikv_jemallocator::Jemalloc; #[global_allocator] static GLOBAL: Jemalloc = Jemalloc; -use std::{collections::HashMap, env, path::Path, str::FromStr, sync::Arc, time::Duration}; +use std::{collections::HashMap, env, fmt, path::Path, str::FromStr, sync::Arc, time::Duration}; use axum::{Router, http::StatusCode, routing::post}; use bottlecap::{ @@ -74,6 +74,7 @@ use dogstatsd::{ aggregator::AggregatorService as MetricsAggregatorService, api_key::ApiKeyFactory, constants::CONTEXTS, flusher::Flusher as MetricsFlusher, metric::EMPTY_TAGS, }; +use futures::future::BoxFuture; use tokio::signal; use tokio_util::sync::CancellationToken; use tracing::error; @@ -159,9 +160,17 @@ async fn main() -> anyhow::Result<()> { // Serializes every flush path. See [`drain_and_flush`]. let flush_lock = Arc::new(tokio::sync::Mutex::new(())); let flush_extension = Arc::new(FlushRouterExtension { - flushing_service: Arc::clone(&flushing_service), - ingest_barrier: ingest_barrier.clone(), - flush_lock: Arc::clone(&flush_lock), + flush_op: { + let fs = Arc::clone(&flushing_service); + let barrier = ingest_barrier.clone(); + let lock = Arc::clone(&flush_lock); + Arc::new(move || { + let fs = Arc::clone(&fs); + let barrier = barrier.clone(); + let lock = Arc::clone(&lock); + Box::pin(async move { drain_and_flush(&lock, &barrier, &fs, true).await }) + }) + }, }); let trace_agent = trace_agent.with_router_extension(flush_extension); // Errors are returned rather than logged so that a startup failure (port @@ -305,11 +314,20 @@ fn stats_url_from_trace_intake(apm_dd_url: &str) -> String { ) } -#[derive(Debug)] +/// The work `POST /flush` performs, returning `true` when payloads were lost. +/// +/// Boxed rather than called directly so the handler's status branches can be +/// tested without standing up a flushing service. +type FlushOp = Arc BoxFuture<'static, bool> + Send + Sync>; + struct FlushRouterExtension { - flushing_service: Arc, - ingest_barrier: IngestBarrier, - flush_lock: Arc>, + flush_op: FlushOp, +} + +impl fmt::Debug for FlushRouterExtension { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("FlushRouterExtension").finish() + } } /// Upper bound on a single `POST /flush`. The flushers already bound their own @@ -323,23 +341,17 @@ const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(10); impl RouterExtension for FlushRouterExtension { fn extend(&self, router: Router) -> Result> { - let fs = Arc::clone(&self.flushing_service); - let barrier = self.ingest_barrier.clone(); - let lock = Arc::clone(&self.flush_lock); + let flush_op = Arc::clone(&self.flush_op); Ok(router.route( "/flush", post(move || { - let fs = Arc::clone(&fs); - let barrier = barrier.clone(); - let lock = Arc::clone(&lock); + let flush_op = Arc::clone(&flush_op); async move { // Isolate panics and bound execution time. flush_blocking_final // expects on the metrics aggregator handle, so a dead aggregator // task would otherwise panic the connection task instead of // returning a status the harness can act on. - let mut task = tokio::task::spawn(async move { - drain_and_flush(&lock, &barrier, &fs, true).await - }); + let mut task = tokio::task::spawn(async move { flush_op().await }); match tokio::time::timeout(FLUSH_REQUEST_TIMEOUT, &mut task).await { Ok(Ok(false)) => StatusCode::NO_CONTENT, // The flush ran, but payloads were lost on the way to @@ -407,6 +419,59 @@ fn enable_logging_subsystem() { mod tests { use super::*; + use axum::body::Body; + use axum::http::Request; + use tower::ServiceExt; + + /// Drives `POST /flush` against an extension backed by `op`. + async fn flush_status(op: FlushOp) -> StatusCode { + let router = FlushRouterExtension { flush_op: op } + .extend(Router::new()) + .expect("extend router"); + + router + .oneshot( + Request::builder() + .method("POST") + .uri("/flush") + .body(Body::empty()) + .expect("build request"), + ) + .await + .expect("route response") + .status() + } + + #[tokio::test] + async fn flush_returns_204_when_nothing_was_lost() { + let status = flush_status(Arc::new(|| Box::pin(async { false }))).await; + assert_eq!(status, StatusCode::NO_CONTENT); + } + + #[tokio::test] + async fn flush_returns_502_when_payloads_were_lost() { + let status = flush_status(Arc::new(|| Box::pin(async { true }))).await; + assert_eq!(status, StatusCode::BAD_GATEWAY); + } + + #[tokio::test] + async fn flush_returns_500_when_the_flush_panics() { + let status = flush_status(Arc::new(|| Box::pin(async { panic!("flush panicked") }))).await; + assert_eq!(status, StatusCode::INTERNAL_SERVER_ERROR); + } + + #[tokio::test(start_paused = true)] + async fn flush_returns_504_when_it_outruns_the_request_timeout() { + let status = flush_status(Arc::new(|| { + Box::pin(async { + tokio::time::sleep(FLUSH_REQUEST_TIMEOUT * 2).await; + false + }) + })) + .await; + assert_eq!(status, StatusCode::GATEWAY_TIMEOUT); + } + #[test] fn stats_url_follows_the_overridden_trace_intake() { assert_eq!( From 2e135e4063f294370d33f4e2e81ea4a3e12368a7 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Thu, 10 Sep 2026 19:11:56 -0400 Subject: [PATCH 14/23] fix(test-mode): stop reporting lost payloads on clean shutdown MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The ingest barrier held only its own acknowledgement channels, while every payload sender for the stats forwarder lived inside the trace agent. On the normal test-mode shutdown path the agent is dropped before the final drain runs, so the stats forwarder had already exited (after draining its queue) and the barrier reported it as dead, logging that accepted payloads were lost on every clean shutdown even though nothing was lost. The barrier now keeps a payload sender for each forwarder, so a live barrier keeps both forwarders running and a barrier error again means a forwarder genuinely died. The stats route is also no longer passed in to the router builder, so it cannot be wired to a channel outside the barrier's coverage. 🤖 --- bottlecap/src/traces/trace_agent.rs | 59 ++++++++++++++++++++++------- 1 file changed, 45 insertions(+), 14 deletions(-) diff --git a/bottlecap/src/traces/trace_agent.rs b/bottlecap/src/traces/trace_agent.rs index 555f89a99..29a8003e2 100644 --- a/bottlecap/src/traces/trace_agent.rs +++ b/bottlecap/src/traces/trace_agent.rs @@ -145,11 +145,21 @@ pub trait RouterExtension: Send + Sync { /// returns once every payload enqueued before the call has reached its /// aggregator, which makes a flush issued afterwards observe it. /// +/// A forwarder task exits once its payload channel closes, which happens as +/// soon as the last payload sender is dropped: the [`TraceAgent`] itself holds +/// two of them, so it drops them when its listener task finishes. The barrier +/// therefore keeps a payload sender of its own for each forwarder, so any live +/// `IngestBarrier` keeps both forwarders running and an `Err` from [`wait`] +/// means a forwarder actually died rather than exited after draining. +/// /// [`wait`]: IngestBarrier::wait #[derive(Clone, Debug)] pub struct IngestBarrier { trace_tx: Sender>, stats_tx: Sender>, + /// Keep-alives only; never sent on. See the note above. + _trace_forwarder_keepalive: Sender, + _stats_forwarder_keepalive: Sender, } /// A forwarder task ended without acknowledging the barrier, so payloads its @@ -282,12 +292,14 @@ impl TraceAgent { invocation_processor_handle, appsec_processor, tags_provider, - tx: trace_tx, - stats_tx, ingest_barrier: IngestBarrier { trace_tx: trace_barrier_tx, stats_tx: stats_barrier_tx, + _trace_forwarder_keepalive: trace_tx.clone(), + _stats_forwarder_keepalive: stats_tx.clone(), }, + tx: trace_tx, + stats_tx, shutdown_token: CancellationToken::new(), stats_concentrator, span_deduper, @@ -309,7 +321,7 @@ impl TraceAgent { pub async fn start(&self) -> Result<(), Box> { let now = Instant::now(); - let router = self.make_router(self.stats_tx.clone())?; + let router = self.make_router()?; let port = u16::try_from(TRACE_AGENT_PORT).expect("TRACE_AGENT_PORT is too large"); let socket = SocketAddr::from(([127, 0, 0, 1], port)); @@ -329,10 +341,12 @@ impl TraceAgent { Ok(()) } - fn make_router( - &self, - stats_tx: Sender, - ) -> Result> { + /// Builds the production router. The stats endpoint is always wired to + /// `self.stats_tx`, the channel [`IngestBarrier`] drains; routing it + /// anywhere else would silently take `/v0.6/stats` out of the barrier's + /// coverage. + fn make_router(&self) -> Result> { + let stats_tx = self.stats_tx.clone(); let stats_generator = Arc::new(StatsGenerator::new(self.stats_concentrator.clone())); let trace_state = TraceState { config: Arc::clone(&self.config), @@ -1143,14 +1157,34 @@ mod tests { assert_eq!(payloads, 3); } + /// Mirrors the test-mode shutdown drain: the listener task owns the + /// `TraceAgent`, so the agent (and every payload sender it holds) is gone + /// by the time the final barrier runs. That must not be reported as a dead + /// forwarder, because the forwarders drain their queues before exiting. + #[tokio::test] + async fn ingest_barrier_succeeds_after_the_agent_is_dropped() { + let (agent, aggregator_handle) = build_test_agent_with_aggregator(); + let trace_tx = agent.get_sender_copy(); + let barrier = agent.ingest_barrier(); + + trace_tx.send(stub_payload()).await.expect("send payload"); + drop(trace_tx); + drop(agent); + + barrier.wait().await.expect("barrier after agent drop"); + + let batches = aggregator_handle.get_batches().await.expect("get_batches"); + let payloads: usize = batches.iter().map(Vec::len).sum(); + assert_eq!(payloads, 1); + } + #[tokio::test] async fn with_router_extension_adds_reachable_route_to_make_router() { let hits = Arc::new(AtomicUsize::new(0)); let agent = build_test_agent().with_router_extension(Arc::new(SpyExtension { hits: Arc::clone(&hits), })); - let (stats_tx, _stats_rx) = mpsc::channel::(1); - let router = agent.make_router(stats_tx).expect("make_router"); + let router = agent.make_router().expect("make_router"); let response = router .oneshot( @@ -1170,10 +1204,8 @@ mod tests { #[tokio::test] async fn make_router_propagates_extension_error() { let agent = build_test_agent().with_router_extension(Arc::new(FailingExtension)); - let (stats_tx, _stats_rx) = mpsc::channel::(1); - let err = agent - .make_router(stats_tx) + .make_router() .expect_err("make_router should surface extension error"); assert!( @@ -1186,8 +1218,7 @@ mod tests { #[tokio::test] async fn make_router_returns_404_for_extension_route_when_none_attached() { let agent = build_test_agent(); - let (stats_tx, _stats_rx) = mpsc::channel::(1); - let router = agent.make_router(stats_tx).expect("make_router"); + let router = agent.make_router().expect("make_router"); let response = router .oneshot( From dfd070933b48d18ad9a74c7b787f455a5d6df146 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Thu, 10 Sep 2026 19:12:03 -0400 Subject: [PATCH 15/23] fix(test-mode): flush pending payloads when stopped with SIGTERM MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Only SIGINT was handled, so a container runtime or harness stopping the binary with SIGTERM killed it outright and the final drain never ran, losing the last accepted payloads. The signal driver is now pulled in by the test-mode feature rather than unconditionally, so the shipped extension no longer carries it. Also corrects the periodic flush comment: the `end` strategy yields the placeholder interval that means "never race a flush", so it is periodic only in name. 🤖 --- bottlecap/Cargo.toml | 8 +++-- bottlecap/src/bin/bottlecap-test-mode/main.rs | 30 +++++++++++++++++-- 2 files changed, 32 insertions(+), 6 deletions(-) diff --git a/bottlecap/Cargo.toml b/bottlecap/Cargo.toml index 2e20bfb24..69954ef14 100644 --- a/bottlecap/Cargo.toml +++ b/bottlecap/Cargo.toml @@ -35,7 +35,7 @@ flate2 = { version = "1.1", default-features = false, features = ["rust_backend" thiserror = { version = "1.0", default-features = false } # Transitive dependency (pulled in via cookie). Pinned to >=0.3.47 so cargo audit / CI passes (RUSTSEC-2026-0009). time = { version = "0.3.47", default-features = false } -tokio = { version = "1.47", default-features = false, features = ["macros", "rt-multi-thread", "signal", "time"] } +tokio = { version = "1.47", default-features = false, features = ["macros", "rt-multi-thread", "time"] } tokio-util = { version = "0.7", default-features = false } tracing = { version = "0.1", default-features = false } tracing-core = { version = "0.1", default-features = false } @@ -182,8 +182,10 @@ fips = [ # `InvocationProcessorHandle::noop()`) to callers that need to drive the # trace-processing surface without Lambda lifecycle state. # Not enabled in `default` or `fips`, so the items it gates do not appear -# in production builds. -test-mode = [] +# in production builds. `tokio/signal` is pulled in here rather than in the +# base dependency so the shipped extension does not carry the signal driver +# for a binary only the test-mode build produces. +test-mode = ["tokio/signal"] [lints.rust] unexpected_cfgs = { level = "warn", check-cfg = ['cfg(coverage,coverage_nightly)'] } diff --git a/bottlecap/src/bin/bottlecap-test-mode/main.rs b/bottlecap/src/bin/bottlecap-test-mode/main.rs index 612480c20..d459d1808 100644 --- a/bottlecap/src/bin/bottlecap-test-mode/main.rs +++ b/bottlecap/src/bin/bottlecap-test-mode/main.rs @@ -194,7 +194,7 @@ async fn main() -> anyhow::Result<()> { // The listener finishing first means it never came up, or stopped serving // without being asked to. Either way there is nothing left to drain. tokio::select! { - result = signal::ctrl_c() => result?, + result = shutdown_signal() => result?, result = &mut listener_task => match result? { Ok(()) => anyhow::bail!("trace agent listener stopped unexpectedly"), Err(e) => return Err(e), @@ -224,9 +224,33 @@ async fn main() -> anyhow::Result<()> { Ok(()) } +/// Resolves on the first signal that should end the process. +/// +/// SIGTERM matters as much as SIGINT here: the harness and any container +/// runtime stop the binary with SIGTERM, whose default disposition kills the +/// process outright, so without this the final drain never runs and the last +/// accepted payloads are lost. +#[cfg(unix)] +async fn shutdown_signal() -> anyhow::Result<()> { + let mut sigterm = signal::unix::signal(signal::unix::SignalKind::terminate())?; + tokio::select! { + result = signal::ctrl_c() => result?, + _ = sigterm.recv() => {} + } + Ok(()) +} + +#[cfg(not(unix))] +async fn shutdown_signal() -> anyhow::Result<()> { + signal::ctrl_c().await?; + Ok(()) +} + /// Spawns the periodic flush driver. Decoupled from managed-instance mode: any -/// non-Default strategy enables it. Manual flushing via `POST /flush` always -/// works regardless. +/// non-Default strategy enables it, using that strategy's interval. Note that +/// `end` yields the 15-minute placeholder interval `FlushControl` uses to mean +/// "never race a flush", so it is periodic only in name. Manual flushing via +/// `POST /flush` always works regardless. fn spawn_periodic_flush( config: &config::Config, flushing_service: &Arc, From 92618efb6646d79207a982f0ec5411d5f2493c0c Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Thu, 10 Sep 2026 19:12:04 -0400 Subject: [PATCH 16/23] refactor(traces): share the trace intake route constant MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The path the config crate appends to the trace intake URL was declared twice, once for data streams and once for the test-mode binary. Both now use a single definition. 🤖 --- bottlecap/src/bin/bottlecap-test-mode/main.rs | 5 +---- bottlecap/src/traces/data_streams/processor.rs | 4 +--- bottlecap/src/traces/mod.rs | 5 +++++ 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/bottlecap/src/bin/bottlecap-test-mode/main.rs b/bottlecap/src/bin/bottlecap-test-mode/main.rs index d459d1808..1e5810342 100644 --- a/bottlecap/src/bin/bottlecap-test-mode/main.rs +++ b/bottlecap/src/bin/bottlecap-test-mode/main.rs @@ -66,7 +66,7 @@ use bottlecap::{ startup::build_trace_agent, tags::{lambda::tags::FUNCTION_ARN_KEY, provider::Provider as TagProvider}, traces::{ - proxy_aggregator, + TRACE_INTAKE_ROUTE, proxy_aggregator, trace_agent::{IngestBarrier, RouterExtension}, }, }; @@ -319,9 +319,6 @@ async fn drain_and_flush( barrier_failed || undelivered } -/// Path the config crate appends to `DD_APM_DD_URL` to build `apm_dd_url`. -const TRACE_INTAKE_ROUTE: &str = "/api/v0.2/traces"; - /// Point stats at the same host as traces. /// /// `DD_APM_DD_URL` only moves the trace intake; stats would otherwise be diff --git a/bottlecap/src/traces/data_streams/processor.rs b/bottlecap/src/traces/data_streams/processor.rs index 7d31f6542..2d1454d9b 100644 --- a/bottlecap/src/traces/data_streams/processor.rs +++ b/bottlecap/src/traces/data_streams/processor.rs @@ -21,6 +21,7 @@ use reqwest::header::{CONTENT_ENCODING, CONTENT_TYPE, HeaderMap, HeaderValue}; use tokio::sync::Mutex as TokioMutex; use tracing::{debug, warn}; +use crate::traces::TRACE_INTAKE_ROUTE; use crate::traces::data_streams::aggregator::Aggregator; use crate::traces::data_streams::checkpoint::compute_consume_checkpoint; use crate::traces::data_streams::context::{ @@ -30,9 +31,6 @@ use crate::traces::proxy_aggregator::{Aggregator as ProxyAggregator, ProxyReques /// gzip level used by the tracer for pipeline stats. const GZIP_LEVEL: u32 = 1; -/// The trace intake path appended to `apm_dd_url` by the upstream config crate. -/// Must be stripped before deriving non-trace endpoints from that field. -const TRACE_INTAKE_ROUTE: &str = "/api/v0.2/traces"; pub struct DsmProcessor { service: String, diff --git a/bottlecap/src/traces/mod.rs b/bottlecap/src/traces/mod.rs index 0d9981805..cbcb7e439 100644 --- a/bottlecap/src/traces/mod.rs +++ b/bottlecap/src/traces/mod.rs @@ -38,6 +38,11 @@ const DNS_LOCAL_HOST_ADDRESS_URL_PREFIX: &str = "127.0.0.1"; // URL from the `_AWS_XRAY_DAEMON_ADDRESS` for DNS traces const AWS_XRAY_DAEMON_ADDRESS_URL_PREFIX: &str = "169.254.79.129"; +/// The trace intake path the config crate appends to `DD_APM_DD_URL` to build +/// `apm_dd_url`. Must be stripped from that field before deriving any +/// non-trace endpoint (stats, DSM) from it. +pub const TRACE_INTAKE_ROUTE: &str = "/api/v0.2/traces"; + // Name of the placeholder invocation span set by Java and Go tracers pub(crate) const INVOCATION_SPAN_RESOURCE: &str = "dd-tracer-serverless-span"; From c8cdbfeb2781d46c6fbc32e957ba66fc6297b01b Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 16 Sep 2026 15:16:30 -0400 Subject: [PATCH 17/23] Report lost metrics instead of aborting when aggregator handle fails --- bottlecap/src/bin/bottlecap-test-mode/main.rs | 7 ++-- bottlecap/src/flushing/service.rs | 36 +++++++++++++------ 2 files changed, 28 insertions(+), 15 deletions(-) diff --git a/bottlecap/src/bin/bottlecap-test-mode/main.rs b/bottlecap/src/bin/bottlecap-test-mode/main.rs index 1e5810342..81b3b9027 100644 --- a/bottlecap/src/bin/bottlecap-test-mode/main.rs +++ b/bottlecap/src/bin/bottlecap-test-mode/main.rs @@ -368,10 +368,9 @@ impl RouterExtension for FlushRouterExtension { post(move || { let flush_op = Arc::clone(&flush_op); async move { - // Isolate panics and bound execution time. flush_blocking_final - // expects on the metrics aggregator handle, so a dead aggregator - // task would otherwise panic the connection task instead of - // returning a status the harness can act on. + // Bound execution time. The flushers bound their own HTTP + // calls, but retries across the five flushers can stack, so + // this caps total wall-clock time for the harness. let mut task = tokio::task::spawn(async move { flush_op().await }); match tokio::time::timeout(FLUSH_REQUEST_TIMEOUT, &mut task).await { Ok(Ok(false)) => StatusCode::NO_CONTENT, diff --git a/bottlecap/src/flushing/service.rs b/bottlecap/src/flushing/service.rs index edd4ac3bf..f87a9334f 100644 --- a/bottlecap/src/flushing/service.rs +++ b/bottlecap/src/flushing/service.rs @@ -5,7 +5,8 @@ use std::sync::Arc; use tracing::{debug, error}; use dogstatsd::{ - aggregator::AggregatorHandle as MetricsAggregatorHandle, flusher::Flusher as MetricsFlusher, + aggregator::{AggregatorHandle as MetricsAggregatorHandle, FlushResponse}, + flusher::Flusher as MetricsFlusher, }; use crate::flushing::handles::{FlushHandles, MetricsRetryBatch}; @@ -331,11 +332,23 @@ impl FlushingService { /// redrive them, so the return value is the only signal that data was /// drained without reaching the intake. async fn flush_blocking_inner(&self, force_stats: bool) -> bool { - let flush_response = self - .metrics_aggr_handle - .flush() - .await - .expect("can't flush metrics aggr handle"); + // A failed handle means the aggregator task is gone and its buffered + // metrics are lost, so report them as undelivered and continue with an + // empty response rather than panicking: the release profile uses + // panic = "abort", so a panic here would take down the process. + let (flush_response, metrics_handle_failed) = match self.metrics_aggr_handle.flush().await { + Ok(response) => (response, false), + Err(e) => { + error!("FLUSHING_SERVICE | Metrics aggregator handle failed to flush: {e}"); + ( + FlushResponse { + series: Vec::new(), + distributions: Vec::new(), + }, + true, + ) + } + }; let metrics_futures: Vec<_> = self .metrics_flushers @@ -365,11 +378,12 @@ impl FlushingService { ("logs", !logs.is_empty()), ( "metrics", - metrics.iter().any(|retry| { - retry.as_ref().is_some_and(|(series, sketches)| { - !series.is_empty() || !sketches.is_empty() - }) - }), + metrics_handle_failed + || metrics.iter().any(|retry| { + retry.as_ref().is_some_and(|(series, sketches)| { + !series.is_empty() || !sketches.is_empty() + }) + }), ), ("traces", traces.is_some_and(|t| !t.is_empty())), ("stats", stats.is_some_and(|s| !s.is_empty())), From 1ecc697bb8096a58b738c4d6f8fb78f8ad8a4f6c Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 16 Sep 2026 15:18:04 -0400 Subject: [PATCH 18/23] Stop trace forwarder when aggregator is gone so flush reports loss --- bottlecap/src/traces/trace_agent.rs | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/bottlecap/src/traces/trace_agent.rs b/bottlecap/src/traces/trace_agent.rs index 29a8003e2..d543a86a3 100644 --- a/bottlecap/src/traces/trace_agent.rs +++ b/bottlecap/src/traces/trace_agent.rs @@ -244,8 +244,13 @@ impl TraceAgent { biased; tracer_payload_info = trace_rx.recv() => { let Some(tracer_payload_info) = tracer_payload_info else { break }; - if let Err(e) = aggregator_handle.insert_payload(tracer_payload_info) { - error!("TRACE_AGENT | Failed to insert payload into aggregator: {e}"); + // An error here means the aggregator task is gone and the + // payload is lost. Stop the forwarder so a pending barrier + // reports the loss via IngestBarrierError instead of + // acknowledging a flush that would silently drop payloads. + if aggregator_handle.insert_payload(tracer_payload_info).is_err() { + error!("TRACE_AGENT | Aggregator stopped, dropping trace forwarder"); + break; } } // Reached only on an iteration where the payload channel is From 447fcaa98229e8fc36958ef40dde6024c1a1047a Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 16 Sep 2026 15:23:51 -0400 Subject: [PATCH 19/23] Abort trace agent listener when shutdown wait times out --- bottlecap/src/bin/bottlecap-test-mode/main.rs | 35 ++++++++++++------- 1 file changed, 23 insertions(+), 12 deletions(-) diff --git a/bottlecap/src/bin/bottlecap-test-mode/main.rs b/bottlecap/src/bin/bottlecap-test-mode/main.rs index 81b3b9027..152cef0df 100644 --- a/bottlecap/src/bin/bottlecap-test-mode/main.rs +++ b/bottlecap/src/bin/bottlecap-test-mode/main.rs @@ -205,18 +205,7 @@ async fn main() -> anyhow::Result<()> { // in-flight /v0.4/traces requests through the aggregator before the flush // reads from it. shutdown_token.cancel(); - // Cancelling only signals the shutdown; awaiting the listener is what - // guarantees in-flight handlers have finished. Bounded so a lingering - // connection cannot wedge shutdown: draining late data beats hanging. - match tokio::time::timeout(SHUTDOWN_TIMEOUT, listener_task).await { - Ok(Ok(Ok(()))) => {} - Ok(Ok(Err(e))) => error!("Trace agent shut down with an error: {e:?}"), - Ok(Err(e)) => error!("Trace agent task failed: {e:?}"), - Err(_) => error!( - "Trace agent did not shut down within {}s, draining anyway", - SHUTDOWN_TIMEOUT.as_secs() - ), - } + await_listener_shutdown(&mut listener_task).await; // Taking the lock here is what makes an in-flight periodic flush finish // before this one starts, so its queued payloads are not counted as // already drained. @@ -224,6 +213,28 @@ async fn main() -> anyhow::Result<()> { Ok(()) } +/// Cancelling only signals the shutdown; awaiting the listener is what +/// guarantees in-flight handlers have finished. Bounded so a lingering +/// connection cannot wedge shutdown: draining late data beats hanging. +async fn await_listener_shutdown(listener_task: &mut tokio::task::JoinHandle>) { + match tokio::time::timeout(SHUTDOWN_TIMEOUT, &mut *listener_task).await { + Ok(Ok(Ok(()))) => {} + Ok(Ok(Err(e))) => error!("Trace agent shut down with an error: {e:?}"), + Ok(Err(e)) => error!("Trace agent task failed: {e:?}"), + Err(_) => { + // Dropping the handle would detach the listener and let an in-flight + // flush keep holding the flush lock past this bound. Aborting and + // awaiting stops it before the drain takes the lock. + listener_task.abort(); + let _ = listener_task.await; + error!( + "Trace agent did not shut down within {}s, aborting and draining anyway", + SHUTDOWN_TIMEOUT.as_secs() + ); + } + } +} + /// Resolves on the first signal that should end the process. /// /// SIGTERM matters as much as SIGINT here: the harness and any container From 42edd75d78246b6d2d755bc0e3820b7e836b1c76 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Fri, 18 Sep 2026 00:19:24 -0400 Subject: [PATCH 20/23] Cancel test-mode flushes on shutdown timeout --- bottlecap/src/bin/bottlecap-test-mode/main.rs | 149 ++++++++++++++++-- 1 file changed, 134 insertions(+), 15 deletions(-) diff --git a/bottlecap/src/bin/bottlecap-test-mode/main.rs b/bottlecap/src/bin/bottlecap-test-mode/main.rs index 152cef0df..24c6f9b5e 100644 --- a/bottlecap/src/bin/bottlecap-test-mode/main.rs +++ b/bottlecap/src/bin/bottlecap-test-mode/main.rs @@ -159,7 +159,9 @@ async fn main() -> anyhow::Result<()> { let ingest_barrier = trace_agent.ingest_barrier(); // Serializes every flush path. See [`drain_and_flush`]. let flush_lock = Arc::new(tokio::sync::Mutex::new(())); + let flush_cancel_token = CancellationToken::new(); let flush_extension = Arc::new(FlushRouterExtension { + cancellation_token: flush_cancel_token.clone(), flush_op: { let fs = Arc::clone(&flushing_service); let barrier = ingest_barrier.clone(); @@ -205,7 +207,7 @@ async fn main() -> anyhow::Result<()> { // in-flight /v0.4/traces requests through the aggregator before the flush // reads from it. shutdown_token.cancel(); - await_listener_shutdown(&mut listener_task).await; + await_listener_shutdown(&mut listener_task, &flush_cancel_token).await; // Taking the lock here is what makes an in-flight periodic flush finish // before this one starts, so its queued payloads are not counted as // already drained. @@ -216,15 +218,18 @@ async fn main() -> anyhow::Result<()> { /// Cancelling only signals the shutdown; awaiting the listener is what /// guarantees in-flight handlers have finished. Bounded so a lingering /// connection cannot wedge shutdown: draining late data beats hanging. -async fn await_listener_shutdown(listener_task: &mut tokio::task::JoinHandle>) { +async fn await_listener_shutdown( + listener_task: &mut tokio::task::JoinHandle>, + flush_cancel_token: &CancellationToken, +) { match tokio::time::timeout(SHUTDOWN_TIMEOUT, &mut *listener_task).await { Ok(Ok(Ok(()))) => {} Ok(Ok(Err(e))) => error!("Trace agent shut down with an error: {e:?}"), Ok(Err(e)) => error!("Trace agent task failed: {e:?}"), Err(_) => { - // Dropping the handle would detach the listener and let an in-flight - // flush keep holding the flush lock past this bound. Aborting and - // awaiting stops it before the drain takes the lock. + // Axum connection tasks can outlive the listener. Cancel their + // flush work explicitly so it releases the lock before the drain. + flush_cancel_token.cancel(); listener_task.abort(); let _ = listener_task.await; error!( @@ -354,6 +359,7 @@ type FlushOp = Arc BoxFuture<'static, bool> + Send + Sync>; struct FlushRouterExtension { flush_op: FlushOp, + cancellation_token: CancellationToken, } impl fmt::Debug for FlushRouterExtension { @@ -374,30 +380,42 @@ const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(10); impl RouterExtension for FlushRouterExtension { fn extend(&self, router: Router) -> Result> { let flush_op = Arc::clone(&self.flush_op); + let cancellation_token = self.cancellation_token.clone(); Ok(router.route( "/flush", post(move || { let flush_op = Arc::clone(&flush_op); + let cancellation_token = cancellation_token.clone(); async move { // Bound execution time. The flushers bound their own HTTP // calls, but retries across the five flushers can stack, so // this caps total wall-clock time for the harness. - let mut task = tokio::task::spawn(async move { flush_op().await }); - match tokio::time::timeout(FLUSH_REQUEST_TIMEOUT, &mut task).await { - Ok(Ok(false)) => StatusCode::NO_CONTENT, + let mut tasks = tokio::task::JoinSet::new(); + tasks.spawn(async move { + tokio::select! { + biased; + () = cancellation_token.cancelled() => { + error!("Flush cancelled after shutdown grace period"); + true + } + result = async { flush_op().await } => result, + } + }); + match tokio::time::timeout(FLUSH_REQUEST_TIMEOUT, tasks.join_next()).await { + Ok(Some(Ok(false))) => StatusCode::NO_CONTENT, // The flush ran, but payloads were lost on the way to // the aggregators or to the intake. Reporting 204 here // would tell the harness the drain succeeded. - Ok(Ok(true)) => { + Ok(Some(Ok(true))) => { error!("Flush completed with lost payloads"); StatusCode::BAD_GATEWAY } - Ok(Err(e)) => { - error!("Flush task failed: {e:?}"); + Ok(result) => { + error!("Flush task failed: {result:?}"); StatusCode::INTERNAL_SERVER_ERROR } Err(_) => { - task.abort(); + tasks.shutdown().await; error!( "Flush timed out after {}s, aborting", FLUSH_REQUEST_TIMEOUT.as_secs() @@ -456,9 +474,19 @@ mod tests { /// Drives `POST /flush` against an extension backed by `op`. async fn flush_status(op: FlushOp) -> StatusCode { - let router = FlushRouterExtension { flush_op: op } - .extend(Router::new()) - .expect("extend router"); + flush_status_with_cancellation(op, CancellationToken::new()).await + } + + async fn flush_status_with_cancellation( + op: FlushOp, + cancellation_token: CancellationToken, + ) -> StatusCode { + let router = FlushRouterExtension { + flush_op: op, + cancellation_token, + } + .extend(Router::new()) + .expect("extend router"); router .oneshot( @@ -473,6 +501,97 @@ mod tests { .status() } + fn blocked_flush( + lock: Arc>, + started: Arc, + ) -> FlushOp { + Arc::new(move || { + let lock = Arc::clone(&lock); + let started = Arc::clone(&started); + Box::pin(async move { + let _guard = lock.lock().await; + started.notify_one(); + std::future::pending().await + }) + }) + } + + #[tokio::test(start_paused = true)] + async fn flush_releases_lock_when_handler_is_cancelled() { + let lock = Arc::new(tokio::sync::Mutex::new(())); + let started = Arc::new(tokio::sync::Notify::new()); + let handler = tokio::spawn(flush_status(blocked_flush( + Arc::clone(&lock), + Arc::clone(&started), + ))); + started.notified().await; + assert!(lock.try_lock().is_err()); + + handler.abort(); + assert!( + handler + .await + .expect_err("handler was aborted") + .is_cancelled() + ); + + let _guard = tokio::time::timeout(Duration::from_secs(1), lock.lock()) + .await + .expect("cancelled handler must release the flush lock"); + } + + #[tokio::test(start_paused = true)] + async fn shutdown_timeout_cancels_flush_in_an_independent_handler() { + let cancellation_token = CancellationToken::new(); + let lock = Arc::new(tokio::sync::Mutex::new(())); + let started = Arc::new(tokio::sync::Notify::new()); + let handler = tokio::spawn(flush_status_with_cancellation( + blocked_flush(Arc::clone(&lock), Arc::clone(&started)), + cancellation_token.clone(), + )); + started.notified().await; + assert!(lock.try_lock().is_err()); + + // Like an Axum connection task, the handler is not owned by the listener. + let mut listener = tokio::spawn(std::future::pending()); + await_listener_shutdown(&mut listener, &cancellation_token).await; + + assert!(cancellation_token.is_cancelled()); + assert!(listener.is_finished()); + assert_eq!( + tokio::time::timeout(Duration::from_secs(1), handler) + .await + .expect("shutdown must cancel the active flush") + .expect("handler completed"), + StatusCode::BAD_GATEWAY + ); + assert!(lock.try_lock().is_ok()); + } + + #[tokio::test] + async fn flush_does_not_start_after_shutdown_cancellation() { + let cancellation_token = CancellationToken::new(); + cancellation_token.cancel(); + let status = flush_status_with_cancellation( + Arc::new(|| panic!("flush must not start after cancellation")), + cancellation_token, + ) + .await; + assert_eq!(status, StatusCode::BAD_GATEWAY); + } + + #[tokio::test(start_paused = true)] + async fn flush_timeout_releases_lock_before_responding() { + let lock = Arc::new(tokio::sync::Mutex::new(())); + let status = flush_status(blocked_flush( + Arc::clone(&lock), + Arc::new(tokio::sync::Notify::new()), + )) + .await; + assert_eq!(status, StatusCode::GATEWAY_TIMEOUT); + assert!(lock.try_lock().is_ok()); + } + #[tokio::test] async fn flush_returns_204_when_nothing_was_lost() { let status = flush_status(Arc::new(|| Box::pin(async { false }))).await; From e2e8731beadf84a4730565aafd511cb25d5a5744 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Fri, 18 Sep 2026 00:23:01 -0400 Subject: [PATCH 21/23] Delay the first test-mode periodic flush --- bottlecap/src/bin/bottlecap-test-mode/main.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/bottlecap/src/bin/bottlecap-test-mode/main.rs b/bottlecap/src/bin/bottlecap-test-mode/main.rs index 24c6f9b5e..96698ec61 100644 --- a/bottlecap/src/bin/bottlecap-test-mode/main.rs +++ b/bottlecap/src/bin/bottlecap-test-mode/main.rs @@ -286,6 +286,7 @@ fn spawn_periodic_flush( let lock = Arc::clone(flush_lock); let token = shutdown_token.clone(); tokio::spawn(async move { + interval.tick().await; // discard the immediate first tick loop { tokio::select! { biased; From 80f16490947ed23159dc333ddd1f14764d2b17bd Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Fri, 18 Sep 2026 00:28:07 -0400 Subject: [PATCH 22/23] Run test-mode endpoint tests in CI --- .github/workflows/rs_ci.yml | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/.github/workflows/rs_ci.yml b/.github/workflows/rs_ci.yml index a5572db18..1dcb41801 100644 --- a/.github/workflows/rs_ci.yml +++ b/.github/workflows/rs_ci.yml @@ -82,8 +82,7 @@ jobs: run: cargo clippy --workspace --all-targets --features default - working-directory: bottlecap run: cargo clippy --workspace --all-targets --no-default-features --features fips - # The test-mode feature gates the bottlecap-test-mode binary via - # required-features, so no other job compiles it. + # Also lint the feature-gated test-mode binary. - working-directory: bottlecap run: cargo clippy --workspace --all-targets --features default,test-mode @@ -145,6 +144,9 @@ jobs: - uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad # v0.0.9 - working-directory: bottlecap run: cargo nextest run --workspace + - name: Test the test-mode binary + working-directory: bottlecap + run: cargo nextest run --features test-mode --bin bottlecap-test-mode format: name: Format From 4013d692d1305635fad621c90abd6bdcc5879d6c Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Mon, 21 Sep 2026 22:52:23 -0400 Subject: [PATCH 23/23] Allow configuring the trace agent listener port The test-mode binary now reads DD_APM_RECEIVER_PORT (the same variable the Go trace agent uses) and binds its listener to that port instead of the default 8126. Unset or empty keeps the default; invalid values fail startup rather than silently listening on the wrong port. Production Lambda behavior is unchanged. --- bottlecap/src/bin/bottlecap-test-mode/main.rs | 69 ++++++++++++++++++- bottlecap/src/traces/trace_agent.rs | 22 +++++- 2 files changed, 85 insertions(+), 6 deletions(-) diff --git a/bottlecap/src/bin/bottlecap-test-mode/main.rs b/bottlecap/src/bin/bottlecap-test-mode/main.rs index 96698ec61..eda6dc846 100644 --- a/bottlecap/src/bin/bottlecap-test-mode/main.rs +++ b/bottlecap/src/bin/bottlecap-test-mode/main.rs @@ -6,7 +6,8 @@ //! workflows that need to point a tracer at bottlecap without standing up a //! Lambda. //! -//! Core endpoints on `127.0.0.1:8126`: +//! Core endpoints on `127.0.0.1:8126` (port configurable, see +//! `DD_APM_RECEIVER_PORT` below): //! //! | Path | Method | Source | //! |----------------|-----------|------------------------------------------| @@ -27,6 +28,7 @@ //! |---------------------------------|-------------------------------------------------------------------------| //! | `DD_APM_DD_URL` | Override trace intake URL; stats follow it (harness points at fake-intake) | //! | `DD_SITE` | Derive trace and stats intake URLs when `DD_APM_DD_URL` is unset | +//! | `DD_APM_RECEIVER_PORT` | TCP port the trace agent listener binds to (default 8126) | //! | `DD_SERVERLESS_FLUSH_STRATEGY` | Enable periodic flushing (e.g. `periodically,5000`); default = manual | //! | `DD_TESTMODE_FUNCTION_ARN` | Override stub function ARN for tag generation | //! | `DD_LOG_LEVEL` | Logging verbosity, parsed by [`bottlecap::config::log_level::LogLevel`] | @@ -67,7 +69,7 @@ use bottlecap::{ tags::{lambda::tags::FUNCTION_ARN_KEY, provider::Provider as TagProvider}, traces::{ TRACE_INTAKE_ROUTE, proxy_aggregator, - trace_agent::{IngestBarrier, RouterExtension}, + trace_agent::{TRACE_AGENT_PORT, IngestBarrier, RouterExtension}, }, }; use dogstatsd::{ @@ -86,6 +88,10 @@ async fn main() -> anyhow::Result<()> { init_ustr(); enable_logging_subsystem(); + // Fail before binding anything so a misconfigured harness errors out + // instead of listening on the wrong port. + let receiver_port = receiver_port_from_env()?; + // Outside Lambda every `AwsConfig` field falls back via `unwrap_or_default()`, // so loading from env is safe with no AWS env vars set. The struct is only // read by the secrets resolver, which test-mode bypasses. @@ -174,7 +180,9 @@ async fn main() -> anyhow::Result<()> { }) }, }); - let trace_agent = trace_agent.with_router_extension(flush_extension); + let trace_agent = trace_agent + .with_receiver_port(receiver_port) + .with_router_extension(flush_extension); // Errors are returned rather than logged so that a startup failure (port // 8126 already bound, for instance) ends the process instead of leaving a // live one with no listener for the harness to connect to. @@ -336,6 +344,33 @@ async fn drain_and_flush( barrier_failed || undelivered } +/// The env var the Go trace agent uses for its listener port. +const ENV_RECEIVER_PORT: &str = "DD_APM_RECEIVER_PORT"; + +/// Resolves the trace agent listener port from [`ENV_RECEIVER_PORT`]. +/// +/// Unset or empty means the default [`TRACE_AGENT_PORT`] (8126): empty is +/// treated as unset because the Go trace agent makes the same choice, and +/// harnesses often export variables with blank values. Any other value must +/// be a port in 1-65535 or startup fails, rather than silently falling back +/// to a port the tracer is not pointed at. +fn receiver_port_from_env() -> anyhow::Result { + receiver_port(env::var(ENV_RECEIVER_PORT).ok().as_deref()) +} + +/// Parses the receiver port from an already-extracted variable value. +/// Split from [`receiver_port_from_env`] so tests need no process env. +fn receiver_port(value: Option<&str>) -> anyhow::Result { + let Some(value) = value.filter(|v| !v.is_empty()) else { + return Ok(u16::try_from(TRACE_AGENT_PORT).expect("default trace agent port fits in u16")); + }; + value + .parse::() + .ok() + .filter(|port| *port != 0) + .ok_or_else(|| anyhow::anyhow!("{ENV_RECEIVER_PORT} must be a port in 1-65535, got {value:?}")) +} + /// Point stats at the same host as traces. /// /// `DD_APM_DD_URL` only moves the trace intake; stats would otherwise be @@ -639,4 +674,32 @@ mod tests { libdd_trace_utils::config_utils::trace_stats_url(site) ); } + + #[test] + fn receiver_port_parses_valid_values_and_defaults_on_unset_or_empty() { + let default = u16::try_from(TRACE_AGENT_PORT).expect("default port fits in u16"); + let test_cases = [ + (None, Ok(default)), + (Some(""), Ok(default)), + (Some("8126"), Ok(8126)), + (Some("9999"), Ok(9999)), + (Some("65535"), Ok(65535)), + (Some("abc"), Err(())), + (Some("0"), Err(())), + (Some("70000"), Err(())), + (Some("-1"), Err(())), + (Some("1.5"), Err(())), + ]; + for (input, expected) in test_cases { + let result = receiver_port(input).map_err(|_| ()); + assert_eq!(result, expected, "input: {input:?}"); + } + } + + #[test] + fn receiver_port_error_names_the_env_var_and_value() { + let err = receiver_port(Some("0")).expect_err("0 is not a valid port"); + assert!(err.to_string().contains("DD_APM_RECEIVER_PORT")); + assert!(err.to_string().contains("\"0\"")); + } } diff --git a/bottlecap/src/traces/trace_agent.rs b/bottlecap/src/traces/trace_agent.rs index d543a86a3..576efc8f3 100644 --- a/bottlecap/src/traces/trace_agent.rs +++ b/bottlecap/src/traces/trace_agent.rs @@ -50,7 +50,8 @@ use libdd_trace_utils::trace_utils::{self}; use crate::traces::stats_concentrator_service::StatsConcentratorHandle; -const TRACE_AGENT_PORT: usize = 8126; +/// Default TCP port the trace agent listener binds to on 127.0.0.1. +pub const TRACE_AGENT_PORT: usize = 8126; // Agent endpoints const V4_TRACE_ENDPOINT_PATH: &str = "/v0.4/traces"; @@ -205,6 +206,9 @@ pub struct TraceAgent { /// `None` when the caller wants no extra routes. See /// [`TraceAgent::with_router_extension`]. router_extension: Option>, + /// `None` when the caller wants the default port. See + /// [`TraceAgent::with_receiver_port`]. + receiver_port: Option, } #[derive(Clone, Copy)] @@ -309,6 +313,7 @@ impl TraceAgent { stats_concentrator, span_deduper, router_extension: None, + receiver_port: None, } } @@ -322,17 +327,28 @@ impl TraceAgent { self } + /// Overrides the listener port, which defaults to [`TRACE_AGENT_PORT`] + /// (8126). For the test-mode binary, which maps `DD_APM_RECEIVER_PORT` + /// to this setting; production callers must keep the default. + #[must_use] + pub fn with_receiver_port(mut self, port: u16) -> Self { + self.receiver_port = Some(port); + self + } + #[allow(clippy::cast_possible_truncation)] pub async fn start(&self) -> Result<(), Box> { let now = Instant::now(); let router = self.make_router()?; - let port = u16::try_from(TRACE_AGENT_PORT).expect("TRACE_AGENT_PORT is too large"); + let port = self + .receiver_port + .unwrap_or(u16::try_from(TRACE_AGENT_PORT).expect("TRACE_AGENT_PORT is too large")); let socket = SocketAddr::from(([127, 0, 0, 1], port)); let listener = tokio::net::TcpListener::bind(&socket).await?; - debug!("TRACE AGENT | Listening on port {TRACE_AGENT_PORT}"); + debug!("TRACE AGENT | Listening on port {port}"); debug!( "TRACE AGENT | Time taken to start: {} ms", now.elapsed().as_millis()