From 84cb56d42b00aef31cd0b68f88e4042e76b57c06 Mon Sep 17 00:00:00 2001 From: Tianning Li Date: Tue, 20 Oct 2026 13:21:20 -0400 Subject: [PATCH] perf(logs-agent): compress each log batch only once Why: the flusher compressed the same batch again for every endpoint. With additional endpoints configured, that multiplied CPU work, which matters for bottlecap on small Lambda functions. What: each batch is now compressed at most once, the first time an endpoint needs it, and the same bytes are reused for every other endpoint. The bytes sent do not change. New tests check that the primary and every additional endpoint receive identical bytes, equal to compressing the batch once, and that each batch in a flush is still compressed on its own. Co-Authored-By: Claude Opus 5.5 --- crates/datadog-logs-agent/src/flusher.rs | 14 +- .../datadog-logs-agent/src/flusher/tests.rs | 120 +++++++++++++++++- 2 files changed, 128 insertions(+), 6 deletions(-) diff --git a/crates/datadog-logs-agent/src/flusher.rs b/crates/datadog-logs-agent/src/flusher.rs index 768dcc5..930a3a1 100644 --- a/crates/datadog-logs-agent/src/flusher.rs +++ b/crates/datadog-logs-agent/src/flusher.rs @@ -4,7 +4,7 @@ use std::future::Future; use std::io::Write as _; use std::pin::Pin; -use std::sync::Arc; +use std::sync::{Arc, OnceLock}; use futures::future::join_all; use reqwest::Client; @@ -167,11 +167,14 @@ impl LogFlusher { let primary_api_key = primary_api_key.as_deref(); async move { let is_primary_datadog = matches!(self.config.mode, Destination::Datadog); + // Compressed at most once, on first use, and reused by every endpoint that + // wants compression. + let compressed = &OnceLock::new(); // Primary endpoint — failures are tracked for cross-invocation retry. let primary = async { match primary_api_key { Some(api_key) => { - self.ship_batch(batch, &primary_url, primary_use_compression, api_key, is_primary_datadog) + self.ship_batch(batch, compressed, &primary_url, primary_use_compression, api_key, is_primary_datadog) .await } None => Ok(()), @@ -186,7 +189,7 @@ impl LogFlusher { let extras = join_all(self.config.additional_endpoints.iter().map(|endpoint| async move { let url = &endpoint.url; let result = self - .ship_batch(batch, url, self.config.use_compression, &endpoint.api_key, true) + .ship_batch(batch, compressed, url, self.config.use_compression, &endpoint.api_key, true) .await; if result.is_err() && !self.policy.return_additional_endpoint_failures { warn!( @@ -239,14 +242,15 @@ impl LogFlusher { async fn ship_batch( &self, batch: &[u8], + compressed: &OnceLock, FlushError>>, url: &str, compress: bool, api_key: &str, is_datadog_intake: bool, ) -> Result<(), Box> { let (body, content_encoding) = if compress { - match compress_zstd(batch, self.config.compression_level) { - Ok(compressed) => (compressed, Some("zstd")), + match compressed.get_or_init(|| compress_zstd(batch, self.config.compression_level)) { + Ok(bytes) => (bytes.clone(), Some("zstd")), Err(e) => { warn!("failed to compress log batch, sending uncompressed: {e}"); (batch.to_vec(), None) diff --git a/crates/datadog-logs-agent/src/flusher/tests.rs b/crates/datadog-logs-agent/src/flusher/tests.rs index a514daf..8f1dae3 100644 --- a/crates/datadog-logs-agent/src/flusher/tests.rs +++ b/crates/datadog-logs-agent/src/flusher/tests.rs @@ -86,7 +86,7 @@ async fn test_flush_sends_post_with_api_key_header() { let url = format!("{}/api/v2/logs", mock_server.url()); let batch = b"[{\"message\":\"test\"}]"; flusher - .ship_batch(batch, &url, false, "test-api-key", true) + .ship_batch(batch, &OnceLock::new(), &url, false, "test-api-key", true) .await .expect("ship_batch should succeed"); @@ -1766,3 +1766,121 @@ async fn test_once_per_request_logging_on_network_errors() { assert_eq!(flusher.flush(vec![]).await.len(), 1); assert_eq!((logged.warnings(), logged.errors()), (0, 1)); } + +/// A compressed Datadog-mode primary and every additional endpoint receive the same bytes, +/// equal to compressing the batch once at the configured level. +#[tokio::test] +async fn test_datadog_primary_and_extras_receive_identical_compressed_bytes() { + 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::Datadog, + use_compression: true, + additional_endpoints: vec![ + LogsAdditionalEndpoint { + api_key: "k1".to_string(), + url: url_1, + is_reliable: true, + }, + LogsAdditionalEndpoint { + api_key: "k2".to_string(), + url: url_2, + is_reliable: true, + }, + ], + ..config_for_mock("http://unused") + }; + let level = config.compression_level; + let client = reqwest::Client::builder().build().expect("client"); + let flusher = + LogFlusher::new(config, client, handle.clone()).with_datadog_intake_url(primary_url); + handle + .insert_batch( + (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 = primary_body.try_recv().expect("primary received a request"); + let extra_1 = body_1.try_recv().expect("extra1 received a request"); + let extra_2 = body_2.try_recv().expect("extra2 received a request"); + assert_eq!(primary, extra_1); + assert_eq!(extra_1, extra_2); + let raw = zstd::decode_all(primary.as_slice()).expect("zstd"); + assert_eq!(primary, compress_zstd(&raw, level).expect("compress")); +} + +/// With several batches in one flush, each request carries its own batch's compressed bytes. +#[tokio::test] +async fn test_each_batch_is_compressed_separately() { + use crate::constants::MAX_BATCH_ENTRIES; + + let bodies = Arc::new(std::sync::Mutex::new(Vec::>::new())); + let recorded = bodies.clone(); + let mut server = mockito::Server::new_async().await; + let mock = server + .mock("POST", "/api/v2/logs") + .match_request(move |req| { + recorded + .lock() + .expect("lock") + .push(req.body().expect("body").clone()); + true + }) + .with_status(202) + .expect(2) + .create_async() + .await; + + let (service, handle) = AggregatorService::new(); + let _task = tokio::spawn(service.run()); + let config = LogFlusherConfig { + mode: Destination::Datadog, + use_compression: true, + ..config_for_mock("http://unused") + }; + let client = reqwest::Client::builder().build().expect("client"); + let flusher = LogFlusher::new(config, client, handle.clone()) + .with_datadog_intake_url(format!("{}/api/v2/logs", server.url())); + let count = MAX_BATCH_ENTRIES + 5; + handle + .insert_batch( + (0..count) + .map(|i| make_entry(&format!("entry {i}"))) + .collect(), + ) + .expect("insert"); + + assert!(flusher.flush(vec![]).await.is_empty()); + mock.assert_async().await; + let bodies = bodies.lock().expect("lock"); + let mut messages: Vec = bodies + .iter() + .flat_map(|body| { + let raw = zstd::decode_all(body.as_slice()).expect("zstd"); + let json: serde_json::Value = serde_json::from_slice(&raw).expect("json"); + json.as_array() + .expect("array") + .iter() + .map(|e| e["message"].as_str().expect("message").to_string()) + .collect::>() + }) + .collect(); + messages.sort(); + let mut expected: Vec = (0..count).map(|i| format!("entry {i}")).collect(); + expected.sort(); + assert_eq!(messages, expected); +}