add shutdown channel to stats flusher

duncanpharvey · duncanpharvey · commit c2191f862eae · 2026-04-07T11:50:29.000-04:00
diff --git a/crates/datadog-serverless-compat/src/main.rs b/crates/datadog-serverless-compat/src/main.rs
@@ -177,7 +177,6 @@ pub async fn main() {
 
     let stats_flusher = Arc::new(stats_flusher::ServerlessStatsFlusher {
         stats_concentrator: stats_concentrator_handle.clone(),
-        force_flush_concentrator: false,
     });
     let stats_processor = Arc::new(stats_processor::ServerlessStatsProcessor {});
 
@@ -199,8 +198,9 @@ pub async fn main() {
         proxy_flusher,
     });
 
+    let (_shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
     tokio::spawn(async move {
-        let res = mini_agent.start_mini_agent().await;
+        let res = mini_agent.start_mini_agent(shutdown_rx).await;
         if let Err(e) = res {
             error!("Error when starting serverless trace mini agent: {e:?}");
         }
diff --git a/crates/datadog-trace-agent/src/mini_agent.rs b/crates/datadog-trace-agent/src/mini_agent.rs
@@ -50,7 +50,10 @@ pub struct MiniAgent {
 }
 
 impl MiniAgent {
-    pub async fn start_mini_agent(&self) -> Result<(), Box<dyn std::error::Error>> {
+    pub async fn start_mini_agent(
+        &self,
+        shutdown_rx: tokio::sync::oneshot::Receiver<()>,
+    ) -> Result<(), Box<dyn std::error::Error>> {
         let now = Instant::now();
 
         // verify we are in a serverless function environment. if not, shut down the mini agent.
@@ -93,7 +96,7 @@ impl MiniAgent {
         let stats_config = self.config.clone();
         let stats_flusher_handle = tokio::spawn(async move {
             stats_flusher
-                .start_stats_flusher(stats_config, stats_rx)
+                .start_stats_flusher(stats_config, stats_rx, shutdown_rx)
                 .await;
         });
 
diff --git a/crates/datadog-trace-agent/src/stats_flusher.rs b/crates/datadog-trace-agent/src/stats_flusher.rs
@@ -3,7 +3,8 @@
 
 use async_trait::async_trait;
 use std::{sync::Arc, time};
-use tokio::sync::{Mutex, mpsc::Receiver};
+use tokio::sync::mpsc::Receiver;
+use tokio::sync::oneshot;
 use tracing::{debug, error};
 
 use libdd_trace_protobuf::pb;
@@ -46,22 +47,28 @@ async fn send_stats_payload(config: &Arc<Config>, payload: pb::StatsPayload) {
 #[async_trait]
 pub trait StatsFlusher {
     /// Starts a stats flusher that listens for stats payloads sent to the tokio mpsc Receiver,
-    /// implementing flushing logic that calls flush_stats.
+    /// implementing flushing logic that calls flush_stats. Runs until the shutdown signal fires,
+    /// at which point it performs a final force flush and returns.
     async fn start_stats_flusher(
         &self,
         config: Arc<Config>,
-        mut rx: Receiver<pb::ClientStatsPayload>,
+        rx: Receiver<pb::ClientStatsPayload>,
+        shutdown_rx: oneshot::Receiver<()>,
     );
     /// Flushes stats to the Datadog trace stats intake.
-    async fn flush_stats(&self, config: Arc<Config>, client_stats: Vec<pb::ClientStatsPayload>);
+    /// `force_flush` controls whether in-progress concentrator buckets are flushed (true on
+    /// shutdown, false on normal interval flushes).
+    async fn flush_stats(
+        &self,
+        config: Arc<Config>,
+        client_stats: Vec<pb::ClientStatsPayload>,
+        force_flush: bool,
+    );
 }
 
 #[derive(Clone)]
 pub struct ServerlessStatsFlusher {
     pub stats_concentrator: Option<StatsConcentrator>,
-    /// When false, flushes are done on completed buckets
-    /// When true, flushes are done on any in progress buckets, useful for integration tests
-    pub force_flush_concentrator: bool,
 }
 
 #[async_trait]
@@ -70,44 +77,51 @@ impl StatsFlusher for ServerlessStatsFlusher {
         &self,
         config: Arc<Config>,
         mut rx: Receiver<pb::ClientStatsPayload>,
+        mut shutdown_rx: oneshot::Receiver<()>,
     ) {
-        let buffer: Arc<Mutex<Vec<pb::ClientStatsPayload>>> = Arc::new(Mutex::new(Vec::new()));
-
-        let buffer_producer = buffer.clone();
-        let buffer_consumer = buffer.clone();
-
-        // Drain the stats channel continuously into the buffer
-        tokio::spawn(async move {
-            while let Some(stats_payload) = rx.recv().await {
-                let mut buffer = buffer_producer.lock().await;
-                buffer.push(stats_payload);
-            }
-        });
+        let mut interval =
+            tokio::time::interval(time::Duration::from_secs(config.stats_flush_interval_secs));
+        let mut buffer: Vec<pb::ClientStatsPayload> = Vec::new();
 
-        // Flush stats from the buffer on a fixed interval
         loop {
-            tokio::time::sleep(time::Duration::from_secs(config.stats_flush_interval_secs)).await;
-
-            let mut buffer = buffer_consumer.lock().await;
-            // Copy the batch for this flush
-            let channel_stats = buffer.to_vec();
-            // Reset the buffer so the next tick only sees new stats
-            buffer.clear();
-            // Release the mutex before flushing stats
-            drop(buffer);
-
-            let should_flush = should_flush_stats_buffer(
-                !channel_stats.is_empty(),
-                self.stats_concentrator.is_some(),
-            );
-            if should_flush {
-                self.flush_stats(config.clone(), channel_stats).await;
+            tokio::select! {
+                // Receive client stats and add them to the buffer
+                Some(stats) = rx.recv() => {
+                    buffer.push(stats);
+                }
+
+                // Drain client stats in buffer and stats from concentrator on interval
+                _ = interval.tick() => {
+                    let client_stats = std::mem::take(&mut buffer);
+                    let should_flush = should_flush_stats_buffer(
+                        !client_stats.is_empty(),
+                        self.stats_concentrator.is_some(),
+                    );
+                    if should_flush {
+                        self.flush_stats(config.clone(), client_stats, false).await;
+                    }
+                }
+
+                _ = &mut shutdown_rx => {
+                    // Drain any client stats that arrived before the shutdown signal
+                    while let Ok(stats) = rx.try_recv() {
+                        buffer.push(stats);
+                    }
+                    // Force flush all in progress concentrator buckets on shutdown signal
+                    self.flush_stats(config.clone(), std::mem::take(&mut buffer), true).await;
+                    return;
+                }
             }
         }
     }
 
     /// Flushes client computed stats from the tracer and serverless computed stats as separate payloads
-    async fn flush_stats(&self, config: Arc<Config>, client_stats: Vec<pb::ClientStatsPayload>) {
+    async fn flush_stats(
+        &self,
+        config: Arc<Config>,
+        client_stats: Vec<pb::ClientStatsPayload>,
+        force_flush: bool,
+    ) {
         // Flush client computed stats from the tracer
         if !client_stats.is_empty() {
             let payload = stats_utils::construct_stats_payload(client_stats);
@@ -116,7 +130,7 @@ impl StatsFlusher for ServerlessStatsFlusher {
 
         // Flush serverless computed stats from the concentrator
         if let Some(ref concentrator) = self.stats_concentrator
-            && let Some(agent_stats) = concentrator.flush(self.force_flush_concentrator)
+            && let Some(agent_stats) = concentrator.flush(force_flush)
         {
             let mut payload = stats_utils::construct_stats_payload(vec![agent_stats]);
             payload.client_computed = false;
diff --git a/crates/datadog-trace-agent/tests/common/mocks.rs b/crates/datadog-trace-agent/tests/common/mocks.rs
@@ -12,6 +12,7 @@ use libdd_trace_protobuf::pb;
 use libdd_trace_utils::trace_utils::{self, MiniAgentMetadata, SendData};
 use std::sync::Arc;
 use tokio::sync::mpsc::{Receiver, Sender};
+use tokio::sync::oneshot;
 
 /// Mock trace processor that returns 200 OK for all requests
 #[allow(dead_code)]
@@ -86,14 +87,20 @@ impl StatsFlusher for MockStatsFlusher {
         &self,
         _config: Arc<Config>,
         mut stats_rx: Receiver<pb::ClientStatsPayload>,
+        _shutdown_rx: oneshot::Receiver<()>,
     ) {
         // Consume messages from the channel without processing them
         while let Some(_stats) = stats_rx.recv().await {
             // Just discard the stats - we're not testing the flusher
         }
     }
 
-    async fn flush_stats(&self, _config: Arc<Config>, _traces: Vec<pb::ClientStatsPayload>) {
+    async fn flush_stats(
+        &self,
+        _config: Arc<Config>,
+        _traces: Vec<pb::ClientStatsPayload>,
+        _force_flush: bool,
+    ) {
         // Do nothing
     }
 }
diff --git a/crates/datadog-trace-agent/tests/integration_test.rs b/crates/datadog-trace-agent/tests/integration_test.rs
@@ -69,7 +69,6 @@ pub fn create_mini_agent_with_real_flushers(config: Arc<Config>) -> MiniAgent {
         stats_processor: Arc::new(ServerlessStatsProcessor {}),
         stats_flusher: Arc::new(ServerlessStatsFlusher {
             stats_concentrator: Some(stats_concentrator_handle),
-            force_flush_concentrator: true,
         }),
         env_verifier: Arc::new(MockEnvVerifier),
         proxy_flusher: Arc::new(ProxyFlusher::new(config.clone())),
@@ -178,7 +177,8 @@ async fn test_mini_agent_tcp_handles_requests() {
 
     // Start the mini agent
     let agent_handle = tokio::spawn(async move {
-        let _ = mini_agent.start_mini_agent().await;
+        let (_shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
+        let _ = mini_agent.start_mini_agent(shutdown_rx).await;
     });
 
     // Give server time to start
@@ -277,7 +277,8 @@ async fn test_mini_agent_named_pipe_handles_requests() {
 
     // Start the mini agent
     let agent_handle = tokio::spawn(async move {
-        let _ = mini_agent.start_mini_agent().await;
+        let (_shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
+        let _ = mini_agent.start_mini_agent(shutdown_rx).await;
     });
 
     // Give server time to create pipe
@@ -348,8 +349,9 @@ async fn test_mini_agent_tcp_with_real_flushers() {
 
     let mini_agent = create_mini_agent_with_real_flushers(config);
 
+    let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
     let agent_handle = tokio::spawn(async move {
-        let _ = mini_agent.start_mini_agent().await;
+        let _ = mini_agent.start_mini_agent(shutdown_rx).await;
     });
 
     // Wait for server to be ready
@@ -376,10 +378,13 @@ async fn test_mini_agent_tcp_with_real_flushers() {
             .expect("Failed to send /v0.4/traces request");
     assert_eq!(trace_response.status(), StatusCode::OK);
 
-    // Wait for flush
+    // Wait for trace flush
     tokio::time::sleep(FLUSH_WAIT_DURATION).await;
-
     verify_trace_request(&mock_server);
+
+    // Trigger shutdown to force flush in progress concentrator buckets
+    let _ = shutdown_tx.send(());
+    tokio::time::sleep(FLUSH_WAIT_DURATION).await;
     verify_stats_request(&mock_server); // Stats generator should generate stats from trace payload
 
     // Clean up
@@ -401,7 +406,8 @@ async fn test_mini_agent_tcp_with_real_flushers_and_tracer_computed_stats() {
     let mini_agent = create_mini_agent_with_real_flushers(config);
 
     let agent_handle = tokio::spawn(async move {
-        let _ = mini_agent.start_mini_agent().await;
+        let (_shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
+        let _ = mini_agent.start_mini_agent(shutdown_rx).await;
     });
 
     // Wait for server to be ready
@@ -459,8 +465,9 @@ async fn test_mini_agent_named_pipe_with_real_flushers() {
 
     let mini_agent = create_mini_agent_with_real_flushers(config);
 
+    let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
     let agent_handle = tokio::spawn(async move {
-        let _ = mini_agent.start_mini_agent().await;
+        let _ = mini_agent.start_mini_agent(shutdown_rx).await;
     });
 
     // Wait for server to be ready
@@ -487,10 +494,13 @@ async fn test_mini_agent_named_pipe_with_real_flushers() {
             .expect("Failed to send /v0.4/traces request over named pipe");
     assert_eq!(trace_response.status(), StatusCode::OK);
 
-    // Wait for flush
+    // Wait for trace flush
     tokio::time::sleep(FLUSH_WAIT_DURATION).await;
-
     verify_trace_request(&mock_server);
+
+    // Trigger shutdown to force flush in progress concentrator buckets
+    let _ = shutdown_tx.send(());
+    tokio::time::sleep(FLUSH_WAIT_DURATION).await;
     verify_stats_request(&mock_server);
 
     // Clean up
