From 7105107498e9e86d52760c791fb17e5f8d67cd89 Mon Sep 17 00:00:00 2001 From: Tianning Li Date: Tue, 20 Oct 2026 12:43:17 -0400 Subject: [PATCH] feat(logs-agent): optional API key resolver and Datadog intake URL override Why: bottlecap gets its API key lazily (it may come from AWS Secrets Manager and can be reloaded) and lets users point logs at a custom intake URL. This crate only took a fixed key and a fixed URL, so bottlecap could not use it without changing behavior for Lambda users. What: the flusher gets two optional builder methods. with_api_key_resolver() calls a key resolver for the primary key on every flush, so a reloaded key is used right away; if no key comes back, the primary is skipped for that flush with one error line and additional endpoints are still sent. with_datadog_intake_url() replaces the default Datadog intake URL. Without them the flusher behaves exactly as before, and no dependencies are added. Co-Authored-By: Claude Opus 5.5 --- crates/datadog-logs-agent/src/flusher.rs | 61 +++++- .../datadog-logs-agent/src/flusher/tests.rs | 206 ++++++++++++++++++ crates/datadog-logs-agent/src/lib.rs | 2 +- 3 files changed, 263 insertions(+), 6 deletions(-) diff --git a/crates/datadog-logs-agent/src/flusher.rs b/crates/datadog-logs-agent/src/flusher.rs index a9d038f..7ed1cae 100644 --- a/crates/datadog-logs-agent/src/flusher.rs +++ b/crates/datadog-logs-agent/src/flusher.rs @@ -1,7 +1,10 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ // SPDX-License-Identifier: Apache-2.0 +use std::future::Future; use std::io::Write as _; +use std::pin::Pin; +use std::sync::Arc; use futures::future::join_all; use reqwest::Client; @@ -15,12 +18,18 @@ use crate::errors::FlushError; /// Maximum number of send attempts before giving up on a batch. const MAX_FLUSH_ATTEMPTS: u32 = 3; +/// Returns the API key for the primary endpoint, or `None` if it cannot be resolved. +pub type ApiKeyResolver = + Arc Pin> + Send>> + Send + Sync>; + /// Drains log batches from an [`AggregatorHandle`] and ships them to Datadog. #[derive(Clone)] pub struct LogFlusher { config: LogFlusherConfig, client: Client, aggregator_handle: AggregatorHandle, + api_key_resolver: Option, + datadog_intake_url: Option, } impl LogFlusher { @@ -38,9 +47,30 @@ impl LogFlusher { config, client, aggregator_handle, + api_key_resolver: None, + datadog_intake_url: None, } } + /// Get the primary endpoint's API key from `resolver` on every flush instead of using + /// `LogFlusherConfig::api_key`. If it returns `None`, the primary endpoint is skipped + /// for that flush (its batches are dropped, with one `error!`); additional endpoints + /// still use their own keys. + #[must_use] + pub fn with_api_key_resolver(mut self, resolver: ApiKeyResolver) -> Self { + self.api_key_resolver = Some(resolver); + self + } + + /// In `Destination::Datadog` mode, send to `url` (full URL, including the path) instead + /// of `https://http-intake.logs.{site}/api/v2/logs`. Ignored in Observability Pipelines + /// Worker mode. + #[must_use] + pub fn with_datadog_intake_url(mut self, url: impl Into) -> Self { + self.datadog_intake_url = Some(url.into()); + self + } + /// Drain the aggregator, ship all pending batches to Datadog, and redrive any /// builders that failed transiently in the previous invocation. /// @@ -55,7 +85,8 @@ impl LogFlusher { /// A vec of `RequestBuilder`s that still failed after all in-call retries. /// The caller should pass these back on the next invocation to re-attempt /// delivery. An empty vec means every batch was delivered successfully - /// (or encountered a permanent error and was dropped — those are logged). + /// (or encountered a permanent error, or the primary API key could not be + /// resolved, and was dropped — those are logged). /// /// Failures on additional endpoints are logged as warnings but their /// builders are not included in the returned vec (best-effort delivery). @@ -88,6 +119,19 @@ impl LogFlusher { } }; + let primary_api_key = match &self.api_key_resolver { + None => Some(self.config.api_key.clone()), + Some(resolve) => { + let key = resolve().await; + if key.is_none() { + error!( + "failed to resolve the logs API key; skipping the primary endpoint for this flush" + ); + } + key + } + }; + if batches.is_empty() { debug!("no log batches to flush"); return failed; @@ -99,12 +143,17 @@ impl LogFlusher { let batch_futures = batches.iter().map(|batch| { let primary_url = primary_url.clone(); + let primary_api_key = primary_api_key.as_deref(); async move { // Primary endpoint — failures are tracked for cross-invocation retry. let is_primary_datadog = matches!(self.config.mode, Destination::Datadog); - let primary_result = self - .ship_batch(batch, &primary_url, primary_use_compression, &self.config.api_key, is_primary_datadog) - .await; + let primary_result = match primary_api_key { + Some(api_key) => { + self.ship_batch(batch, &primary_url, primary_use_compression, api_key, is_primary_datadog) + .await + } + None => Ok(()), + }; // Additional endpoints — best-effort; failures are only logged. // Additional endpoints are always Datadog intakes regardless of the @@ -145,7 +194,9 @@ impl LogFlusher { fn resolve_endpoint(&self) -> (String, bool) { match &self.config.mode { Destination::Datadog => { - let url = format!("https://http-intake.logs.{}/api/v2/logs", self.config.site); + let url = self.datadog_intake_url.clone().unwrap_or_else(|| { + format!("https://http-intake.logs.{}/api/v2/logs", self.config.site) + }); (url, self.config.use_compression) } Destination::ObservabilityPipelinesWorker { url } => { diff --git a/crates/datadog-logs-agent/src/flusher/tests.rs b/crates/datadog-logs-agent/src/flusher/tests.rs index c75bb25..efede8d 100644 --- a/crates/datadog-logs-agent/src/flusher/tests.rs +++ b/crates/datadog-logs-agent/src/flusher/tests.rs @@ -1158,3 +1158,209 @@ async fn test_default_logging_warns_once_per_dropped_oversized_entry() { assert_eq!(logged.warnings(), 2); assert_eq!(logged.errors(), 0, "default logging uses warn!, not error!"); } + +/// A resolver that counts its calls and returns `keys` in order; the last key repeats. +fn resolver(keys: &[Option<&'static str>], calls: Arc) -> ApiKeyResolver { + let keys: Vec> = keys.iter().map(|k| k.map(str::to_string)).collect(); + Arc::new( + move || -> std::pin::Pin> + Send>> { + let n = calls.fetch_add(1, Ordering::SeqCst); + let key = keys.get(n).or(keys.last()).cloned().flatten(); + Box::pin(async move { key }) + }, + ) +} + +/// In Datadog mode, `with_datadog_intake_url` replaces the public intake URL; headers +/// and compression are unchanged. +#[tokio::test] +async fn test_datadog_intake_url_override_is_used_in_datadog_mode() { + let (service, handle) = AggregatorService::new(); + let _task = tokio::spawn(service.run()); + let mut server = mockito::Server::new_async().await; + let mock = server + .mock("POST", "/custom/path") + .match_header("DD-API-KEY", "test-api-key") + .match_header("DD-PROTOCOL", "agent-json") + .match_header("Content-Encoding", "zstd") + .with_status(202) + .expect(1) + .create_async() + .await; + let config = LogFlusherConfig { + mode: Destination::Datadog, + use_compression: true, + ..config_for_mock("http://unused") + }; + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new(config, client, handle.clone()) + .with_datadog_intake_url(format!("{}/custom/path", server.url())); + handle + .insert_batch(vec![make_entry("override")]) + .expect("insert"); + + assert!(flusher.flush(vec![]).await.is_empty()); + mock.assert_async().await; +} + +/// `with_datadog_intake_url` does not affect Observability Pipelines Worker mode. +#[tokio::test] +async fn test_datadog_intake_url_override_is_ignored_for_opw() { + let (service, handle) = AggregatorService::new(); + let _task = tokio::spawn(service.run()); + let mut server = mockito::Server::new_async().await; + let mock = server + .mock("POST", "/api/v2/logs") + .with_status(202) + .expect(1) + .create_async() + .await; + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new(config_for_mock(&server.url()), client, handle.clone()) + .with_datadog_intake_url("http://127.0.0.1:1/never"); + handle + .insert_batch(vec![make_entry("opw")]) + .expect("insert"); + + assert!(flusher.flush(vec![]).await.is_empty()); + mock.assert_async().await; +} + +/// With a resolver, the primary uses the resolved key, not `LogFlusherConfig::api_key`. +#[tokio::test] +async fn test_api_key_resolver_key_is_used_for_the_primary() { + let (service, handle) = AggregatorService::new(); + let _task = tokio::spawn(service.run()); + let mut server = mockito::Server::new_async().await; + let mock = server + .mock("POST", "/api/v2/logs") + .match_header("DD-API-KEY", "resolved-key") + .with_status(202) + .expect(1) + .create_async() + .await; + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new(config_for_mock(&server.url()), client, handle.clone()) + .with_api_key_resolver(resolver( + &[Some("resolved-key")], + Arc::new(AtomicUsize::new(0)), + )); + handle + .insert_batch(vec![make_entry("keyed")]) + .expect("insert"); + + assert!(flusher.flush(vec![]).await.is_empty()); + mock.assert_async().await; +} + +/// The resolver is called on every flush, even when there is nothing to send. +#[tokio::test] +async fn test_api_key_resolver_is_called_on_every_flush() { + let calls = Arc::new(AtomicUsize::new(0)); + let (service, handle) = AggregatorService::new(); + let _task = tokio::spawn(service.run()); + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new(config_for_mock("http://127.0.0.1:1"), client, handle) + .with_api_key_resolver(resolver(&[Some("k")], calls.clone())); + + for _ in 0..3 { + assert!(flusher.flush(vec![]).await.is_empty()); + } + assert_eq!(calls.load(Ordering::SeqCst), 3); +} + +/// A key that changes between flushes (e.g. a reloaded secret) is used on the next flush. +#[tokio::test] +async fn test_reloaded_api_key_is_used_on_the_next_flush() { + let (service, handle) = AggregatorService::new(); + let _task = tokio::spawn(service.run()); + let mut server = mockito::Server::new_async().await; + let first = server + .mock("POST", "/api/v2/logs") + .match_header("DD-API-KEY", "key-1") + .with_status(202) + .expect(1) + .create_async() + .await; + let second = server + .mock("POST", "/api/v2/logs") + .match_header("DD-API-KEY", "key-2") + .with_status(202) + .expect(1) + .create_async() + .await; + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new(config_for_mock(&server.url()), client, handle.clone()) + .with_api_key_resolver(resolver( + &[Some("key-1"), Some("key-2")], + Arc::new(AtomicUsize::new(0)), + )); + + handle + .insert_batch(vec![make_entry("one")]) + .expect("insert"); + assert!(flusher.flush(vec![]).await.is_empty()); + handle + .insert_batch(vec![make_entry("two")]) + .expect("insert"); + assert!(flusher.flush(vec![]).await.is_empty()); + first.assert_async().await; + second.assert_async().await; +} + +/// If the key cannot be resolved, the primary is skipped with one error line, and +/// additional endpoints are still sent with their own keys. +#[tokio::test(flavor = "current_thread")] +async fn test_unresolved_api_key_skips_the_primary_but_sends_extras() { + let (logged, _guard) = LevelCounter::install(); + let (service, handle) = AggregatorService::new(); + let _task = tokio::spawn(service.run()); + let mut primary = mockito::Server::new_async().await; + let mut extra = mockito::Server::new_async().await; + let primary_mock = primary + .mock("POST", "/api/v2/logs") + .expect(0) + .create_async() + .await; + let extra_mock = extra + .mock("POST", "/extra") + .match_header("DD-API-KEY", "extra-key") + .with_status(202) + .expect(1) + .create_async() + .await; + let config = LogFlusherConfig { + additional_endpoints: vec![LogsAdditionalEndpoint { + api_key: "extra-key".to_string(), + url: format!("{}/extra", extra.url()), + is_reliable: true, + }], + ..config_for_mock(&primary.url()) + }; + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new(config, client, handle.clone()) + .with_api_key_resolver(resolver(&[None], Arc::new(AtomicUsize::new(0)))); + handle + .insert_batch(vec![make_entry("log")]) + .expect("insert"); + + assert!(flusher.flush(vec![]).await.is_empty()); + primary_mock.assert_async().await; + extra_mock.assert_async().await; + assert_eq!((logged.warnings(), logged.errors()), (0, 1)); +} + +/// The unresolved-key error is logged on every flush, even with nothing to send. +#[tokio::test(flavor = "current_thread")] +async fn test_unresolved_api_key_logs_an_error_even_without_batches() { + let (logged, _guard) = LevelCounter::install(); + let (service, handle) = AggregatorService::new(); + let _task = tokio::spawn(service.run()); + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new(config_for_mock("http://127.0.0.1:1"), client, handle) + .with_api_key_resolver(resolver(&[None], Arc::new(AtomicUsize::new(0)))); + + assert!(flusher.flush(vec![]).await.is_empty()); + assert!(flusher.flush(vec![]).await.is_empty()); + assert_eq!(logged.errors(), 2); +} diff --git a/crates/datadog-logs-agent/src/lib.rs b/crates/datadog-logs-agent/src/lib.rs index 5f2b0f6..b33ae52 100644 --- a/crates/datadog-logs-agent/src/lib.rs +++ b/crates/datadog-logs-agent/src/lib.rs @@ -22,7 +22,7 @@ mod test_support; // Re-export the most commonly used types at the crate root pub use aggregator::{AggregatorHandle, AggregatorService, OversizePolicy}; pub use config::{Destination, LogFlusherConfig}; -pub use flusher::LogFlusher; +pub use flusher::{ApiKeyResolver, LogFlusher}; pub use intake_entry::IntakeEntry; pub use logs_additional_endpoint::LogsAdditionalEndpoint; pub use server::{LogServer, LogServerConfig};