diff --git a/.github/workflows/rs_ci.yml b/.github/workflows/rs_ci.yml index 0518386dc..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 - # No other job compiles the test-mode feature: it gates test-only - # constructors that are absent from default and fips builds. + # 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 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 }} diff --git a/bottlecap/Cargo.toml b/bottlecap/Cargo.toml index f14e48730..69954ef14 100644 --- a/bottlecap/Cargo.toml +++ b/bottlecap/Cargo.toml @@ -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 @@ -178,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 new file mode 100644 index 000000000..96698ec61 --- /dev/null +++ b/bottlecap/src/bin/bottlecap-test-mode/main.rs @@ -0,0 +1,642 @@ +//! 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. +//! +//! Core endpoints 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` | +//! +//! 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 | +//! |---------------------------------|-------------------------------------------------------------------------| +//! | `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`] | +//! +//! [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, fmt, path::Path, str::FromStr, sync::Arc, time::Duration}; + +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::{ + TRACE_INTAKE_ROUTE, proxy_aggregator, + trace_agent::{IngestBarrier, RouterExtension}, + }, +}; +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; +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), + 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); + 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 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(); + 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 + // 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}")) + }); + + 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. + tokio::select! { + result = shutdown_signal() => 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. + shutdown_token.cancel(); + 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. + drain_and_flush(&flush_lock, &ingest_barrier, &flushing_service, true).await; + 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>, + 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(_) => { + // 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!( + "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 +/// 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, 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, + 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 { + interval.tick().await; // discard the immediate first tick + 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 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, + 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. + 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 +} + +/// 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), + ) +} + +/// 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 { + flush_op: FlushOp, + cancellation_token: CancellationToken, +} + +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 +/// 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); + +/// 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 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 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(Some(Ok(true))) => { + error!("Flush completed with lost payloads"); + StatusCode::BAD_GATEWAY + } + Ok(result) => { + error!("Flush task failed: {result:?}"); + StatusCode::INTERNAL_SERVER_ERROR + } + Err(_) => { + tasks.shutdown().await; + error!( + "Flush timed out after {}s, aborting", + FLUSH_REQUEST_TIMEOUT.as_secs() + ); + StatusCode::GATEWAY_TIMEOUT + } + } + } + }), + )) + } +} + +// 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"); +} + +#[cfg(test)] +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 { + 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( + Request::builder() + .method("POST") + .uri("/flush") + .body(Body::empty()) + .expect("build request"), + ) + .await + .expect("route response") + .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; + 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!( + 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/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/flushing/service.rs b/bottlecap/src/flushing/service.rs index a8181d7e9..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}; @@ -300,8 +301,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,19 +315,40 @@ 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) { - let flush_response = self - .metrics_aggr_handle - .flush() - .await - .expect("can't flush metrics aggr handle"); + /// + /// 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 { + // 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 @@ -341,13 +366,43 @@ 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_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())), + ("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() } } 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..b1ff66d16 --- /dev/null +++ b/bottlecap/src/startup.rs @@ -0,0 +1,237 @@ +// 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 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 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 +/// 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, + tags_provider: &Arc, + invocation_processor_handle: InvocationProcessorHandle, + 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(). + 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(), + stats_url_override + .unwrap_or_else(|| 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, + None, + ); + + // 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/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"; diff --git a/bottlecap/src/traces/trace_agent.rs b/bottlecap/src/traces/trace_agent.rs index aea1d084d..d543a86a3 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; @@ -112,10 +114,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 @@ -134,6 +137,56 @@ 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. +/// +/// 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 +/// 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) -> 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>) -> Result<(), IngestBarrierError> { + let (ack_tx, ack_rx) = oneshot::channel(); + // 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) + } +} + pub struct TraceAgent { pub config: Arc, pub trace_processor: Arc, @@ -145,6 +198,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 @@ -179,12 +234,56 @@ 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 }; + // 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 + // 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(()); + } } } }); @@ -198,7 +297,14 @@ impl TraceAgent { invocation_processor_handle, appsec_processor, tags_provider, + 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, @@ -220,22 +326,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()?; let port = u16::try_from(TRACE_AGENT_PORT).expect("TRACE_AGENT_PORT is too large"); let socket = SocketAddr::from(([127, 0, 0, 1], port)); @@ -255,10 +346,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), @@ -784,6 +877,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) { @@ -853,13 +953,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; @@ -981,6 +1085,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()); @@ -1002,9 +1110,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 {}), @@ -1014,7 +1122,65 @@ 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.expect("barrier"); + + let batches = aggregator_handle.get_batches().await.expect("get_batches"); + let payloads: usize = batches.iter().map(Vec::len).sum(); + 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] @@ -1023,8 +1189,7 @@ mod tests { 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( @@ -1044,10 +1209,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!( @@ -1060,8 +1223,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(