From 179bcd8adcd25b0500ca5581208403131f4962f5 Mon Sep 17 00:00:00 2001 From: Tianning Li Date: Wed, 21 Oct 2026 15:51:58 -0400 Subject: [PATCH] feat(logs-agent): accept pre-serialized logs and an oversize policy Why: bottlecap already turns each Lambda log into JSON in its own format, which this crate's entry type cannot express, and it sends logs over 1 MiB (with a warning) instead of dropping them. To use this crate without changing what Lambda users see, the aggregator needs both, as options. What: - The aggregator handle gets insert_serialized_batch(), which queues JSON strings as-is in the same queue as normal entries. - The aggregator and its service get an oversize policy: the default still drops entries over 1 MiB; the new "send with warning" option keeps them with one warning unless they are too big for any payload. Nothing changes for existing users. - The flusher's test module moves unchanged to src/flusher/tests.rs (it had grown to about 1,500 lines, mostly tests), and the counter tests use to count warning and error lines moves to a small test-only helper so other test modules can use it too. Most of this change's diff is that move. Co-Authored-By: Claude Opus 5.5 --- .../datadog-logs-agent/src/aggregator/core.rs | 152 ++- .../datadog-logs-agent/src/aggregator/mod.rs | 2 +- .../src/aggregator/service.rs | 112 +- crates/datadog-logs-agent/src/flusher.rs | 1206 +---------------- .../datadog-logs-agent/src/flusher/tests.rs | 1160 ++++++++++++++++ crates/datadog-logs-agent/src/lib.rs | 4 +- crates/datadog-logs-agent/src/test_support.rs | 52 + 7 files changed, 1474 insertions(+), 1214 deletions(-) create mode 100644 crates/datadog-logs-agent/src/flusher/tests.rs create mode 100644 crates/datadog-logs-agent/src/test_support.rs diff --git a/crates/datadog-logs-agent/src/aggregator/core.rs b/crates/datadog-logs-agent/src/aggregator/core.rs index 8f39d23d..70c283e8 100644 --- a/crates/datadog-logs-agent/src/aggregator/core.rs +++ b/crates/datadog-logs-agent/src/aggregator/core.rs @@ -3,10 +3,25 @@ use std::collections::VecDeque; +use tracing::warn; + use crate::constants::{MAX_BATCH_ENTRIES, MAX_CONTENT_BYTES, MAX_LOG_BYTES}; use crate::errors::AggregatorError; use crate::intake_entry::IntakeEntry; +/// What the aggregator does with a log entry larger than `MAX_LOG_BYTES`. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +#[non_exhaustive] +pub enum OversizePolicy { + /// Drop the entry: `insert` returns `EntryTooLarge`, and the aggregator service logs + /// one warning per dropped entry. + #[default] + Drop, + /// Keep the entry, with one warning, as long as it still fits in a payload on its own. + /// Entries too big for any payload are dropped, with one warning from the service. + SendWithWarning, +} + /// In-memory log batch accumulator. /// /// Stores pre-serialized JSON strings in a FIFO queue. A batch is "full" @@ -15,6 +30,7 @@ use crate::intake_entry::IntakeEntry; pub struct Aggregator { messages: VecDeque, current_size_bytes: usize, + oversize_policy: OversizePolicy, } impl Default for Aggregator { @@ -29,23 +45,54 @@ impl Aggregator { Self { messages: VecDeque::new(), current_size_bytes: 0, + oversize_policy: OversizePolicy::default(), } } + /// Set how entries larger than `MAX_LOG_BYTES` are handled. Default: drop them. + #[must_use] + pub fn with_oversize_policy(mut self, policy: OversizePolicy) -> Self { + self.oversize_policy = policy; + self + } + /// Insert a log entry into the batch. /// /// Returns `Ok(true)` if the batch is now full and ready to flush. - /// Returns `Err(AggregatorError::EntryTooLarge)` if the serialized - /// entry exceeds `MAX_LOG_BYTES` — the entry is dropped. + /// Returns `Err(AggregatorError::EntryTooLarge)` if the entry is dropped under the + /// oversize policy. pub fn insert(&mut self, entry: &IntakeEntry) -> Result { let serialized = serde_json::to_string(entry)?; + self.insert_serialized(serialized) + } + + /// Insert an entry that is already serialized as a JSON object. It is batched as-is, + /// without re-serializing or validating it: the caller must pass valid JSON, or the + /// whole payload it lands in is rejected by the intake. + /// + /// Same return values as [`Aggregator::insert`]. + pub fn insert_serialized(&mut self, serialized: String) -> Result { let len = serialized.len(); if len > MAX_LOG_BYTES { - return Err(AggregatorError::EntryTooLarge { - size: len, - max: MAX_LOG_BYTES, - }); + match self.oversize_policy { + OversizePolicy::Drop => { + return Err(AggregatorError::EntryTooLarge { + size: len, + max: MAX_LOG_BYTES, + }); + } + OversizePolicy::SendWithWarning => { + // The entry must fit in a payload on its own: `[` + entry + `]`. + let max = MAX_CONTENT_BYTES - 2; + if len > max { + return Err(AggregatorError::EntryTooLarge { size: len, max }); + } + warn!( + "log entry of {len} bytes exceeds {MAX_LOG_BYTES} bytes; sending it anyway" + ); + } + } } self.messages.push_back(serialized); @@ -291,4 +338,97 @@ mod tests { )); assert_eq!(agg.len(), 1); } + + /// A JSON string of exactly `len` bytes, for size tests. + fn json_string_of_len(len: usize) -> String { + format!("\"{}\"", "a".repeat(len - 2)) + } + + #[test] + fn test_insert_serialized_keeps_entry_bytes_as_is() { + let mut agg = Aggregator::new(); + let raw = r#"{"message":{"message":"x","lambda":{"arn":"a","request_id":null}},"ddsource":"lambda"}"#; + agg.insert_serialized(raw.to_string()).expect("insert"); + assert_eq!( + agg.get_batch().expect("batch"), + format!("[{raw}]").into_bytes() + ); + } + + #[test] + fn test_insert_serialized_drops_oversized_entry_by_default() { + let mut agg = Aggregator::new(); + let result = agg.insert_serialized(json_string_of_len(MAX_LOG_BYTES + 1)); + assert!(matches!( + result, + Err(AggregatorError::EntryTooLarge { size, max }) + if size == MAX_LOG_BYTES + 1 && max == MAX_LOG_BYTES + )); + assert!(agg.is_empty()); + } + + #[test] + fn test_send_with_warning_keeps_entry_over_max_log_bytes() { + let mut agg = Aggregator::new().with_oversize_policy(OversizePolicy::SendWithWarning); + let big = json_string_of_len(2 * MAX_LOG_BYTES); + agg.insert_serialized(big.clone()).expect("kept"); + assert_eq!( + agg.get_batch().expect("batch"), + format!("[{big}]").into_bytes() + ); + } + + #[test] + fn test_send_with_warning_drops_only_entries_that_cannot_fit_a_payload() { + let mut agg = Aggregator::new().with_oversize_policy(OversizePolicy::SendWithWarning); + let largest = MAX_CONTENT_BYTES - 2; + agg.insert_serialized(json_string_of_len(largest)) + .expect("fits exactly"); + assert_eq!(agg.get_batch().expect("batch").len(), MAX_CONTENT_BYTES); + + let result = agg.insert_serialized(json_string_of_len(largest + 1)); + assert!(matches!( + result, + Err(AggregatorError::EntryTooLarge { size, max }) + if size == largest + 1 && max == largest + )); + } + + /// Three 2 MiB entries cannot share one 5 MiB payload: they are split across batches, + /// every batch stays within the limit, and nothing is lost. + #[test] + fn test_send_with_warning_splits_oversized_entries_across_batches() { + let mut agg = Aggregator::new().with_oversize_policy(OversizePolicy::SendWithWarning); + for _ in 0..3 { + agg.insert_serialized(json_string_of_len(2 * MAX_LOG_BYTES)) + .expect("kept"); + } + agg.insert_serialized(r#"{"after":1}"#.to_string()) + .expect("kept"); + + let batches = agg.get_all_batches(); + assert!( + batches.len() >= 2, + "three 2 MiB entries need at least two payloads" + ); + assert!(batches.iter().all(|b| b.len() <= MAX_CONTENT_BYTES)); + let entries: usize = batches + .iter() + .map(|b| { + serde_json::from_slice::>(b) + .expect("json") + .len() + }) + .sum(); + assert_eq!(entries, 4); + assert!(agg.is_empty()); + } + + #[test] + fn test_insert_entry_honors_send_with_warning() { + let mut agg = Aggregator::new().with_oversize_policy(OversizePolicy::SendWithWarning); + agg.insert(&make_entry(&"a".repeat(MAX_LOG_BYTES + 1))) + .expect("kept"); + assert_eq!(agg.len(), 1); + } } diff --git a/crates/datadog-logs-agent/src/aggregator/mod.rs b/crates/datadog-logs-agent/src/aggregator/mod.rs index c191ab66..784936c5 100644 --- a/crates/datadog-logs-agent/src/aggregator/mod.rs +++ b/crates/datadog-logs-agent/src/aggregator/mod.rs @@ -2,7 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 pub mod core; -pub use core::Aggregator; +pub use core::{Aggregator, OversizePolicy}; pub mod service; pub use service::{AggregatorHandle, AggregatorService}; diff --git a/crates/datadog-logs-agent/src/aggregator/service.rs b/crates/datadog-logs-agent/src/aggregator/service.rs index c4ffc644..34777f37 100644 --- a/crates/datadog-logs-agent/src/aggregator/service.rs +++ b/crates/datadog-logs-agent/src/aggregator/service.rs @@ -4,12 +4,13 @@ use tokio::sync::{mpsc, oneshot}; use tracing::{debug, error, warn}; -use crate::aggregator::Aggregator; +use crate::aggregator::{Aggregator, OversizePolicy}; use crate::intake_entry::IntakeEntry; #[derive(Debug)] enum AggregatorCommand { InsertBatch(Vec), + InsertSerializedBatch(Vec), GetBatches(oneshot::Sender>>), Shutdown, } @@ -30,6 +31,17 @@ impl AggregatorHandle { .map_err(|e| format!("failed to send InsertBatch: {e}")) } + /// Queue log entries that are already serialized as JSON objects. They are batched + /// as-is, without validation; the caller must pass valid JSON. The oversize policy + /// still applies. + /// + /// Returns an error only if the service has already stopped. + pub fn insert_serialized_batch(&self, entries: Vec) -> Result<(), String> { + self.tx + .send(AggregatorCommand::InsertSerializedBatch(entries)) + .map_err(|e| format!("failed to send InsertSerializedBatch: {e}")) + } + /// Retrieve and drain all accumulated log batches as JSON arrays. /// /// Returns an empty `Vec` if the aggregator holds no logs. @@ -71,6 +83,13 @@ impl AggregatorService { (service, handle) } + /// Set how entries larger than `MAX_LOG_BYTES` are handled. Default: drop them. + #[must_use] + pub fn with_oversize_policy(mut self, policy: OversizePolicy) -> Self { + self.aggregator = std::mem::take(&mut self.aggregator).with_oversize_policy(policy); + self + } + /// Run the service event loop. /// /// Returns when a `Shutdown` command is received or the last handle is dropped. @@ -87,6 +106,14 @@ impl AggregatorService { } } + AggregatorCommand::InsertSerializedBatch(entries) => { + for entry in entries { + if let Err(e) = self.aggregator.insert_serialized(entry) { + warn!("dropping log entry: {e}"); + } + } + } + AggregatorCommand::GetBatches(response_tx) => { let batches = self.aggregator.get_all_batches(); if response_tx.send(batches).is_err() { @@ -182,4 +209,87 @@ mod tests { handle.shutdown().expect("shutdown"); task.await.expect("task"); } + + use crate::test_support::LevelCounter; + + fn json_string_of_len(len: usize) -> String { + format!("\"{}\"", "a".repeat(len - 2)) + } + + #[tokio::test] + async fn test_insert_serialized_batch_roundtrip() { + let (service, handle) = AggregatorService::new(); + let _task = tokio::spawn(service.run()); + handle + .insert_serialized_batch(vec![r#"{"a":1}"#.to_string(), r#"{"b":2}"#.to_string()]) + .expect("insert"); + let batches = handle.get_batches().await.expect("batches"); + assert_eq!(batches, vec![br#"[{"a":1},{"b":2}]"#.to_vec()]); + } + + #[tokio::test] + async fn test_entry_and_serialized_inserts_share_one_queue() { + let (service, handle) = AggregatorService::new(); + let _task = tokio::spawn(service.run()); + handle + .insert_batch(vec![IntakeEntry::from_message("first", 1)]) + .expect("insert"); + handle + .insert_serialized_batch(vec![r#"{"second":true}"#.to_string()]) + .expect("insert"); + let batches = handle.get_batches().await.expect("batches"); + assert_eq!(batches.len(), 1); + let text = String::from_utf8(batches[0].clone()).expect("utf8"); + let first = text.find("first").expect("first entry"); + let second = text.find("second").expect("second entry"); + assert!( + first < second, + "entries keep insertion order across both insert paths" + ); + } + + /// Default policy: an oversized pre-serialized entry is dropped with one warning. + #[tokio::test(flavor = "current_thread")] + async fn test_insert_serialized_batch_drops_oversized_entry_by_default() { + let (logged, _guard) = LevelCounter::install(); + let (service, handle) = AggregatorService::new(); + let _task = tokio::spawn(service.run()); + handle + .insert_serialized_batch(vec![json_string_of_len( + crate::constants::MAX_LOG_BYTES + 1, + )]) + .expect("insert"); + assert!(handle.get_batches().await.expect("batches").is_empty()); + assert_eq!((logged.warnings(), logged.errors()), (1, 0)); + } + + /// With `SendWithWarning`, each oversized entry that is kept logs exactly one warning. + #[tokio::test(flavor = "current_thread")] + async fn test_send_with_warning_warns_once_per_kept_oversized_entry() { + let (logged, _guard) = LevelCounter::install(); + let (service, handle) = AggregatorService::new(); + let service = service.with_oversize_policy(OversizePolicy::SendWithWarning); + let _task = tokio::spawn(service.run()); + let big = json_string_of_len(2 * crate::constants::MAX_LOG_BYTES); + handle + .insert_serialized_batch(vec![big.clone(), big]) + .expect("insert"); + assert!(!handle.get_batches().await.expect("batches").is_empty()); + assert_eq!((logged.warnings(), logged.errors()), (2, 0)); + } + + /// With `SendWithWarning`, an entry too big for any payload is dropped with one warning. + #[tokio::test(flavor = "current_thread")] + async fn test_send_with_warning_warns_once_per_entry_too_big_for_any_payload() { + let (logged, _guard) = LevelCounter::install(); + let (service, handle) = AggregatorService::new(); + let service = service.with_oversize_policy(OversizePolicy::SendWithWarning); + let _task = tokio::spawn(service.run()); + let too_big = json_string_of_len(crate::constants::MAX_CONTENT_BYTES); + handle + .insert_serialized_batch(vec![too_big]) + .expect("insert"); + assert!(handle.get_batches().await.expect("batches").is_empty()); + assert_eq!((logged.warnings(), logged.errors()), (1, 0)); + } } diff --git a/crates/datadog-logs-agent/src/flusher.rs b/crates/datadog-logs-agent/src/flusher.rs index 45e063fa..a9d038f6 100644 --- a/crates/datadog-logs-agent/src/flusher.rs +++ b/crates/datadog-logs-agent/src/flusher.rs @@ -284,1208 +284,4 @@ fn compress_zstd(data: &[u8], level: i32) -> Result, FlushError> { #[cfg(test)] // Tests use plain reqwest client to connect to local mock server #[allow(clippy::disallowed_methods)] -mod tests { - use super::*; - use crate::aggregator::AggregatorService; - use crate::config::{Destination, LogFlusherConfig}; - use crate::intake_entry::IntakeEntry; - use crate::logs_additional_endpoint::LogsAdditionalEndpoint; - use mockito::Matcher; - use std::sync::Arc; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::time::Duration; - - fn make_entry(msg: &str) -> IntakeEntry { - IntakeEntry::from_message(msg, 1_700_000_000_000) - } - - fn config_for_mock(mock_url: &str) -> LogFlusherConfig { - // Use OPW mode pointing at the mock server to avoid HTTPS - LogFlusherConfig { - api_key: "test-api-key".to_string(), - site: "datadoghq.com".to_string(), - mode: Destination::ObservabilityPipelinesWorker { - url: format!("{mock_url}/api/v2/logs"), - }, - additional_endpoints: Vec::new(), - use_compression: false, - compression_level: 3, - flush_timeout: Duration::from_secs(5), - } - } - - #[tokio::test] - async fn test_flush_empty_aggregator_does_not_call_api() { - let (service, handle) = AggregatorService::new(); - let _task = tokio::spawn(service.run()); - - // Server with no routes — any request would cause test failure - let mock_server = mockito::Server::new_async().await; - let config = config_for_mock(&mock_server.url()); - let client = reqwest::Client::builder().build().expect("client"); - let flusher = LogFlusher::new(config, client, handle); - - assert!( - flusher.flush(vec![]).await.is_empty(), - "empty flush should succeed" - ); - // No mock assertions needed — absence of request is the assertion - } - - #[tokio::test] - async fn test_flush_sends_post_with_api_key_header() { - // Verify that Datadog mode sends both DD-API-KEY and DD-PROTOCOL: - // agent-json headers. We call ship_batch directly to bypass - // resolve_endpoint (which builds an HTTPS URL incompatible with the - // HTTP mock server). - let (service, handle) = AggregatorService::new(); - let _task = tokio::spawn(service.run()); - - let mut mock_server = mockito::Server::new_async().await; - let mock = mock_server - .mock("POST", "/api/v2/logs") - .match_header("DD-API-KEY", "test-api-key") - .match_header("DD-PROTOCOL", "agent-json") - .with_status(202) - .create_async() - .await; - - let config = LogFlusherConfig { - api_key: "test-api-key".to_string(), - site: "datadoghq.com".to_string(), - mode: Destination::Datadog, - additional_endpoints: Vec::new(), - use_compression: false, - compression_level: 3, - flush_timeout: Duration::from_secs(5), - }; - - let client = reqwest::Client::builder().build().expect("client"); - let flusher = LogFlusher::new(config, client, handle); - - // Call ship_batch directly to use the mock server's HTTP URL instead - // of the HTTPS URL that resolve_endpoint would produce. - let url = format!("{}/api/v2/logs", mock_server.url()); - let batch = b"[{\"message\":\"test\"}]"; - flusher - .ship_batch(batch, &url, false, "test-api-key", true) - .await - .expect("ship_batch should succeed"); - - mock.assert_async().await; - } - - #[tokio::test] - async fn test_flush_opw_mode_omits_dd_protocol_header() { - let (service, handle) = AggregatorService::new(); - let _task = tokio::spawn(service.run()); - - let mut mock_server = mockito::Server::new_async().await; - let opw_url = format!("{}/logs", mock_server.url()); - - // Verify DD-PROTOCOL is NOT present in OPW requests - let mock = mock_server - .mock("POST", "/logs") - .match_header("DD-API-KEY", "test-api-key") - .match_header("DD-PROTOCOL", Matcher::Missing) - .with_status(200) - .expect(1) - .create_async() - .await; - - let config = LogFlusherConfig { - api_key: "test-api-key".to_string(), - site: "unused".to_string(), - mode: Destination::ObservabilityPipelinesWorker { url: opw_url }, - additional_endpoints: Vec::new(), - use_compression: false, - compression_level: 3, - flush_timeout: Duration::from_secs(5), - }; - - let client = reqwest::Client::builder().build().expect("client"); - let flusher = LogFlusher::new(config, client, handle.clone()); - - handle - .insert_batch(vec![make_entry("opw log")]) - .expect("insert"); - let result = flusher.flush(vec![]).await; - assert!(result.is_empty(), "OPW flush should return empty on 200"); - mock.assert_async().await; - } - - /// When the primary destination is OPW, additional endpoints are still Datadog - /// intakes and must receive the `DD-PROTOCOL: agent-json` header. - #[tokio::test] - async fn test_opw_primary_additional_endpoint_receives_dd_protocol_header() { - 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; - - // OPW primary must NOT have DD-PROTOCOL header - let _primary_mock = primary - .mock("POST", "/logs") - .match_header("DD-PROTOCOL", Matcher::Missing) - .with_status(200) - .expect(1) - .create_async() - .await; - - // Additional endpoint (Datadog intake) MUST have DD-PROTOCOL header - let extra_mock = extra - .mock("POST", "/extra") - .match_header("DD-PROTOCOL", "agent-json") - .with_status(202) - .expect(1) - .create_async() - .await; - - let config = LogFlusherConfig { - api_key: "key".to_string(), - site: "datadoghq.com".to_string(), - mode: Destination::ObservabilityPipelinesWorker { - url: format!("{}/logs", primary.url()), - }, - additional_endpoints: vec![LogsAdditionalEndpoint { - api_key: "extra-key".to_string(), - url: format!("{}/extra", extra.url()), - is_reliable: true, - }], - use_compression: false, - compression_level: 3, - flush_timeout: Duration::from_secs(5), - }; - - let client = reqwest::Client::builder().build().expect("client"); - let flusher = LogFlusher::new(config, client, handle.clone()); - handle - .insert_batch(vec![make_entry("test")]) - .expect("insert"); - - assert!(flusher.flush(vec![]).await.is_empty()); - extra_mock.assert_async().await; - } - - #[tokio::test] - async fn test_flush_does_not_retry_on_403() { - let (service, handle) = AggregatorService::new(); - let _task = tokio::spawn(service.run()); - - let mut mock_server = mockito::Server::new_async().await; - // expect(1) means exactly one call — if retried, the test will fail - let mock = mock_server - .mock("POST", "/api/v2/logs") - .with_status(403) - .expect(1) - .create_async() - .await; - - let config = config_for_mock(&mock_server.url()); - let client = reqwest::Client::builder().build().expect("client"); - let flusher = LogFlusher::new(config, client, handle.clone()); - - handle - .insert_batch(vec![make_entry("log")]) - .expect("insert"); - let result = flusher.flush(vec![]).await; - // 403 is a permanent error — the batch is dropped, no builder to retry. - assert!( - result.is_empty(), - "403 is a permanent error; no builder to retry" - ); - mock.assert_async().await; - } - - /// 429 (Too Many Requests) is a retryable 4xx — the retry loop must - /// continue rather than short-circuiting with a permanent failure. - #[tokio::test] - async fn test_flush_retries_on_429_then_succeeds() { - let (service, handle) = AggregatorService::new(); - let _task = tokio::spawn(service.run()); - - let mut mock_server = mockito::Server::new_async().await; - // First call → 429, second call → 200 - let _throttled = mock_server - .mock("POST", "/api/v2/logs") - .with_status(429) - .expect(1) - .create_async() - .await; - let _ok = mock_server - .mock("POST", "/api/v2/logs") - .with_status(200) - .expect(1) - .create_async() - .await; - - let config = config_for_mock(&mock_server.url()); - let client = reqwest::Client::builder().build().expect("client"); - let flusher = LogFlusher::new(config, client, handle.clone()); - - handle - .insert_batch(vec![make_entry("throttled log")]) - .expect("insert"); - let result = flusher.flush(vec![]).await; - assert!(result.is_empty(), "should succeed after 429 retry"); - } - - #[tokio::test] - async fn test_flush_retries_on_5xx_then_succeeds() { - let (service, handle) = AggregatorService::new(); - let _task = tokio::spawn(service.run()); - - let mut mock_server = mockito::Server::new_async().await; - // First call → 500, second call → 202 - let _fail_mock = mock_server - .mock("POST", "/api/v2/logs") - .with_status(500) - .expect(1) - .create_async() - .await; - let _ok_mock = mock_server - .mock("POST", "/api/v2/logs") - .with_status(202) - .expect(1) - .create_async() - .await; - - let config = config_for_mock(&mock_server.url()); - let client = reqwest::Client::builder().build().expect("client"); - let flusher = LogFlusher::new(config, client, handle.clone()); - - handle - .insert_batch(vec![make_entry("log")]) - .expect("insert"); - let result = flusher.flush(vec![]).await; - assert!(result.is_empty(), "should succeed on second attempt"); - } - - /// All additional endpoints receive the same batch when flush() is called. - #[tokio::test] - async fn test_additional_endpoints_all_receive_batch() { - let (service, handle) = AggregatorService::new(); - let _task = tokio::spawn(service.run()); - - let mut primary = mockito::Server::new_async().await; - let mut extra1 = mockito::Server::new_async().await; - let mut extra2 = mockito::Server::new_async().await; - - let primary_mock = primary - .mock("POST", "/api/v2/logs") - .with_status(202) - .expect(1) - .create_async() - .await; - let extra1_mock = extra1 - .mock("POST", "/extra") - .with_status(200) - .expect(1) - .create_async() - .await; - let extra2_mock = extra2 - .mock("POST", "/extra") - .with_status(200) - .expect(1) - .create_async() - .await; - - let config = LogFlusherConfig { - api_key: "key".to_string(), - site: "datadoghq.com".to_string(), - mode: Destination::ObservabilityPipelinesWorker { - url: format!("{}/api/v2/logs", primary.url()), - }, - additional_endpoints: vec![ - LogsAdditionalEndpoint { - api_key: "extra-key-1".to_string(), - url: format!("{}/extra", extra1.url()), - is_reliable: true, - }, - LogsAdditionalEndpoint { - api_key: "extra-key-2".to_string(), - url: format!("{}/extra", extra2.url()), - is_reliable: true, - }, - ], - use_compression: false, - compression_level: 3, - flush_timeout: Duration::from_secs(5), - }; - - let client = reqwest::Client::builder().build().expect("client"); - let flusher = LogFlusher::new(config, client, handle.clone()); - handle.insert_batch(vec![make_entry("hi")]).expect("insert"); - - assert!(flusher.flush(vec![]).await.is_empty()); - primary_mock.assert_async().await; - extra1_mock.assert_async().await; - extra2_mock.assert_async().await; - } - - /// Additional endpoints are dispatched concurrently: if they were sequential, - /// two endpoints each waiting at a Barrier(2) would deadlock — only concurrent - /// dispatch lets both handlers reach the barrier simultaneously. - #[tokio::test] - async fn test_additional_endpoints_dispatched_concurrently() { - use std::sync::Arc; - use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; - use tokio::net::TcpListener; - use tokio::sync::Barrier; - - let (service, handle) = AggregatorService::new(); - let _task = tokio::spawn(service.run()); - - let barrier = Arc::new(Barrier::new(2)); - - // Spawn a minimal HTTP server that waits at the barrier before - // responding, so both must be in-flight at the same time to complete. - async fn serve_once(barrier: Arc) -> String { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - tokio::spawn(async move { - let (mut stream, _) = listener.accept().await.unwrap(); - let mut buf = vec![0u8; 4096]; - let _ = stream.read(&mut buf).await; - barrier.wait().await; - let _ = stream - .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n") - .await; - }); - format!("http://127.0.0.1:{}/logs", addr.port()) - } - - let url1 = serve_once(barrier.clone()).await; - let url2 = serve_once(barrier.clone()).await; - - let mut primary = mockito::Server::new_async().await; - let _primary_mock = primary - .mock("POST", "/api/v2/logs") - .with_status(202) - .create_async() - .await; - - let config = LogFlusherConfig { - api_key: "key".to_string(), - site: "datadoghq.com".to_string(), - mode: Destination::ObservabilityPipelinesWorker { - url: format!("{}/api/v2/logs", primary.url()), - }, - additional_endpoints: vec![ - LogsAdditionalEndpoint { - api_key: "extra-key-1".to_string(), - url: url1, - is_reliable: true, - }, - LogsAdditionalEndpoint { - api_key: "extra-key-2".to_string(), - url: url2, - is_reliable: true, - }, - ], - use_compression: false, - compression_level: 3, - flush_timeout: Duration::from_secs(5), - }; - - let client = reqwest::Client::builder().build().expect("client"); - let flusher = LogFlusher::new(config, client, handle.clone()); - handle - .insert_batch(vec![make_entry("concurrent")]) - .expect("insert"); - - assert!(flusher.flush(vec![]).await.is_empty()); - } - - /// A builder returned by `flush` can be redriven on the next call. - /// - /// The mock fails on the first 3 attempts (exhausting the per-invocation - /// retry budget), then succeeds on the 4th attempt (the next invocation). - /// This proves the cross-invocation retry path end-to-end. - #[tokio::test] - async fn test_cross_invocation_retry_delivers_on_redrive() { - let (service, handle) = AggregatorService::new(); - let _task = tokio::spawn(service.run()); - - let mut mock_server = mockito::Server::new_async().await; - // First 3 calls: transient 503 → exhausts per-invocation retry budget - let _fail_mock = mock_server - .mock("POST", "/api/v2/logs") - .with_status(503) - .expect(3) - .create_async() - .await; - // 4th call: redriven on the next flush → succeeds - let _ok_mock = mock_server - .mock("POST", "/api/v2/logs") - .with_status(200) - .expect(1) - .create_async() - .await; - - let config = config_for_mock(&mock_server.url()); - let client = reqwest::Client::builder().build().expect("client"); - let flusher = LogFlusher::new(config, client, handle.clone()); - handle - .insert_batch(vec![make_entry("retry-me")]) - .expect("insert"); - - // First flush: all 3 attempts fail → returns the builder for retry. - let failed = flusher.flush(vec![]).await; - assert_eq!(failed.len(), 1, "one builder should be returned for retry"); - - // Second flush: aggregator is empty; redrives the failed builder → succeeds. - let result = flusher.flush(failed).await; - assert!( - result.is_empty(), - "redriven builder should succeed on the next invocation" - ); - } - - /// When the primary is OPW (compression disabled for OPW transport) and - /// `use_compression` is true, additional Datadog endpoints must still - /// receive `Content-Encoding: zstd` — they are Datadog intakes and honour - /// `DD_LOGS_CONFIG_USE_COMPRESSION` independently of the primary. - #[tokio::test] - async fn test_opw_primary_additional_endpoint_compresses_when_enabled() { - 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; - - // OPW primary must NOT have Content-Encoding - let _primary_mock = primary - .mock("POST", "/logs") - .match_header("Content-Encoding", Matcher::Missing) - .with_status(200) - .expect(1) - .create_async() - .await; - - // Additional Datadog endpoint MUST receive Content-Encoding: zstd - let extra_mock = extra - .mock("POST", "/extra") - .match_header("Content-Encoding", "zstd") - .with_status(202) - .expect(1) - .create_async() - .await; - - let config = LogFlusherConfig { - api_key: "key".to_string(), - site: "datadoghq.com".to_string(), - mode: Destination::ObservabilityPipelinesWorker { - url: format!("{}/logs", primary.url()), - }, - additional_endpoints: vec![LogsAdditionalEndpoint { - api_key: "extra-key".to_string(), - url: format!("{}/extra", extra.url()), - is_reliable: true, - }], - use_compression: true, - compression_level: 3, - flush_timeout: Duration::from_secs(5), - }; - - let client = reqwest::Client::builder().build().expect("client"); - let flusher = LogFlusher::new(config, client, handle.clone()); - handle - .insert_batch(vec![make_entry("compressed extra")]) - .expect("insert"); - - assert!(flusher.flush(vec![]).await.is_empty()); - extra_mock.assert_async().await; - } - - /// Even when `use_compression: true`, the OPW primary endpoint must never - /// receive a `Content-Encoding` header — OPW does not support zstd. - #[tokio::test] - async fn test_opw_primary_never_compressed_even_when_flag_set() { - let (service, handle) = AggregatorService::new(); - let _task = tokio::spawn(service.run()); - - let mut mock_server = mockito::Server::new_async().await; - let mock = mock_server - .mock("POST", "/logs") - .match_header("Content-Encoding", Matcher::Missing) - .with_status(200) - .expect(1) - .create_async() - .await; - - let config = LogFlusherConfig { - api_key: "key".to_string(), - site: "datadoghq.com".to_string(), - mode: Destination::ObservabilityPipelinesWorker { - url: format!("{}/logs", mock_server.url()), - }, - additional_endpoints: vec![], - use_compression: true, // flag set but must not reach OPW - compression_level: 3, - flush_timeout: Duration::from_secs(5), - }; - - let client = reqwest::Client::builder().build().expect("client"); - let flusher = LogFlusher::new(config, client, handle.clone()); - handle - .insert_batch(vec![make_entry("opw no compress")]) - .expect("insert"); - - assert!(flusher.flush(vec![]).await.is_empty()); - mock.assert_async().await; - } - - /// Additional-endpoint failures are best-effort: their builders are NOT - /// included in the returned retry vec, even when they exhaust all retries. - /// Only primary-endpoint failures are tracked for cross-invocation retry. - #[tokio::test] - async fn test_additional_endpoint_failures_not_tracked_for_retry() { - 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") - .with_status(200) - .expect(1) - .create_async() - .await; - // Additional endpoint always returns 503 — exhausts per-invocation retries. - let _extra_mock = extra - .mock("POST", "/extra") - .with_status(503) - .expect(3) // MAX_FLUSH_ATTEMPTS - .create_async() - .await; - - let config = LogFlusherConfig { - api_key: "key".to_string(), - site: "datadoghq.com".to_string(), - mode: Destination::ObservabilityPipelinesWorker { - url: format!("{}/api/v2/logs", primary.url()), - }, - additional_endpoints: vec![LogsAdditionalEndpoint { - api_key: "extra-key".to_string(), - url: format!("{}/extra", extra.url()), - is_reliable: true, - }], - use_compression: false, - compression_level: 3, - flush_timeout: Duration::from_secs(5), - }; - - let client = reqwest::Client::builder().build().expect("client"); - let flusher = LogFlusher::new(config, client, handle.clone()); - handle - .insert_batch(vec![make_entry("test")]) - .expect("insert"); - - let result = flusher.flush(vec![]).await; - assert!( - result.is_empty(), - "additional-endpoint failures are best-effort and must not be tracked for retry" - ); - } - - /// Default policy: client errors other than 408/425/429 are dropped after one attempt. - #[tokio::test] - async fn test_default_policy_drops_permanent_client_errors_after_one_attempt() { - for status in [400_usize, 401, 404, 413] { - 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(status) - .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()); - handle - .insert_batch(vec![make_entry("log")]) - .expect("insert"); - - assert!( - flusher.flush(vec![]).await.is_empty(), - "status {status} must be dropped, not handed back" - ); - mock.assert_async().await; - } - } - - /// Default policy: 408, 425 and 429 are retried like server errors, then handed back. - #[tokio::test] - async fn test_default_policy_retries_408_425_and_429() { - for status in [408_usize, 425, 429] { - 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(status) - .expect(3) - .create_async() - .await; - let client = reqwest::Client::builder().build().expect("client"); - let flusher = LogFlusher::new(config_for_mock(&server.url()), client, handle.clone()); - handle - .insert_batch(vec![make_entry("log")]) - .expect("insert"); - - assert_eq!( - flusher.flush(vec![]).await.len(), - 1, - "status {status} is handed back" - ); - mock.assert_async().await; - } - } - - /// Connection errors are retried and then handed back. - #[tokio::test] - async fn test_default_policy_hands_back_connection_errors() { - let (service, handle) = AggregatorService::new(); - let _task = tokio::spawn(service.run()); - let listener = tokio::net::TcpListener::bind("127.0.0.1:0") - .await - .expect("bind"); - let port = listener.local_addr().expect("addr").port(); - let accepts = Arc::new(AtomicUsize::new(0)); - let accepts_in_task = Arc::clone(&accepts); - let _server = tokio::spawn(async move { - while let Ok((stream, _)) = listener.accept().await { - accepts_in_task.fetch_add(1, Ordering::SeqCst); - drop(stream); - } - }); - let client = reqwest::Client::builder().build().expect("client"); - let flusher = LogFlusher::new( - config_for_mock(&format!("http://127.0.0.1:{port}")), - client, - handle.clone(), - ); - handle - .insert_batch(vec![make_entry("log")]) - .expect("insert"); - - assert_eq!(flusher.flush(vec![]).await.len(), 1); - assert_eq!(accepts.load(Ordering::SeqCst), 3); - } - - /// A redriven request that fails again gets a fresh retry budget and is handed back again. - #[tokio::test] - async fn test_redriven_request_that_fails_again_is_handed_back_again() { - 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(503) - .expect(6) - .create_async() - .await; - let client = reqwest::Client::builder().build().expect("client"); - let flusher = LogFlusher::new(config_for_mock(&server.url()), client, handle.clone()); - handle - .insert_batch(vec![make_entry("log")]) - .expect("insert"); - - let failed = flusher.flush(vec![]).await; - assert_eq!(failed.len(), 1); - let failed_again = flusher.flush(failed).await; - assert_eq!(failed_again.len(), 1); - mock.assert_async().await; - } - - /// Sends one batch in Datadog mode to a listener that drops every connection, so the - /// request is handed back and can be inspected. - async fn handed_back_datadog_request(use_compression: bool) -> reqwest::Request { - let listener = tokio::net::TcpListener::bind("127.0.0.1:0") - .await - .expect("bind"); - let addr = listener.local_addr().expect("addr"); - let _server = tokio::spawn(async move { - while let Ok((stream, _)) = listener.accept().await { - drop(stream); - } - }); - let (service, handle) = AggregatorService::new(); - let _task = tokio::spawn(service.run()); - let config = LogFlusherConfig { - site: "datadoghq.eu".to_string(), - mode: Destination::Datadog, - use_compression, - ..config_for_mock("http://unused") - }; - // Route the public intake hostname to the local listener. - let client = reqwest::Client::builder() - .resolve("http-intake.logs.datadoghq.eu", addr) - .build() - .expect("client"); - let flusher = LogFlusher::new(config, client, handle.clone()); - handle - .insert_batch(vec![make_entry("default path")]) - .expect("insert"); - - let mut failed = flusher.flush(vec![]).await; - assert_eq!(failed.len(), 1); - failed.remove(0).build().expect("request") - } - - /// Datadog mode by default sends compressed batches to the site's public intake with the - /// Datadog-intake headers. - #[tokio::test] - async fn test_default_datadog_request_when_compressed() { - let req = handed_back_datadog_request(true).await; - - assert_eq!( - req.url().as_str(), - "https://http-intake.logs.datadoghq.eu/api/v2/logs" - ); - assert_eq!( - req.timeout(), - Some(&config_for_mock("http://unused").flush_timeout) - ); - let headers = req.headers(); - assert_eq!(headers["DD-API-KEY"], "test-api-key"); - assert_eq!(headers["DD-PROTOCOL"], "agent-json"); - assert_eq!(headers["Content-Type"], "application/json"); - assert_eq!(headers["Content-Encoding"], "zstd"); - let body = - zstd::decode_all(req.body().expect("body").as_bytes().expect("bytes")).expect("zstd"); - let json: serde_json::Value = serde_json::from_slice(&body).expect("json"); - assert_eq!(json[0]["message"], "default path"); - } - - /// With compression off, Datadog mode sends the raw batch without Content-Encoding. - #[tokio::test] - async fn test_default_datadog_request_when_uncompressed() { - let req = handed_back_datadog_request(false).await; - - assert_eq!( - req.url().as_str(), - "https://http-intake.logs.datadoghq.eu/api/v2/logs" - ); - assert_eq!( - req.timeout(), - Some(&config_for_mock("http://unused").flush_timeout) - ); - let headers = req.headers(); - assert_eq!(headers["DD-API-KEY"], "test-api-key"); - assert_eq!(headers["DD-PROTOCOL"], "agent-json"); - assert_eq!(headers["Content-Type"], "application/json"); - assert!(headers.get("Content-Encoding").is_none()); - let json: serde_json::Value = - serde_json::from_slice(req.body().expect("body").as_bytes().expect("bytes")) - .expect("json"); - assert_eq!(json[0]["message"], "default path"); - } - - /// One-shot HTTP server. Records ":received" when the full request has arrived and - /// ":responded" just before answering 200 after `delay`; hands back the body. - async fn capture_once( - name: &'static str, - delay: Duration, - events: Arc>>, - ) -> (String, tokio::sync::oneshot::Receiver>) { - use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; - - let listener = tokio::net::TcpListener::bind("127.0.0.1:0") - .await - .expect("bind"); - let port = listener.local_addr().expect("addr").port(); - let (tx, rx) = tokio::sync::oneshot::channel(); - tokio::spawn(async move { - let (mut stream, _) = listener.accept().await.expect("accept"); - let mut raw = Vec::new(); - let mut buf = [0_u8; 8192]; - let body = loop { - let n = stream.read(&mut buf).await.expect("read"); - assert!(n > 0, "connection closed before the request was complete"); - raw.extend_from_slice(&buf[..n]); - let Some(end) = raw.windows(4).position(|w| w == b"\r\n\r\n") else { - continue; - }; - let head = String::from_utf8_lossy(&raw[..end]).to_ascii_lowercase(); - let len: usize = head - .lines() - .find_map(|l| l.strip_prefix("content-length:")) - .expect("request has content-length") - .trim() - .parse() - .expect("content-length"); - while raw.len() < end + 4 + len { - let n = stream.read(&mut buf).await.expect("read body"); - assert!(n > 0, "connection closed before the body was complete"); - raw.extend_from_slice(&buf[..n]); - } - break raw[end + 4..end + 4 + len].to_vec(); - }; - events - .lock() - .expect("lock") - .push(format!("{name}:received")); - tokio::time::sleep(delay).await; - let _ = tx.send(body); - events - .lock() - .expect("lock") - .push(format!("{name}:responded")); - let _ = stream - .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n") - .await; - }); - (format!("http://127.0.0.1:{port}/logs"), rx) - } - - /// Default policy: per batch, additional endpoints are sent only after the primary - /// request has finished. - #[tokio::test] - async fn test_default_policy_sends_additional_endpoints_after_primary() { - let events = Arc::new(std::sync::Mutex::new(Vec::new())); - let (primary_url, _primary_body) = - capture_once("primary", Duration::from_millis(200), events.clone()).await; - let (extra_url, _extra_body) = capture_once("extra", Duration::ZERO, events.clone()).await; - - let (service, handle) = AggregatorService::new(); - let _task = tokio::spawn(service.run()); - let config = LogFlusherConfig { - mode: Destination::ObservabilityPipelinesWorker { url: primary_url }, - additional_endpoints: vec![LogsAdditionalEndpoint { - api_key: "extra-key".to_string(), - url: extra_url, - is_reliable: true, - }], - ..config_for_mock("http://unused") - }; - let client = reqwest::Client::builder().build().expect("client"); - let flusher = LogFlusher::new(config, client, handle.clone()); - handle - .insert_batch(vec![make_entry("ordered")]) - .expect("insert"); - - assert!(flusher.flush(vec![]).await.is_empty()); - assert_eq!( - *events.lock().expect("lock"), - vec![ - "primary:received", - "primary:responded", - "extra:received", - "extra:responded" - ] - ); - } - - /// Default policy: when the primary fails, additional endpoints are still sent with their - /// own API key, and only the primary request is handed back. - #[tokio::test] - async fn test_default_extras_still_sent_with_their_own_key_when_primary_fails() { - 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") - .match_header("DD-API-KEY", "test-api-key") - .with_status(503) - .expect(3) - .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()); - handle - .insert_batch(vec![make_entry("log")]) - .expect("insert"); - - let failed = flusher.flush(vec![]).await; - assert_eq!(failed.len(), 1); - let req = failed - .into_iter() - .next() - .expect("handed-back request") - .build() - .expect("request"); - assert_eq!(req.url().as_str(), format!("{}/api/v2/logs", primary.url())); - assert_eq!(req.headers()["DD-API-KEY"], "test-api-key"); - primary_mock.assert_async().await; - extra_mock.assert_async().await; - } - - /// Every additional endpoint receives byte-identical compressed bodies that decompress to - /// exactly what the primary received. - #[tokio::test] - async fn test_additional_endpoints_receive_identical_compressed_bodies() { - let events = Arc::new(std::sync::Mutex::new(Vec::new())); - let (primary_url, mut primary_body) = - capture_once("primary", Duration::ZERO, events.clone()).await; - let (url_1, mut body_1) = capture_once("extra1", Duration::ZERO, events.clone()).await; - let (url_2, mut body_2) = capture_once("extra2", Duration::ZERO, events.clone()).await; - - let (service, handle) = AggregatorService::new(); - let _task = tokio::spawn(service.run()); - let config = LogFlusherConfig { - mode: Destination::ObservabilityPipelinesWorker { url: primary_url }, - additional_endpoints: vec![ - LogsAdditionalEndpoint { - api_key: "extra-key-1".to_string(), - url: url_1, - is_reliable: true, - }, - LogsAdditionalEndpoint { - api_key: "extra-key-2".to_string(), - url: url_2, - is_reliable: true, - }, - ], - use_compression: true, - ..config_for_mock("http://unused") - }; - let compression_level = config.compression_level; - let client = reqwest::Client::builder().build().expect("client"); - let flusher = LogFlusher::new(config, client, handle.clone()); - handle - .insert_batch( - // Varied enough that zstd levels 3 and 4 produce different bytes. - (0..200_u64) - .map(|i| { - let words: Vec = (0..40) - .map(|j| format!("{:x}", (i * 40 + j) * 2_654_435_761 % 1_000_003)) - .collect(); - make_entry(&format!("same bytes {i} {}", words.join(" "))) - }) - .collect(), - ) - .expect("insert"); - - assert!(flusher.flush(vec![]).await.is_empty()); - let primary_body = primary_body.try_recv().expect("primary received a request"); - let body_1 = body_1.try_recv().expect("extra1 received a request"); - let body_2 = body_2.try_recv().expect("extra2 received a request"); - assert_eq!(body_1, body_2); - assert_eq!( - body_1, - zstd::encode_all(primary_body.as_slice(), compression_level).expect("zstd"), - "additional endpoints use the configured compression level" - ); - let decompressed = zstd::decode_all(body_1.as_slice()).expect("zstd"); - assert_eq!(decompressed, primary_body); - let decoded: serde_json::Value = serde_json::from_slice(&decompressed).expect("json"); - assert_eq!(decoded.as_array().expect("array").len(), 200); - } - - /// Counts `warn!` and `error!` events emitted by this crate on the current thread. - #[derive(Clone, Default)] - struct LevelCounter { - warnings: Arc, - errors: Arc, - } - - impl tracing_subscriber::Layer for LevelCounter { - fn on_event( - &self, - event: &tracing::Event<'_>, - _ctx: tracing_subscriber::layer::Context<'_, S>, - ) { - let meta = event.metadata(); - if !meta.target().starts_with("datadog_logs_agent") { - return; - } - match *meta.level() { - tracing::Level::WARN => self.warnings.fetch_add(1, Ordering::SeqCst), - tracing::Level::ERROR => self.errors.fetch_add(1, Ordering::SeqCst), - _ => return, - }; - } - } - - impl LevelCounter { - /// Installs the counter for this thread. Use the current-thread runtime so tasks - /// spawned by the test run on this thread. - fn install() -> (Self, tracing::subscriber::DefaultGuard) { - use tracing_subscriber::layer::SubscriberExt as _; - let counter = Self::default(); - let guard = tracing::subscriber::set_default( - tracing_subscriber::registry().with(counter.clone()), - ); - (counter, guard) - } - - fn warnings(&self) -> usize { - self.warnings.load(Ordering::SeqCst) - } - - fn errors(&self) -> usize { - self.errors.load(Ordering::SeqCst) - } - } - - /// Default: a request that keeps failing logs one warning per attempt plus a final one. - #[tokio::test(flavor = "current_thread")] - async fn test_default_logging_warns_on_every_failed_attempt() { - let (logged, _guard) = LevelCounter::install(); - 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(503) - .expect(3) - .create_async() - .await; - let client = reqwest::Client::builder().build().expect("client"); - let flusher = LogFlusher::new(config_for_mock(&server.url()), client, handle.clone()); - handle - .insert_batch(vec![make_entry("log")]) - .expect("insert"); - - assert_eq!(flusher.flush(vec![]).await.len(), 1); - mock.assert_async().await; - assert_eq!( - logged.warnings(), - 4, - "3 transient-error warnings + 1 gave-up warning" - ); - assert_eq!(logged.errors(), 0, "default logging uses warn!, not error!"); - } - - /// Default: a connection that keeps dropping logs one warning per attempt plus a final one. - #[tokio::test(flavor = "current_thread")] - async fn test_default_logging_warns_on_every_failed_network_attempt() { - let (logged, _guard) = LevelCounter::install(); - let (service, handle) = AggregatorService::new(); - let _task = tokio::spawn(service.run()); - let listener = tokio::net::TcpListener::bind("127.0.0.1:0") - .await - .expect("bind"); - let port = listener.local_addr().expect("addr").port(); - let _server = tokio::spawn(async move { - while let Ok((stream, _)) = listener.accept().await { - drop(stream); - } - }); - let client = reqwest::Client::builder().build().expect("client"); - let flusher = LogFlusher::new( - config_for_mock(&format!("http://127.0.0.1:{port}")), - client, - handle.clone(), - ); - handle - .insert_batch(vec![make_entry("log")]) - .expect("insert"); - - assert_eq!(flusher.flush(vec![]).await.len(), 1); - assert_eq!( - logged.warnings(), - 4, - "3 network-error warnings + 1 gave-up warning" - ); - assert_eq!(logged.errors(), 0, "default logging uses warn!, not error!"); - } - - /// Default: a permanent client error logs exactly one warning. - #[tokio::test(flavor = "current_thread")] - async fn test_default_logging_warns_once_on_permanent_error() { - let (logged, _guard) = LevelCounter::install(); - 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(404) - .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()); - handle - .insert_batch(vec![make_entry("log")]) - .expect("insert"); - - assert!(flusher.flush(vec![]).await.is_empty()); - mock.assert_async().await; - assert_eq!(logged.warnings(), 1); - assert_eq!(logged.errors(), 0, "default logging uses warn!, not error!"); - } - - /// Default: an additional endpoint that keeps failing adds one more warning on top of - /// the per-attempt ones. - #[tokio::test(flavor = "current_thread")] - async fn test_default_logging_for_failing_additional_endpoint() { - 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") - .with_status(202) - .expect(1) - .create_async() - .await; - let extra_mock = extra - .mock("POST", "/extra") - .with_status(503) - .expect(3) - .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()); - 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(), - 5, - "3 transient + 1 gave-up + 1 additional-endpoint warning" - ); - assert_eq!(logged.errors(), 0, "default logging uses warn!, not error!"); - } - - /// Default: each entry over the 1 MiB limit is dropped with exactly one warning. - #[tokio::test(flavor = "current_thread")] - async fn test_default_logging_warns_once_per_dropped_oversized_entry() { - let (logged, _guard) = LevelCounter::install(); - let (service, handle) = AggregatorService::new(); - let _task = tokio::spawn(service.run()); - - let oversized = "a".repeat(crate::constants::MAX_LOG_BYTES + 1); - handle - .insert_batch(vec![make_entry(&oversized), make_entry(&oversized)]) - .expect("insert"); - assert!(handle.get_batches().await.expect("batches").is_empty()); - assert_eq!(logged.warnings(), 2); - assert_eq!(logged.errors(), 0, "default logging uses warn!, not error!"); - } -} +mod tests; diff --git a/crates/datadog-logs-agent/src/flusher/tests.rs b/crates/datadog-logs-agent/src/flusher/tests.rs new file mode 100644 index 00000000..c75bb251 --- /dev/null +++ b/crates/datadog-logs-agent/src/flusher/tests.rs @@ -0,0 +1,1160 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +use super::*; +use crate::aggregator::AggregatorService; +use crate::config::{Destination, LogFlusherConfig}; +use crate::intake_entry::IntakeEntry; +use crate::logs_additional_endpoint::LogsAdditionalEndpoint; +use crate::test_support::LevelCounter; +use mockito::Matcher; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::Duration; + +fn make_entry(msg: &str) -> IntakeEntry { + IntakeEntry::from_message(msg, 1_700_000_000_000) +} + +fn config_for_mock(mock_url: &str) -> LogFlusherConfig { + // Use OPW mode pointing at the mock server to avoid HTTPS + LogFlusherConfig { + api_key: "test-api-key".to_string(), + site: "datadoghq.com".to_string(), + mode: Destination::ObservabilityPipelinesWorker { + url: format!("{mock_url}/api/v2/logs"), + }, + additional_endpoints: Vec::new(), + use_compression: false, + compression_level: 3, + flush_timeout: Duration::from_secs(5), + } +} + +#[tokio::test] +async fn test_flush_empty_aggregator_does_not_call_api() { + let (service, handle) = AggregatorService::new(); + let _task = tokio::spawn(service.run()); + + // Server with no routes — any request would cause test failure + let mock_server = mockito::Server::new_async().await; + let config = config_for_mock(&mock_server.url()); + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new(config, client, handle); + + assert!( + flusher.flush(vec![]).await.is_empty(), + "empty flush should succeed" + ); + // No mock assertions needed — absence of request is the assertion +} + +#[tokio::test] +async fn test_flush_sends_post_with_api_key_header() { + // Verify that Datadog mode sends both DD-API-KEY and DD-PROTOCOL: + // agent-json headers. We call ship_batch directly to bypass + // resolve_endpoint (which builds an HTTPS URL incompatible with the + // HTTP mock server). + let (service, handle) = AggregatorService::new(); + let _task = tokio::spawn(service.run()); + + let mut mock_server = mockito::Server::new_async().await; + let mock = mock_server + .mock("POST", "/api/v2/logs") + .match_header("DD-API-KEY", "test-api-key") + .match_header("DD-PROTOCOL", "agent-json") + .with_status(202) + .create_async() + .await; + + let config = LogFlusherConfig { + api_key: "test-api-key".to_string(), + site: "datadoghq.com".to_string(), + mode: Destination::Datadog, + additional_endpoints: Vec::new(), + use_compression: false, + compression_level: 3, + flush_timeout: Duration::from_secs(5), + }; + + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new(config, client, handle); + + // Call ship_batch directly to use the mock server's HTTP URL instead + // of the HTTPS URL that resolve_endpoint would produce. + let url = format!("{}/api/v2/logs", mock_server.url()); + let batch = b"[{\"message\":\"test\"}]"; + flusher + .ship_batch(batch, &url, false, "test-api-key", true) + .await + .expect("ship_batch should succeed"); + + mock.assert_async().await; +} + +#[tokio::test] +async fn test_flush_opw_mode_omits_dd_protocol_header() { + let (service, handle) = AggregatorService::new(); + let _task = tokio::spawn(service.run()); + + let mut mock_server = mockito::Server::new_async().await; + let opw_url = format!("{}/logs", mock_server.url()); + + // Verify DD-PROTOCOL is NOT present in OPW requests + let mock = mock_server + .mock("POST", "/logs") + .match_header("DD-API-KEY", "test-api-key") + .match_header("DD-PROTOCOL", Matcher::Missing) + .with_status(200) + .expect(1) + .create_async() + .await; + + let config = LogFlusherConfig { + api_key: "test-api-key".to_string(), + site: "unused".to_string(), + mode: Destination::ObservabilityPipelinesWorker { url: opw_url }, + additional_endpoints: Vec::new(), + use_compression: false, + compression_level: 3, + flush_timeout: Duration::from_secs(5), + }; + + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new(config, client, handle.clone()); + + handle + .insert_batch(vec![make_entry("opw log")]) + .expect("insert"); + let result = flusher.flush(vec![]).await; + assert!(result.is_empty(), "OPW flush should return empty on 200"); + mock.assert_async().await; +} + +/// When the primary destination is OPW, additional endpoints are still Datadog +/// intakes and must receive the `DD-PROTOCOL: agent-json` header. +#[tokio::test] +async fn test_opw_primary_additional_endpoint_receives_dd_protocol_header() { + 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; + + // OPW primary must NOT have DD-PROTOCOL header + let _primary_mock = primary + .mock("POST", "/logs") + .match_header("DD-PROTOCOL", Matcher::Missing) + .with_status(200) + .expect(1) + .create_async() + .await; + + // Additional endpoint (Datadog intake) MUST have DD-PROTOCOL header + let extra_mock = extra + .mock("POST", "/extra") + .match_header("DD-PROTOCOL", "agent-json") + .with_status(202) + .expect(1) + .create_async() + .await; + + let config = LogFlusherConfig { + api_key: "key".to_string(), + site: "datadoghq.com".to_string(), + mode: Destination::ObservabilityPipelinesWorker { + url: format!("{}/logs", primary.url()), + }, + additional_endpoints: vec![LogsAdditionalEndpoint { + api_key: "extra-key".to_string(), + url: format!("{}/extra", extra.url()), + is_reliable: true, + }], + use_compression: false, + compression_level: 3, + flush_timeout: Duration::from_secs(5), + }; + + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new(config, client, handle.clone()); + handle + .insert_batch(vec![make_entry("test")]) + .expect("insert"); + + assert!(flusher.flush(vec![]).await.is_empty()); + extra_mock.assert_async().await; +} + +#[tokio::test] +async fn test_flush_does_not_retry_on_403() { + let (service, handle) = AggregatorService::new(); + let _task = tokio::spawn(service.run()); + + let mut mock_server = mockito::Server::new_async().await; + // expect(1) means exactly one call — if retried, the test will fail + let mock = mock_server + .mock("POST", "/api/v2/logs") + .with_status(403) + .expect(1) + .create_async() + .await; + + let config = config_for_mock(&mock_server.url()); + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new(config, client, handle.clone()); + + handle + .insert_batch(vec![make_entry("log")]) + .expect("insert"); + let result = flusher.flush(vec![]).await; + // 403 is a permanent error — the batch is dropped, no builder to retry. + assert!( + result.is_empty(), + "403 is a permanent error; no builder to retry" + ); + mock.assert_async().await; +} + +/// 429 (Too Many Requests) is a retryable 4xx — the retry loop must +/// continue rather than short-circuiting with a permanent failure. +#[tokio::test] +async fn test_flush_retries_on_429_then_succeeds() { + let (service, handle) = AggregatorService::new(); + let _task = tokio::spawn(service.run()); + + let mut mock_server = mockito::Server::new_async().await; + // First call → 429, second call → 200 + let _throttled = mock_server + .mock("POST", "/api/v2/logs") + .with_status(429) + .expect(1) + .create_async() + .await; + let _ok = mock_server + .mock("POST", "/api/v2/logs") + .with_status(200) + .expect(1) + .create_async() + .await; + + let config = config_for_mock(&mock_server.url()); + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new(config, client, handle.clone()); + + handle + .insert_batch(vec![make_entry("throttled log")]) + .expect("insert"); + let result = flusher.flush(vec![]).await; + assert!(result.is_empty(), "should succeed after 429 retry"); +} + +#[tokio::test] +async fn test_flush_retries_on_5xx_then_succeeds() { + let (service, handle) = AggregatorService::new(); + let _task = tokio::spawn(service.run()); + + let mut mock_server = mockito::Server::new_async().await; + // First call → 500, second call → 202 + let _fail_mock = mock_server + .mock("POST", "/api/v2/logs") + .with_status(500) + .expect(1) + .create_async() + .await; + let _ok_mock = mock_server + .mock("POST", "/api/v2/logs") + .with_status(202) + .expect(1) + .create_async() + .await; + + let config = config_for_mock(&mock_server.url()); + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new(config, client, handle.clone()); + + handle + .insert_batch(vec![make_entry("log")]) + .expect("insert"); + let result = flusher.flush(vec![]).await; + assert!(result.is_empty(), "should succeed on second attempt"); +} + +/// All additional endpoints receive the same batch when flush() is called. +#[tokio::test] +async fn test_additional_endpoints_all_receive_batch() { + let (service, handle) = AggregatorService::new(); + let _task = tokio::spawn(service.run()); + + let mut primary = mockito::Server::new_async().await; + let mut extra1 = mockito::Server::new_async().await; + let mut extra2 = mockito::Server::new_async().await; + + let primary_mock = primary + .mock("POST", "/api/v2/logs") + .with_status(202) + .expect(1) + .create_async() + .await; + let extra1_mock = extra1 + .mock("POST", "/extra") + .with_status(200) + .expect(1) + .create_async() + .await; + let extra2_mock = extra2 + .mock("POST", "/extra") + .with_status(200) + .expect(1) + .create_async() + .await; + + let config = LogFlusherConfig { + api_key: "key".to_string(), + site: "datadoghq.com".to_string(), + mode: Destination::ObservabilityPipelinesWorker { + url: format!("{}/api/v2/logs", primary.url()), + }, + additional_endpoints: vec![ + LogsAdditionalEndpoint { + api_key: "extra-key-1".to_string(), + url: format!("{}/extra", extra1.url()), + is_reliable: true, + }, + LogsAdditionalEndpoint { + api_key: "extra-key-2".to_string(), + url: format!("{}/extra", extra2.url()), + is_reliable: true, + }, + ], + use_compression: false, + compression_level: 3, + flush_timeout: Duration::from_secs(5), + }; + + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new(config, client, handle.clone()); + handle.insert_batch(vec![make_entry("hi")]).expect("insert"); + + assert!(flusher.flush(vec![]).await.is_empty()); + primary_mock.assert_async().await; + extra1_mock.assert_async().await; + extra2_mock.assert_async().await; +} + +/// Additional endpoints are dispatched concurrently: if they were sequential, +/// two endpoints each waiting at a Barrier(2) would deadlock — only concurrent +/// dispatch lets both handlers reach the barrier simultaneously. +#[tokio::test] +async fn test_additional_endpoints_dispatched_concurrently() { + use std::sync::Arc; + use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; + use tokio::net::TcpListener; + use tokio::sync::Barrier; + + let (service, handle) = AggregatorService::new(); + let _task = tokio::spawn(service.run()); + + let barrier = Arc::new(Barrier::new(2)); + + // Spawn a minimal HTTP server that waits at the barrier before + // responding, so both must be in-flight at the same time to complete. + async fn serve_once(barrier: Arc) -> String { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut buf = vec![0u8; 4096]; + let _ = stream.read(&mut buf).await; + barrier.wait().await; + let _ = stream + .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n") + .await; + }); + format!("http://127.0.0.1:{}/logs", addr.port()) + } + + let url1 = serve_once(barrier.clone()).await; + let url2 = serve_once(barrier.clone()).await; + + let mut primary = mockito::Server::new_async().await; + let _primary_mock = primary + .mock("POST", "/api/v2/logs") + .with_status(202) + .create_async() + .await; + + let config = LogFlusherConfig { + api_key: "key".to_string(), + site: "datadoghq.com".to_string(), + mode: Destination::ObservabilityPipelinesWorker { + url: format!("{}/api/v2/logs", primary.url()), + }, + additional_endpoints: vec![ + LogsAdditionalEndpoint { + api_key: "extra-key-1".to_string(), + url: url1, + is_reliable: true, + }, + LogsAdditionalEndpoint { + api_key: "extra-key-2".to_string(), + url: url2, + is_reliable: true, + }, + ], + use_compression: false, + compression_level: 3, + flush_timeout: Duration::from_secs(5), + }; + + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new(config, client, handle.clone()); + handle + .insert_batch(vec![make_entry("concurrent")]) + .expect("insert"); + + assert!(flusher.flush(vec![]).await.is_empty()); +} + +/// A builder returned by `flush` can be redriven on the next call. +/// +/// The mock fails on the first 3 attempts (exhausting the per-invocation +/// retry budget), then succeeds on the 4th attempt (the next invocation). +/// This proves the cross-invocation retry path end-to-end. +#[tokio::test] +async fn test_cross_invocation_retry_delivers_on_redrive() { + let (service, handle) = AggregatorService::new(); + let _task = tokio::spawn(service.run()); + + let mut mock_server = mockito::Server::new_async().await; + // First 3 calls: transient 503 → exhausts per-invocation retry budget + let _fail_mock = mock_server + .mock("POST", "/api/v2/logs") + .with_status(503) + .expect(3) + .create_async() + .await; + // 4th call: redriven on the next flush → succeeds + let _ok_mock = mock_server + .mock("POST", "/api/v2/logs") + .with_status(200) + .expect(1) + .create_async() + .await; + + let config = config_for_mock(&mock_server.url()); + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new(config, client, handle.clone()); + handle + .insert_batch(vec![make_entry("retry-me")]) + .expect("insert"); + + // First flush: all 3 attempts fail → returns the builder for retry. + let failed = flusher.flush(vec![]).await; + assert_eq!(failed.len(), 1, "one builder should be returned for retry"); + + // Second flush: aggregator is empty; redrives the failed builder → succeeds. + let result = flusher.flush(failed).await; + assert!( + result.is_empty(), + "redriven builder should succeed on the next invocation" + ); +} + +/// When the primary is OPW (compression disabled for OPW transport) and +/// `use_compression` is true, additional Datadog endpoints must still +/// receive `Content-Encoding: zstd` — they are Datadog intakes and honour +/// `DD_LOGS_CONFIG_USE_COMPRESSION` independently of the primary. +#[tokio::test] +async fn test_opw_primary_additional_endpoint_compresses_when_enabled() { + 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; + + // OPW primary must NOT have Content-Encoding + let _primary_mock = primary + .mock("POST", "/logs") + .match_header("Content-Encoding", Matcher::Missing) + .with_status(200) + .expect(1) + .create_async() + .await; + + // Additional Datadog endpoint MUST receive Content-Encoding: zstd + let extra_mock = extra + .mock("POST", "/extra") + .match_header("Content-Encoding", "zstd") + .with_status(202) + .expect(1) + .create_async() + .await; + + let config = LogFlusherConfig { + api_key: "key".to_string(), + site: "datadoghq.com".to_string(), + mode: Destination::ObservabilityPipelinesWorker { + url: format!("{}/logs", primary.url()), + }, + additional_endpoints: vec![LogsAdditionalEndpoint { + api_key: "extra-key".to_string(), + url: format!("{}/extra", extra.url()), + is_reliable: true, + }], + use_compression: true, + compression_level: 3, + flush_timeout: Duration::from_secs(5), + }; + + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new(config, client, handle.clone()); + handle + .insert_batch(vec![make_entry("compressed extra")]) + .expect("insert"); + + assert!(flusher.flush(vec![]).await.is_empty()); + extra_mock.assert_async().await; +} + +/// Even when `use_compression: true`, the OPW primary endpoint must never +/// receive a `Content-Encoding` header — OPW does not support zstd. +#[tokio::test] +async fn test_opw_primary_never_compressed_even_when_flag_set() { + let (service, handle) = AggregatorService::new(); + let _task = tokio::spawn(service.run()); + + let mut mock_server = mockito::Server::new_async().await; + let mock = mock_server + .mock("POST", "/logs") + .match_header("Content-Encoding", Matcher::Missing) + .with_status(200) + .expect(1) + .create_async() + .await; + + let config = LogFlusherConfig { + api_key: "key".to_string(), + site: "datadoghq.com".to_string(), + mode: Destination::ObservabilityPipelinesWorker { + url: format!("{}/logs", mock_server.url()), + }, + additional_endpoints: vec![], + use_compression: true, // flag set but must not reach OPW + compression_level: 3, + flush_timeout: Duration::from_secs(5), + }; + + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new(config, client, handle.clone()); + handle + .insert_batch(vec![make_entry("opw no compress")]) + .expect("insert"); + + assert!(flusher.flush(vec![]).await.is_empty()); + mock.assert_async().await; +} + +/// Additional-endpoint failures are best-effort: their builders are NOT +/// included in the returned retry vec, even when they exhaust all retries. +/// Only primary-endpoint failures are tracked for cross-invocation retry. +#[tokio::test] +async fn test_additional_endpoint_failures_not_tracked_for_retry() { + 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") + .with_status(200) + .expect(1) + .create_async() + .await; + // Additional endpoint always returns 503 — exhausts per-invocation retries. + let _extra_mock = extra + .mock("POST", "/extra") + .with_status(503) + .expect(3) // MAX_FLUSH_ATTEMPTS + .create_async() + .await; + + let config = LogFlusherConfig { + api_key: "key".to_string(), + site: "datadoghq.com".to_string(), + mode: Destination::ObservabilityPipelinesWorker { + url: format!("{}/api/v2/logs", primary.url()), + }, + additional_endpoints: vec![LogsAdditionalEndpoint { + api_key: "extra-key".to_string(), + url: format!("{}/extra", extra.url()), + is_reliable: true, + }], + use_compression: false, + compression_level: 3, + flush_timeout: Duration::from_secs(5), + }; + + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new(config, client, handle.clone()); + handle + .insert_batch(vec![make_entry("test")]) + .expect("insert"); + + let result = flusher.flush(vec![]).await; + assert!( + result.is_empty(), + "additional-endpoint failures are best-effort and must not be tracked for retry" + ); +} + +/// Default policy: client errors other than 408/425/429 are dropped after one attempt. +#[tokio::test] +async fn test_default_policy_drops_permanent_client_errors_after_one_attempt() { + for status in [400_usize, 401, 404, 413] { + 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(status) + .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()); + handle + .insert_batch(vec![make_entry("log")]) + .expect("insert"); + + assert!( + flusher.flush(vec![]).await.is_empty(), + "status {status} must be dropped, not handed back" + ); + mock.assert_async().await; + } +} + +/// Default policy: 408, 425 and 429 are retried like server errors, then handed back. +#[tokio::test] +async fn test_default_policy_retries_408_425_and_429() { + for status in [408_usize, 425, 429] { + 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(status) + .expect(3) + .create_async() + .await; + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new(config_for_mock(&server.url()), client, handle.clone()); + handle + .insert_batch(vec![make_entry("log")]) + .expect("insert"); + + assert_eq!( + flusher.flush(vec![]).await.len(), + 1, + "status {status} is handed back" + ); + mock.assert_async().await; + } +} + +/// Connection errors are retried and then handed back. +#[tokio::test] +async fn test_default_policy_hands_back_connection_errors() { + let (service, handle) = AggregatorService::new(); + let _task = tokio::spawn(service.run()); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind"); + let port = listener.local_addr().expect("addr").port(); + let accepts = Arc::new(AtomicUsize::new(0)); + let accepts_in_task = Arc::clone(&accepts); + let _server = tokio::spawn(async move { + while let Ok((stream, _)) = listener.accept().await { + accepts_in_task.fetch_add(1, Ordering::SeqCst); + drop(stream); + } + }); + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new( + config_for_mock(&format!("http://127.0.0.1:{port}")), + client, + handle.clone(), + ); + handle + .insert_batch(vec![make_entry("log")]) + .expect("insert"); + + assert_eq!(flusher.flush(vec![]).await.len(), 1); + assert_eq!(accepts.load(Ordering::SeqCst), 3); +} + +/// A redriven request that fails again gets a fresh retry budget and is handed back again. +#[tokio::test] +async fn test_redriven_request_that_fails_again_is_handed_back_again() { + 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(503) + .expect(6) + .create_async() + .await; + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new(config_for_mock(&server.url()), client, handle.clone()); + handle + .insert_batch(vec![make_entry("log")]) + .expect("insert"); + + let failed = flusher.flush(vec![]).await; + assert_eq!(failed.len(), 1); + let failed_again = flusher.flush(failed).await; + assert_eq!(failed_again.len(), 1); + mock.assert_async().await; +} + +/// Sends one batch in Datadog mode to a listener that drops every connection, so the +/// request is handed back and can be inspected. +async fn handed_back_datadog_request(use_compression: bool) -> reqwest::Request { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind"); + let addr = listener.local_addr().expect("addr"); + let _server = tokio::spawn(async move { + while let Ok((stream, _)) = listener.accept().await { + drop(stream); + } + }); + let (service, handle) = AggregatorService::new(); + let _task = tokio::spawn(service.run()); + let config = LogFlusherConfig { + site: "datadoghq.eu".to_string(), + mode: Destination::Datadog, + use_compression, + ..config_for_mock("http://unused") + }; + // Route the public intake hostname to the local listener. + let client = reqwest::Client::builder() + .resolve("http-intake.logs.datadoghq.eu", addr) + .build() + .expect("client"); + let flusher = LogFlusher::new(config, client, handle.clone()); + handle + .insert_batch(vec![make_entry("default path")]) + .expect("insert"); + + let mut failed = flusher.flush(vec![]).await; + assert_eq!(failed.len(), 1); + failed.remove(0).build().expect("request") +} + +/// Datadog mode by default sends compressed batches to the site's public intake with the +/// Datadog-intake headers. +#[tokio::test] +async fn test_default_datadog_request_when_compressed() { + let req = handed_back_datadog_request(true).await; + + assert_eq!( + req.url().as_str(), + "https://http-intake.logs.datadoghq.eu/api/v2/logs" + ); + assert_eq!( + req.timeout(), + Some(&config_for_mock("http://unused").flush_timeout) + ); + let headers = req.headers(); + assert_eq!(headers["DD-API-KEY"], "test-api-key"); + assert_eq!(headers["DD-PROTOCOL"], "agent-json"); + assert_eq!(headers["Content-Type"], "application/json"); + assert_eq!(headers["Content-Encoding"], "zstd"); + let body = + zstd::decode_all(req.body().expect("body").as_bytes().expect("bytes")).expect("zstd"); + let json: serde_json::Value = serde_json::from_slice(&body).expect("json"); + assert_eq!(json[0]["message"], "default path"); +} + +/// With compression off, Datadog mode sends the raw batch without Content-Encoding. +#[tokio::test] +async fn test_default_datadog_request_when_uncompressed() { + let req = handed_back_datadog_request(false).await; + + assert_eq!( + req.url().as_str(), + "https://http-intake.logs.datadoghq.eu/api/v2/logs" + ); + assert_eq!( + req.timeout(), + Some(&config_for_mock("http://unused").flush_timeout) + ); + let headers = req.headers(); + assert_eq!(headers["DD-API-KEY"], "test-api-key"); + assert_eq!(headers["DD-PROTOCOL"], "agent-json"); + assert_eq!(headers["Content-Type"], "application/json"); + assert!(headers.get("Content-Encoding").is_none()); + let json: serde_json::Value = + serde_json::from_slice(req.body().expect("body").as_bytes().expect("bytes")).expect("json"); + assert_eq!(json[0]["message"], "default path"); +} + +/// One-shot HTTP server. Records ":received" when the full request has arrived and +/// ":responded" just before answering 200 after `delay`; hands back the body. +async fn capture_once( + name: &'static str, + delay: Duration, + events: Arc>>, +) -> (String, tokio::sync::oneshot::Receiver>) { + use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind"); + let port = listener.local_addr().expect("addr").port(); + let (tx, rx) = tokio::sync::oneshot::channel(); + tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.expect("accept"); + let mut raw = Vec::new(); + let mut buf = [0_u8; 8192]; + let body = loop { + let n = stream.read(&mut buf).await.expect("read"); + assert!(n > 0, "connection closed before the request was complete"); + raw.extend_from_slice(&buf[..n]); + let Some(end) = raw.windows(4).position(|w| w == b"\r\n\r\n") else { + continue; + }; + let head = String::from_utf8_lossy(&raw[..end]).to_ascii_lowercase(); + let len: usize = head + .lines() + .find_map(|l| l.strip_prefix("content-length:")) + .expect("request has content-length") + .trim() + .parse() + .expect("content-length"); + while raw.len() < end + 4 + len { + let n = stream.read(&mut buf).await.expect("read body"); + assert!(n > 0, "connection closed before the body was complete"); + raw.extend_from_slice(&buf[..n]); + } + break raw[end + 4..end + 4 + len].to_vec(); + }; + events + .lock() + .expect("lock") + .push(format!("{name}:received")); + tokio::time::sleep(delay).await; + let _ = tx.send(body); + events + .lock() + .expect("lock") + .push(format!("{name}:responded")); + let _ = stream + .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n") + .await; + }); + (format!("http://127.0.0.1:{port}/logs"), rx) +} + +/// Default policy: per batch, additional endpoints are sent only after the primary +/// request has finished. +#[tokio::test] +async fn test_default_policy_sends_additional_endpoints_after_primary() { + let events = Arc::new(std::sync::Mutex::new(Vec::new())); + let (primary_url, _primary_body) = + capture_once("primary", Duration::from_millis(200), events.clone()).await; + let (extra_url, _extra_body) = capture_once("extra", Duration::ZERO, events.clone()).await; + + let (service, handle) = AggregatorService::new(); + let _task = tokio::spawn(service.run()); + let config = LogFlusherConfig { + mode: Destination::ObservabilityPipelinesWorker { url: primary_url }, + additional_endpoints: vec![LogsAdditionalEndpoint { + api_key: "extra-key".to_string(), + url: extra_url, + is_reliable: true, + }], + ..config_for_mock("http://unused") + }; + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new(config, client, handle.clone()); + handle + .insert_batch(vec![make_entry("ordered")]) + .expect("insert"); + + assert!(flusher.flush(vec![]).await.is_empty()); + assert_eq!( + *events.lock().expect("lock"), + vec![ + "primary:received", + "primary:responded", + "extra:received", + "extra:responded" + ] + ); +} + +/// Default policy: when the primary fails, additional endpoints are still sent with their +/// own API key, and only the primary request is handed back. +#[tokio::test] +async fn test_default_extras_still_sent_with_their_own_key_when_primary_fails() { + 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") + .match_header("DD-API-KEY", "test-api-key") + .with_status(503) + .expect(3) + .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()); + handle + .insert_batch(vec![make_entry("log")]) + .expect("insert"); + + let failed = flusher.flush(vec![]).await; + assert_eq!(failed.len(), 1); + let req = failed + .into_iter() + .next() + .expect("handed-back request") + .build() + .expect("request"); + assert_eq!(req.url().as_str(), format!("{}/api/v2/logs", primary.url())); + assert_eq!(req.headers()["DD-API-KEY"], "test-api-key"); + primary_mock.assert_async().await; + extra_mock.assert_async().await; +} + +/// Every additional endpoint receives byte-identical compressed bodies that decompress to +/// exactly what the primary received. +#[tokio::test] +async fn test_additional_endpoints_receive_identical_compressed_bodies() { + let events = Arc::new(std::sync::Mutex::new(Vec::new())); + let (primary_url, mut primary_body) = + capture_once("primary", Duration::ZERO, events.clone()).await; + let (url_1, mut body_1) = capture_once("extra1", Duration::ZERO, events.clone()).await; + let (url_2, mut body_2) = capture_once("extra2", Duration::ZERO, events.clone()).await; + + let (service, handle) = AggregatorService::new(); + let _task = tokio::spawn(service.run()); + let config = LogFlusherConfig { + mode: Destination::ObservabilityPipelinesWorker { url: primary_url }, + additional_endpoints: vec![ + LogsAdditionalEndpoint { + api_key: "extra-key-1".to_string(), + url: url_1, + is_reliable: true, + }, + LogsAdditionalEndpoint { + api_key: "extra-key-2".to_string(), + url: url_2, + is_reliable: true, + }, + ], + use_compression: true, + ..config_for_mock("http://unused") + }; + let compression_level = config.compression_level; + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new(config, client, handle.clone()); + handle + .insert_batch( + // Varied enough that zstd levels 3 and 4 produce different bytes. + (0..200_u64) + .map(|i| { + let words: Vec = (0..40) + .map(|j| format!("{:x}", (i * 40 + j) * 2_654_435_761 % 1_000_003)) + .collect(); + make_entry(&format!("same bytes {i} {}", words.join(" "))) + }) + .collect(), + ) + .expect("insert"); + + assert!(flusher.flush(vec![]).await.is_empty()); + let primary_body = primary_body.try_recv().expect("primary received a request"); + let body_1 = body_1.try_recv().expect("extra1 received a request"); + let body_2 = body_2.try_recv().expect("extra2 received a request"); + assert_eq!(body_1, body_2); + assert_eq!( + body_1, + zstd::encode_all(primary_body.as_slice(), compression_level).expect("zstd"), + "additional endpoints use the configured compression level" + ); + let decompressed = zstd::decode_all(body_1.as_slice()).expect("zstd"); + assert_eq!(decompressed, primary_body); + let decoded: serde_json::Value = serde_json::from_slice(&decompressed).expect("json"); + assert_eq!(decoded.as_array().expect("array").len(), 200); +} + +/// Default: a request that keeps failing logs one warning per attempt plus a final one. +#[tokio::test(flavor = "current_thread")] +async fn test_default_logging_warns_on_every_failed_attempt() { + let (logged, _guard) = LevelCounter::install(); + 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(503) + .expect(3) + .create_async() + .await; + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new(config_for_mock(&server.url()), client, handle.clone()); + handle + .insert_batch(vec![make_entry("log")]) + .expect("insert"); + + assert_eq!(flusher.flush(vec![]).await.len(), 1); + mock.assert_async().await; + assert_eq!( + logged.warnings(), + 4, + "3 transient-error warnings + 1 gave-up warning" + ); + assert_eq!(logged.errors(), 0, "default logging uses warn!, not error!"); +} + +/// Default: a connection that keeps dropping logs one warning per attempt plus a final one. +#[tokio::test(flavor = "current_thread")] +async fn test_default_logging_warns_on_every_failed_network_attempt() { + let (logged, _guard) = LevelCounter::install(); + let (service, handle) = AggregatorService::new(); + let _task = tokio::spawn(service.run()); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind"); + let port = listener.local_addr().expect("addr").port(); + let _server = tokio::spawn(async move { + while let Ok((stream, _)) = listener.accept().await { + drop(stream); + } + }); + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new( + config_for_mock(&format!("http://127.0.0.1:{port}")), + client, + handle.clone(), + ); + handle + .insert_batch(vec![make_entry("log")]) + .expect("insert"); + + assert_eq!(flusher.flush(vec![]).await.len(), 1); + assert_eq!( + logged.warnings(), + 4, + "3 network-error warnings + 1 gave-up warning" + ); + assert_eq!(logged.errors(), 0, "default logging uses warn!, not error!"); +} + +/// Default: a permanent client error logs exactly one warning. +#[tokio::test(flavor = "current_thread")] +async fn test_default_logging_warns_once_on_permanent_error() { + let (logged, _guard) = LevelCounter::install(); + 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(404) + .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()); + handle + .insert_batch(vec![make_entry("log")]) + .expect("insert"); + + assert!(flusher.flush(vec![]).await.is_empty()); + mock.assert_async().await; + assert_eq!(logged.warnings(), 1); + assert_eq!(logged.errors(), 0, "default logging uses warn!, not error!"); +} + +/// Default: an additional endpoint that keeps failing adds one more warning on top of +/// the per-attempt ones. +#[tokio::test(flavor = "current_thread")] +async fn test_default_logging_for_failing_additional_endpoint() { + 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") + .with_status(202) + .expect(1) + .create_async() + .await; + let extra_mock = extra + .mock("POST", "/extra") + .with_status(503) + .expect(3) + .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()); + 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(), + 5, + "3 transient + 1 gave-up + 1 additional-endpoint warning" + ); + assert_eq!(logged.errors(), 0, "default logging uses warn!, not error!"); +} + +/// Default: each entry over the 1 MiB limit is dropped with exactly one warning. +#[tokio::test(flavor = "current_thread")] +async fn test_default_logging_warns_once_per_dropped_oversized_entry() { + let (logged, _guard) = LevelCounter::install(); + let (service, handle) = AggregatorService::new(); + let _task = tokio::spawn(service.run()); + + let oversized = "a".repeat(crate::constants::MAX_LOG_BYTES + 1); + handle + .insert_batch(vec![make_entry(&oversized), make_entry(&oversized)]) + .expect("insert"); + assert!(handle.get_batches().await.expect("batches").is_empty()); + assert_eq!(logged.warnings(), 2); + assert_eq!(logged.errors(), 0, "default logging uses warn!, not error!"); +} diff --git a/crates/datadog-logs-agent/src/lib.rs b/crates/datadog-logs-agent/src/lib.rs index ef1c4542..5f2b0f63 100644 --- a/crates/datadog-logs-agent/src/lib.rs +++ b/crates/datadog-logs-agent/src/lib.rs @@ -16,9 +16,11 @@ pub mod intake_entry; pub mod logs_additional_endpoint; pub mod server; +#[cfg(test)] +mod test_support; // Re-export the most commonly used types at the crate root -pub use aggregator::{AggregatorHandle, AggregatorService}; +pub use aggregator::{AggregatorHandle, AggregatorService, OversizePolicy}; pub use config::{Destination, LogFlusherConfig}; pub use flusher::LogFlusher; pub use intake_entry::IntakeEntry; diff --git a/crates/datadog-logs-agent/src/test_support.rs b/crates/datadog-logs-agent/src/test_support.rs new file mode 100644 index 00000000..c529d8a7 --- /dev/null +++ b/crates/datadog-logs-agent/src/test_support.rs @@ -0,0 +1,52 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +//! Helpers shared by this crate's unit tests. + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +/// Counts `warn!` and `error!` events emitted by this crate on the current thread. +#[derive(Clone, Default)] +pub(crate) struct LevelCounter { + warnings: Arc, + errors: Arc, +} + +impl tracing_subscriber::Layer for LevelCounter { + fn on_event( + &self, + event: &tracing::Event<'_>, + _ctx: tracing_subscriber::layer::Context<'_, S>, + ) { + let meta = event.metadata(); + if !meta.target().starts_with("datadog_logs_agent") { + return; + } + match *meta.level() { + tracing::Level::WARN => self.warnings.fetch_add(1, Ordering::SeqCst), + tracing::Level::ERROR => self.errors.fetch_add(1, Ordering::SeqCst), + _ => return, + }; + } +} + +impl LevelCounter { + /// Installs the counter for this thread. Use the current-thread runtime so tasks + /// spawned by the test run on this thread. + pub(crate) fn install() -> (Self, tracing::subscriber::DefaultGuard) { + use tracing_subscriber::layer::SubscriberExt as _; + let counter = Self::default(); + let guard = + tracing::subscriber::set_default(tracing_subscriber::registry().with(counter.clone())); + (counter, guard) + } + + pub(crate) fn warnings(&self) -> usize { + self.warnings.load(Ordering::SeqCst) + } + + pub(crate) fn errors(&self) -> usize { + self.errors.load(Ordering::SeqCst) + } +}