diff --git a/bottlecap/src/bin/bottlecap/main.rs b/bottlecap/src/bin/bottlecap/main.rs index 6cb53476a..e000b7a34 100644 --- a/bottlecap/src/bin/bottlecap/main.rs +++ b/bottlecap/src/bin/bottlecap/main.rs @@ -141,6 +141,13 @@ async fn main() -> anyhow::Result<()> { .map_err(|e| anyhow::anyhow!("Failed to create client: {e:?}"))?; cold_start_durations.create_extensions_api_client_ms = elapsed_ms(step_start_time); + // Lambda starts the runtime once every extension has registered, and + // tracers probe `/info` early in their init. Bind the trace agent's port + // before registering, so that probe queues instead of being refused. + let trace_listener = bottlecap::traces::trace_agent::TraceAgent::bind_listener(None) + .inspect_err(|e| error!("Error starting trace agent: {e:?}")) + .ok(); + let cloned_client = client.clone(); let runtime_api = aws_config.runtime_api.clone(); let managed_instance_mode = aws_config.is_managed_instance_mode(); @@ -193,6 +200,7 @@ async fn main() -> anyhow::Result<()> { &r, Arc::clone(&api_key_factory), cold_start_durations, + trace_listener, ) .await { @@ -326,6 +334,7 @@ async fn extension_loop_idle( } #[allow(clippy::too_many_lines)] +#[allow(clippy::too_many_arguments)] async fn extension_loop_active( aws_config: Arc, config: &Arc, @@ -334,6 +343,7 @@ async fn extension_loop_active( r: &RegisterResponse, api_key_factory: Arc, mut cold_start_durations: ColdStartDurations, + trace_listener: Option, ) -> anyhow::Result<()> { let (mut event_bus, event_bus_tx) = EventBus::run(); @@ -462,6 +472,7 @@ async fn extension_loop_active( appsec_processor.clone(), &shared_client, Arc::clone(&proxy_aggregator), + trace_listener, ); let api_runtime_proxy_shutdown_signal = start_api_runtime_proxy( diff --git a/bottlecap/src/startup.rs b/bottlecap/src/startup.rs index cf12b0939..8a19072c5 100644 --- a/bottlecap/src/startup.rs +++ b/bottlecap/src/startup.rs @@ -223,15 +223,19 @@ pub fn build_trace_agent( } /// 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 +/// the [`trace_agent::TraceAgent`] HTTP server on `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; 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. +/// `listener` comes from [`trace_agent::TraceAgent::bind_listener`], called +/// before the extension registers. `None` means that bind failed (the caller +/// logged it), and no trace agent is served. +/// +/// Errors from `TraceAgent::start_with_listener` (router-extension failures, +/// axum serve errors) are logged and discarded; 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. +#[allow(clippy::too_many_arguments)] pub fn start_trace_agent( config: &Arc, api_key_factory: &Arc, @@ -240,6 +244,7 @@ pub fn start_trace_agent( appsec_processor: Option>>, client: &reqwest::Client, proxy_aggregator: Arc>, + listener: Option, ) -> TraceAgentPipeline { let (trace_agent, pipeline) = build_trace_agent( config, @@ -252,11 +257,13 @@ pub fn start_trace_agent( None, ); - tokio::spawn(async move { - if let Err(e) = trace_agent.start().await { - error!("Error starting trace agent: {e:?}"); - } - }); + if let Some(listener) = listener { + tokio::spawn(async move { + if let Err(e) = trace_agent.start_with_listener(listener).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 bb7bdf6cf..4bf906a0f 100644 --- a/bottlecap/src/traces/trace_agent.rs +++ b/bottlecap/src/traces/trace_agent.rs @@ -53,6 +53,8 @@ use crate::traces::stats_concentrator_service::StatsConcentratorHandle; /// Default TCP port the trace agent listener binds to on 127.0.0.1. pub const TRACE_AGENT_PORT: usize = 8126; +/// Listen backlog; the kernel caps it at `net.core.somaxconn`. +const LISTEN_BACKLOG: u32 = 1024; // Agent endpoints const V4_TRACE_ENDPOINT_PATH: &str = "/v0.4/traces"; @@ -64,6 +66,11 @@ const LLM_OBS_EVAL_METRIC_ENDPOINT_PATH: &str = "/evp_proxy/v2/api/intake/llm-ob const LLM_OBS_EVAL_METRIC_ENDPOINT_PATH_V2: &str = "/evp_proxy/v2/api/intake/llm-obs/v2/eval-metric"; const LLM_OBS_SPANS_ENDPOINT_PATH: &str = "/evp_proxy/v2/api/v2/llmobs"; +// Agentless flag discovery and Java/Go LLM Obs find the EVP proxy by this base path +// in /info; dd-trace-java and dd-trace-go LLM Obs also need the trailing slash. +const EVP_PROXY_V2_BASE_PATH: &str = "/evp_proxy/v2/"; +const FFE_FLAG_EVALUATION_ENDPOINT_PATH: &str = "/evp_proxy/v2/api/v2/flagevaluation"; +const FFE_EXPOSURES_ENDPOINT_PATH: &str = "/evp_proxy/v2/api/v2/exposures"; const INFO_ENDPOINT_PATH: &str = "/info"; const V1_DEBUGGER_ENDPOINT_PATH: &str = "/debugger/v1/input"; const V2_DEBUGGER_ENDPOINT_PATH: &str = "/debugger/v2/input"; @@ -73,6 +80,8 @@ const INSTRUMENTATION_ENDPOINT_PATH: &str = "/telemetry/proxy/api/v2/apmtelemetr // Intake endpoints const DSM_INTAKE_PATH: &str = "/api/v0.1/pipeline_stats"; const LLM_OBS_SPANS_INTAKE_PATH: &str = "/api/v2/llmobs"; +const FFE_FLAG_EVALUATION_INTAKE_PATH: &str = "/api/v2/flagevaluation"; +const FFE_EXPOSURES_INTAKE_PATH: &str = "/api/v2/exposures"; const LLM_OBS_EVAL_METRIC_INTAKE_PATH: &str = "/api/intake/llm-obs/v1/eval-metric"; const LLM_OBS_EVAL_METRIC_INTAKE_PATH_V2: &str = "/api/intake/llm-obs/v2/eval-metric"; const PROFILING_INTAKE_PATH: &str = "/api/v2/profile"; @@ -85,6 +94,8 @@ 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; +/// Tracers send EVP batches of about 5 MiB at most; the Agent's EVP proxy accepts 10 MiB. +const EVP_PROXY_REQUEST_BODY_LIMIT: usize = 10 * 1024 * 1024; pub const MAX_CONTENT_LENGTH: usize = 50 * 1024 * 1024; const LAMBDA_LOAD_SPAN: &str = "aws.lambda.load"; @@ -115,11 +126,12 @@ pub struct ProxyState { /// deterministic drain endpoint) without adding a dedicated field per /// route to [`TraceAgent`]. /// -/// Returning `Err` propagates out of [`TraceAgent::start`], aborting the -/// 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 +/// Returning `Err` propagates out of [`TraceAgent::start_with_listener`] +/// (and [`TraceAgent::start`]), aborting the HTTP listener task. Note that +/// the production convenience entry point +/// [`crate::startup::start_trace_agent`] spawns `start_with_listener` 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 @@ -367,21 +379,47 @@ impl TraceAgent { self } - #[allow(clippy::cast_possible_truncation)] + /// Binds the trace agent's listening socket on `port`, or + /// [`TRACE_AGENT_PORT`] when `None`. Connections are accepted (and queued) + /// from here on, before the agent serves. + /// + /// The extension calls this before it registers: Lambda starts the runtime + /// once every extension has registered, and tracers probe `/info` early in + /// their init, so the socket must already accept connections by then. + /// + /// Sets `SO_REUSEADDR` like `tokio::net::TcpListener::bind`. Must run + /// inside the tokio runtime. + pub fn bind_listener(port: Option) -> std::io::Result { + let port = + port.unwrap_or(u16::try_from(TRACE_AGENT_PORT).expect("TRACE_AGENT_PORT is too large")); + let socket = tokio::net::TcpSocket::new_v4()?; + socket.set_reuseaddr(true)?; + socket.bind(SocketAddr::from(([127, 0, 0, 1], port)))?; + let listener = socket.listen(LISTEN_BACKLOG)?; + debug!( + "TRACE AGENT | Listening on port {}", + listener.local_addr()?.port() + ); + Ok(listener) + } + pub async fn start(&self) -> Result<(), Box> { + let listener = Self::bind_listener(self.receiver_port)?; + self.start_with_listener(listener).await + } + + /// Serves on a listener from [`TraceAgent::bind_listener`]. + #[allow(clippy::cast_possible_truncation)] + pub async fn start_with_listener( + &self, + listener: tokio::net::TcpListener, + ) -> Result<(), Box> { let now = Instant::now(); let router = self.make_router()?; - 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 {port}"); debug!( - "TRACE AGENT | Time taken to start: {} ms", + "TRACE AGENT | Time taken to build the router: {} ms", now.elapsed().as_millis() ); @@ -443,15 +481,6 @@ impl TraceAgent { let proxy_router = Router::new() .route(DSM_AGENT_PATH, post(Self::dsm_proxy)) .route(PROFILING_ENDPOINT_PATH, post(Self::profiling_proxy)) - .route( - LLM_OBS_EVAL_METRIC_ENDPOINT_PATH, - post(Self::llm_obs_eval_metric_proxy), - ) - .route( - LLM_OBS_EVAL_METRIC_ENDPOINT_PATH_V2, - post(Self::llm_obs_eval_metric_proxy_v2), - ) - .route(LLM_OBS_SPANS_ENDPOINT_PATH, post(Self::llm_obs_spans_proxy)) .route(V1_DEBUGGER_ENDPOINT_PATH, post(Self::debugger_logs_proxy)) .route(V2_DEBUGGER_ENDPOINT_PATH, post(Self::debugger_intake_proxy)) .route( @@ -463,6 +492,24 @@ impl TraceAgent { post(Self::instrumentation_proxy), ) .layer(RequestBodyLimitLayer::new(DEFAULT_REQUEST_BODY_LIMIT)) + .with_state(proxy_state.clone()); + + let evp_proxy_router = Router::new() + .route( + LLM_OBS_EVAL_METRIC_ENDPOINT_PATH, + post(Self::llm_obs_eval_metric_proxy), + ) + .route( + LLM_OBS_EVAL_METRIC_ENDPOINT_PATH_V2, + post(Self::llm_obs_eval_metric_proxy_v2), + ) + .route(LLM_OBS_SPANS_ENDPOINT_PATH, post(Self::llm_obs_spans_proxy)) + .route( + FFE_FLAG_EVALUATION_ENDPOINT_PATH, + post(Self::ffe_flag_evaluation_proxy), + ) + .route(FFE_EXPOSURES_ENDPOINT_PATH, post(Self::ffe_exposures_proxy)) + .layer(RequestBodyLimitLayer::new(EVP_PROXY_REQUEST_BODY_LIMIT)) .with_state(proxy_state); let info_router = Router::new().route(INFO_ENDPOINT_PATH, any(Self::info)); @@ -471,6 +518,7 @@ impl TraceAgent { .merge(trace_router) .merge(stats_router) .merge(proxy_router) + .merge(evp_proxy_router) .merge(info_router); if let Some(extension) = &self.router_extension { @@ -594,6 +642,33 @@ impl TraceAgent { .await } + async fn ffe_flag_evaluation_proxy( + State(state): State, + request: Request, + ) -> Response { + Self::handle_proxy( + state.config, + state.proxy_aggregator, + request, + "event-platform-intake", + FFE_FLAG_EVALUATION_INTAKE_PATH, + "ffe_flag_evaluation", + ) + .await + } + + async fn ffe_exposures_proxy(State(state): State, request: Request) -> Response { + Self::handle_proxy( + state.config, + state.proxy_aggregator, + request, + "event-platform-intake", + FFE_EXPOSURES_INTAKE_PATH, + "ffe_exposures", + ) + .await + } + // Used for `/debugger/v1/input` in Exception Replay async fn debugger_logs_proxy(State(state): State, request: Request) -> Response { Self::handle_proxy( @@ -646,10 +721,16 @@ impl TraceAgent { LLM_OBS_EVAL_METRIC_ENDPOINT_PATH, LLM_OBS_EVAL_METRIC_ENDPOINT_PATH_V2, LLM_OBS_SPANS_ENDPOINT_PATH, + EVP_PROXY_V2_BASE_PATH, + FFE_FLAG_EVALUATION_ENDPOINT_PATH, + FFE_EXPOSURES_ENDPOINT_PATH, V1_DEBUGGER_ENDPOINT_PATH, V2_DEBUGGER_ENDPOINT_PATH, DEBUGGER_DIAGNOSTICS_ENDPOINT_PATH, ], + // The proxy flusher keeps these two headers. Agentless flag discovery needs + // both listed before it uses the local route. + "evp_proxy_allowed_headers": ["DD-EVP-ORIGIN", "DD-EVP-ORIGIN-VERSION"], "client_drop_p0s": false, } ); @@ -1189,6 +1270,15 @@ mod tests { SendDataBuilderInfo::new(builder, size, owned_tags) } + /// Once `bind_listener` returns, clients can connect before the agent serves. + #[tokio::test] + async fn bind_accepts_connections_before_serving() { + let listener = TraceAgent::bind_listener(Some(0)).expect("bind"); + let addr = listener.local_addr().expect("local_addr"); + + std::net::TcpStream::connect(addr).expect("connect before serving"); + } + /// 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 @@ -1286,4 +1376,150 @@ mod tests { assert_eq!(response.status(), StatusCode::NOT_FOUND); } + + /// The `/info` fields tracers read to pick a local EVP route. Discovery prefers + /// v4 when it is listed; dd-trace-java and dd-trace-go LLM Obs need the trailing + /// slash; agentless flag discovery needs both identity headers. + #[tokio::test] + async fn info_advertises_evp_proxy_for_feature_flags() { + let response = TraceAgent::info().await; + let body = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .expect("info body"); + let info: serde_json::Value = serde_json::from_slice(&body).expect("info json"); + let endpoints: Vec<&str> = info["endpoints"] + .as_array() + .expect("endpoints") + .iter() + .filter_map(serde_json::Value::as_str) + .collect(); + // Literals, not the production constants, so a typo fails here. + for path in [ + "/evp_proxy/v2/", + "/evp_proxy/v2/api/v2/flagevaluation", + "/evp_proxy/v2/api/v2/exposures", + ] { + assert!(endpoints.contains(&path), "missing {path} in {endpoints:?}"); + } + assert!( + !endpoints + .iter() + .any(|e| e.trim_start_matches('/').starts_with("evp_proxy/v4")), + "v4 is listed, so tracers would send to unrouted v4 paths: {endpoints:?}" + ); + let headers: Vec = info["evp_proxy_allowed_headers"] + .as_array() + .expect("evp_proxy_allowed_headers") + .iter() + .filter_map(serde_json::Value::as_str) + .map(str::to_ascii_lowercase) + .collect(); + for header in ["dd-evp-origin", "dd-evp-origin-version"] { + assert!( + headers.iter().any(|h| h == header), + "missing {header} in {headers:?}" + ); + } + } + + /// EVP paths without a route return 404 and are not buffered. + #[tokio::test] + async fn unrouted_evp_paths_return_not_found() { + let agent = build_test_agent(); + let router = agent.make_router().expect("make_router"); + let response = router + .oneshot( + Request::builder() + .method("POST") + .uri("/evp_proxy/v2/api/v2/not-a-flag-route") + .header("X-Datadog-EVP-Subdomain", "event-platform-intake") + .body(Body::from("{}")) + .expect("build request"), + ) + .await + .expect("route response"); + + assert_eq!(response.status(), StatusCode::NOT_FOUND); + assert!(agent.proxy_aggregator.lock().await.get_batch().is_empty()); + } + + /// Tracers send EVP batches of about 5 MiB at most (the envelope can add a + /// little). Sent with Content-Length, such a batch gets 413 under the default + /// 2 MiB proxy limit, and tracers drop the rejected batch. + #[tokio::test] + async fn evp_proxy_routes_accept_five_mib_batches() { + for path in [ + LLM_OBS_SPANS_ENDPOINT_PATH, + LLM_OBS_EVAL_METRIC_ENDPOINT_PATH, + LLM_OBS_EVAL_METRIC_ENDPOINT_PATH_V2, + FFE_FLAG_EVALUATION_ENDPOINT_PATH, + FFE_EXPOSURES_ENDPOINT_PATH, + ] { + for (size, expected) in [ + (5 * 1024 * 1024 + 64 * 1024, StatusCode::OK), + ( + EVP_PROXY_REQUEST_BODY_LIMIT + 1, + StatusCode::PAYLOAD_TOO_LARGE, + ), + ] { + let agent = build_test_agent(); + let router = agent.make_router().expect("make_router"); + let response = router + .oneshot( + Request::builder() + .method("POST") + .uri(path) + .header("Content-Type", "application/json") + .header("Content-Length", size.to_string()) + .body(Body::from(vec![b' '; size])) + .expect("build request"), + ) + .await + .expect("route response"); + assert_eq!(response.status(), expected, "{path} with {size} bytes"); + let queued = agent.proxy_aggregator.lock().await.get_batch().len(); + assert_eq!(queued, usize::from(expected == StatusCode::OK), "{path}"); + } + } + } + + #[tokio::test] + async fn feature_flag_events_are_relayed_to_event_platform_intake() { + let agent = build_test_agent(); + let site = agent.config.site.clone(); + let router = agent.make_router().expect("make_router"); + // Literals: the tracer paths and the intake paths are the contract. + for (path, intake) in [ + ( + "/evp_proxy/v2/api/v2/flagevaluation", + "/api/v2/flagevaluation", + ), + ("/evp_proxy/v2/api/v2/exposures", "/api/v2/exposures"), + ] { + let response = router + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri(path) + .header("Content-Type", "application/json") + .header("DD-EVP-ORIGIN", "dd-trace-py") + .header("DD-EVP-ORIGIN-VERSION", "4.17.0") + .header("X-Datadog-EVP-Subdomain", "event-platform-intake") + .body(Body::from(r#"{"events":[]}"#)) + .expect("build request"), + ) + .await + .expect("route response"); + assert_eq!(response.status(), StatusCode::OK, "{path}"); + let batch = agent.proxy_aggregator.lock().await.get_batch(); + assert_eq!(batch.len(), 1, "{path}"); + assert_eq!( + batch[0].target_url, + format!("https://event-platform-intake.{site}{intake}") + ); + assert_eq!(batch[0].headers["dd-evp-origin"], "dd-trace-py"); + assert_eq!(batch[0].headers["dd-evp-origin-version"], "4.17.0"); + } + } } diff --git a/bottlecap/tests/common/fake_extensions_api.rs b/bottlecap/tests/common/fake_extensions_api.rs index 1e39d8eb9..4b370f968 100644 --- a/bottlecap/tests/common/fake_extensions_api.rs +++ b/bottlecap/tests/common/fake_extensions_api.rs @@ -74,6 +74,9 @@ struct Shared { telemetry_url: Arc>>, /// `types` of the last accepted subscribe body. telemetry_types: Arc>>, + /// Whether the trace agent port accepted a connection while the first + /// register request was handled, `None` until then. + trace_agent_listening_at_register: Arc>>, } pub struct FakeExtensionsApi { @@ -121,6 +124,7 @@ impl FakeExtensionsApi { subscribe_status: Arc::new(AtomicU16::new(StatusCode::OK.as_u16())), telemetry_url: Arc::new(Mutex::new(None)), telemetry_types: Arc::new(Mutex::new(Vec::new())), + trace_agent_listening_at_register: Arc::new(Mutex::new(None)), }; let router = Router::new() .route("/2020-01-01/extension/register", post(register)) @@ -239,6 +243,16 @@ impl FakeExtensionsApi { self.shared.requests.lock().expect("requests lock").clone() } + /// Whether the trace agent port accepted a connection while the first + /// register request was handled, `None` before any register. + pub fn trace_agent_listening_at_register(&self) -> Option { + *self + .shared + .trace_agent_listening_at_register + .lock() + .expect("trace agent probe lock") + } + // --- Telemetry API sender ---------------------------------------------- /// Port the extension asked Lambda to deliver telemetry to, `None` until @@ -345,6 +359,21 @@ async fn register( body: Bytes, ) -> Result { record(&shared, "POST", "/2020-01-01/extension/register"); + // Lambda starts the runtime, and with it the tracer's early `/info` + // probe, once every extension has registered. Probe the trace agent + // once, before answering, the way that tracer would. + { + let mut at_register = shared + .trace_agent_listening_at_register + .lock() + .expect("trace agent probe lock"); + if at_register.is_none() { + let addr = SocketAddr::from(([127, 0, 0, 1], super::port_guard::TRACE_AGENT_PORT)); + let listening = + std::net::TcpStream::connect_timeout(&addr, Duration::from_millis(200)).is_ok(); + *at_register = Some(listening); + } + } // Lambda answers `403` without a name. Any name but the file under // `/opt/extensions` registers an internal extension, which cannot take SHUTDOWN. if headers diff --git a/bottlecap/tests/evp_proxy_flush_test.rs b/bottlecap/tests/evp_proxy_flush_test.rs new file mode 100644 index 000000000..36c11d0e5 --- /dev/null +++ b/bottlecap/tests/evp_proxy_flush_test.rs @@ -0,0 +1,91 @@ +// Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 +// Integration tests are exempt from the production lint policy in Cargo.toml. +#![allow( + unreachable_pub, + clippy::pedantic, + clippy::panic, + clippy::unwrap_used, + clippy::expect_used +)] + +//! The proxy flusher sends a queued flag event batch with the API key and the +//! tracer's identity headers. The trace agent builds the intake URL; here the +//! batch targets a local intake directly. + +mod fixtures; +#[path = "common/harness.rs"] +mod harness; + +use std::collections::HashMap; +use std::sync::Arc; + +use bottlecap::LAMBDA_RUNTIME_SLUG; +use bottlecap::config::Config; +use bottlecap::tags::provider::Provider; +use bottlecap::traces::proxy_aggregator::{Aggregator as ProxyAggregator, ProxyRequest}; +use bottlecap::traces::proxy_flusher::Flusher as ProxyFlusher; +use datadog_fips::reqwest_adapter::create_reqwest_client_builder; +use dogstatsd::api_key::ApiKeyFactory; +use harness::prelude::*; +use tokio::sync::Mutex; + +#[tokio::test] +async fn flag_event_batch_is_flushed_with_api_key_and_origin_headers() { + let intake = FaultIntake::start(FaultMode::Ok).await; + let config = Arc::new(Config { + api_key: TEST_API_KEY.to_string(), + site: "datadoghq.com".to_string(), + ..Config::default() + }); + let tags = Arc::new(Provider::new( + Arc::clone(&config), + LAMBDA_RUNTIME_SLUG.to_string(), + &HashMap::from([( + "function_arn".to_string(), + "arn:aws:lambda:us-east-1:123456789012:function:test".to_string(), + )]), + )); + let aggregator = Arc::new(Mutex::new(ProxyAggregator::default())); + let mut headers = reqwest::header::HeaderMap::new(); + headers.insert("content-type", "application/json".parse().unwrap()); + headers.insert("dd-evp-origin", "dd-trace-js".parse().unwrap()); + headers.insert("dd-evp-origin-version", "5.80.0".parse().unwrap()); + aggregator.lock().await.add(ProxyRequest { + headers, + body: r#"{"exposures":[{"flag":{"key":"harness-exposed-flag"}}]}"#.into(), + target_url: format!("{}/api/v2/exposures", intake.url()), + }); + // Test-only client towards a local intake; the FIPS adapter is what production uses. + let client = create_reqwest_client_builder() + .expect("client builder") + .no_proxy() + .build() + .expect("client"); + let flusher = ProxyFlusher::new( + Arc::new(ApiKeyFactory::new(TEST_API_KEY)), + Arc::clone(&aggregator), + tags, + Arc::clone(&config), + client, + ); + + let failed = flusher.flush(None).await; + assert!( + failed.is_none(), + "flush failed: {failed:?}. {}", + intake.summary() + ); + + let sent = intake.requests_to("/api/v2/exposures"); + assert_eq!(sent.len(), 1, "{}", intake.summary()); + assert_eq!(sent[0].api_key(), Some(TEST_API_KEY)); + assert_eq!(sent[0].header("dd-evp-origin"), Some("dd-trace-js")); + assert_eq!(sent[0].header("dd-evp-origin-version"), Some("5.80.0")); + assert!( + sent[0] + .body_text() + .unwrap() + .contains("harness-exposed-flag") + ); +} diff --git a/bottlecap/tests/lifecycle_contract_test.rs b/bottlecap/tests/lifecycle_contract_test.rs index 1882b3f1c..f955edbf4 100644 --- a/bottlecap/tests/lifecycle_contract_test.rs +++ b/bottlecap/tests/lifecycle_contract_test.rs @@ -89,3 +89,46 @@ async fn invoke_then_shutdown_exits_cleanly() { s.intake.summary() ); } + +/// Lambda starts the runtime, and the tracer's first `/info` call, once every +/// extension has registered, not at the extension's first `/next`. The fake +/// probes the trace agent port while it handles register, so the trace agent +/// must already be listening by then. +async fn assert_trace_agent_listens_before_register(managed_instance: bool) { + let s = Scenario::start(FaultMode::Ok).await; + let builder = s.bottlecap(); + let mut bc = if managed_instance { + builder.managed_instance().spawn() + } else { + builder.spawn() + }; + + assert!( + s.api.wait_for_next_calls(1, bc.startup_deadline()).await, + "no /next poll. {}", + bc.log_tail(40) + ); + assert_eq!( + s.api.trace_agent_listening_at_register(), + Some(true), + "the trace agent was not listening on 8126 when the extension registered. {}", + bc.log_tail(40) + ); + // Only the startup order matters here. The process is killed rather than + // shut down because the On-Demand loop does not yet act on a SHUTDOWN that + // answers the first `/next`. + bc.assert_no_panic(); + bc.kill(); +} + +#[tokio::test(flavor = "multi_thread")] +#[serial_test::serial(bottlecap_ports)] +async fn trace_agent_listens_before_register() { + assert_trace_agent_listens_before_register(false).await; +} + +#[tokio::test(flavor = "multi_thread")] +#[serial_test::serial(bottlecap_ports)] +async fn trace_agent_listens_before_register_managed_instance() { + assert_trace_agent_listens_before_register(true).await; +}