APMSP-1969 fix dropped trace payloads for the sidecar (#1047)

- The MAX_PAYLOAD_SIZE used in trace_utils::coalesce_send_data() was
  50mb. The agent drops payloads greater than > 25mb and returns a 413.
  So it was potentially combining payloads that would result in an
  error and drop.

- In the sidecar's trace_flusher payloads were being dropped if the
  queue's size exceeded the min drop size. The correct behavior is to
  check if the payload being enqueued exceeds min drop size and log an
  error.

- When the sidecar's trace_flusher was dropping a payload that was too
  large it was still adding that payload's size to the queue size. This
  could have lead to subsequent payloads being dropped due to an incorrectly
  lar ge queue size until a time based flush was done.

- A bug was discovered in the test helper function poll_for_mock_hits()
  where we were incorrectly returning true always when expected hits was
  0.

Author: ekump

datadog-sidecar/src/service/tracing/trace_flusher.rs

@@ -25,7 +25,6 @@ use tracing::{debug, error, info};
 
 const DEFAULT_FLUSH_INTERVAL_MS: u64 = 5_000;
 const DEFAULT_MIN_FORCE_FLUSH_SIZE_BYTES: u32 = 1_000_000;
-const DEFAULT_MIN_FORCE_DROP_SIZE_BYTES: u32 = 10_000_000;
 
 /// `TraceFlusherStats` holds stats of the trace flusher like the count of allocated shared memory
 /// for agent config, agent config writers, last used entries in agent configs, and the size of send
@@ -103,7 +102,7 @@ impl Default for TraceFlusher {
             inner: Mutex::new(TraceFlusherData::default()),
             interval_ms: AtomicU64::new(DEFAULT_FLUSH_INTERVAL_MS),
             min_force_flush_size_bytes: AtomicU32::new(DEFAULT_MIN_FORCE_FLUSH_SIZE_BYTES),
-            min_force_drop_size_bytes: AtomicU32::new(DEFAULT_MIN_FORCE_DROP_SIZE_BYTES),
+            min_force_drop_size_bytes: AtomicU32::new(trace_utils::MAX_PAYLOAD_SIZE as u32),
             remote_config: Mutex::new(Default::default()),
             metrics: Mutex::new(Default::default()),
         }
@@ -120,20 +119,24 @@ impl TraceFlusher {
         let mut flush_data = self.inner.lock_or_panic();
         let flush_data = flush_data.deref_mut();
 
-        flush_data.traces.send_data_size += data.len();
-
-        if flush_data.traces.send_data_size
-            > self.min_force_drop_size_bytes.load(Ordering::Relaxed) as usize
-        {
+        if data.len() > self.min_force_drop_size_bytes.load(Ordering::Relaxed) as usize {
+            error!(
+                "Error sending trace. Individual trace size of {}B exceeds {}B limit",
+                data.len(),
+                self.min_force_drop_size_bytes.load(Ordering::Relaxed) as usize
+            );
             return;
         }
 
+        flush_data.traces.send_data_size += data.len();
         flush_data.traces.send_data.push(data);
+
         if flush_data.flusher.is_none() {
             let (force_flush, completer) = ManualFuture::new();
             flush_data.flusher = Some(self.clone().start_trace_flusher(force_flush));
             flush_data.traces.force_flush = Some(completer);
         }
+
         if flush_data.traces.send_data_size
             > self.min_force_flush_size_bytes.load(Ordering::Relaxed) as usize
         {
@@ -321,7 +324,11 @@ mod tests {
     // observe that a request to the trace agent is not made. Then enqueue a third trace exceeding
     // the min force flush size, and observe that a request to the trace agent is made.
     async fn test_min_flush_size() {
-        let trace_flusher = Arc::new(TraceFlusher::default());
+        // Set the interval high enough that it can't cause a false positive
+        let trace_flusher = Arc::new(TraceFlusher {
+            interval_ms: AtomicU64::new(20_000),
+            ..TraceFlusher::default()
+        });
 
         let server = MockServer::start();
 
@@ -399,10 +406,12 @@ mod tests {
 
     #[cfg_attr(miri, ignore)]
     #[tokio::test]
-    async fn test_flush_drop_size() {
+    // Test scenario: Enqueue a trace with a size greater than the minimum force drop size, and
+    // observe that it is not sent.
+    async fn test_drop_size_no_flush() {
         // Set the interval high enough that it can't cause a false positive
         let trace_flusher = Arc::new(TraceFlusher {
-            interval_ms: AtomicU64::new(10_000),
+            interval_ms: AtomicU64::new(20_000),
             ..TraceFlusher::default()
         });
         let server = MockServer::start();
@@ -413,6 +422,7 @@ mod tests {
                     .body(r#"{"status":"ok"}"#);
             })
             .await;
+
         let size = trace_flusher
             .min_force_drop_size_bytes
             .load(Ordering::Relaxed) as usize

datadog-trace-utils/src/test_utils/mod.rs

@@ -409,14 +409,21 @@ pub async fn poll_for_mock_hit(
     expected_hits: usize,
     delete_after_hits: bool,
 ) -> bool {
-    let mut mock_hit = mock.hits_async().await == expected_hits;
+    let mut mock_hit = false;
 
     let mut mock_observations_remaining = poll_attempts;
 
     while !mock_hit {
         sleep(Duration::from_millis(sleep_interval_ms)).await;
-        mock_hit = mock.hits_async().await == expected_hits;
         mock_observations_remaining -= 1;
+        mock_hit = if expected_hits > 0 {
+            mock.hits_async().await == expected_hits
+        } else {
+            // If we are polling for 0 hits, we need to ensure we do all observations, otherwise
+            // this will always be true
+            mock.hits_async().await == 0 && mock_observations_remaining == 0
+        };
+
         if mock_observations_remaining == 0 || mock_hit {
             if delete_after_hits {
                 mock.delete();

datadog-trace-utils/src/trace_utils.rs

@@ -23,14 +23,16 @@ use std::cmp::Ordering;
 use std::collections::{HashMap, HashSet};
 use std::env;
 
+/// The maximum payload size for a single request that can be sent to the trace agent. Payloads
+/// larger than this size will be dropped and the agent will return a 413 error if
+/// `datadog-send-real-http-status` is set.
+pub const MAX_PAYLOAD_SIZE: usize = 25 * 1024 * 1024;
 /// Span metric the mini agent must set for the backend to recognize top level span
 const TOP_LEVEL_KEY: &str = "_top_level";
 /// Span metric the tracer sets to denote a top level span
 const TRACER_TOP_LEVEL_KEY: &str = "_dd.top_level";
 const MEASURED_KEY: &str = "_dd.measured";
 const PARTIAL_VERSION_KEY: &str = "_dd.partial_version";
-
-const MAX_PAYLOAD_SIZE: usize = 50 * 1024 * 1024;
 const MAX_STRING_DICT_SIZE: u32 = 25_000_000;
 const SPAN_ELEMENT_COUNT: usize = 12;