[SVLS-8713] feat(logs-agent): accept pre-serialized logs and an oversize policy

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<String>,
     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<bool, AggregatorError> {
         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<bool, AggregatorError> {
         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::<Vec<serde_json::Value>>(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);
+    }
 }

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};

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<IntakeEntry>),
+    InsertSerializedBatch(Vec<String>),
     GetBatches(oneshot::Sender<Vec<Vec<u8>>>),
     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<String>) -> 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));
+    }
 }