implement trace stats for serverless compat

Author: duncanpharvey

Cargo.lock

@@ -18,7 +18,8 @@ use zstd::zstd_safe::CompressionLevel;
 
 use datadog_trace_agent::{
     aggregator::TraceAggregator,
-    config, env_verifier, mini_agent, proxy_flusher, stats_flusher, stats_processor,
+    config, env_verifier, mini_agent, proxy_flusher, stats_concentrator_service, stats_flusher,
+    stats_generator, stats_processor,
     trace_flusher::{self, TraceFlusher},
     trace_processor,
 };
@@ -107,6 +108,12 @@ pub async fn main() {
     let https_proxy = env::var("DD_PROXY_HTTPS")
         .or_else(|_| env::var("HTTPS_PROXY"))
         .ok();
+
+    let dd_serverless_stats_computation_enabled =
+        env::var("DD_SERVERLESS_STATS_COMPUTATION_ENABLED")
+            .map(|val| val.to_lowercase() != "false")
+            .unwrap_or(true);
+
     debug!("Starting serverless trace mini agent");
 
     let env_filter = format!("h2=off,hyper=off,rustls=off,{}", log_level);
@@ -132,11 +139,6 @@ pub async fn main() {
 
     let env_verifier = Arc::new(env_verifier::ServerlessEnvVerifier::default());
 
-    let trace_processor = Arc::new(trace_processor::ServerlessTraceProcessor {});
-
-    let stats_flusher = Arc::new(stats_flusher::ServerlessStatsFlusher {});
-    let stats_processor = Arc::new(stats_processor::ServerlessStatsProcessor {});
-
     let config = match config::Config::new() {
         Ok(c) => Arc::new(c),
         Err(e) => {
@@ -145,6 +147,30 @@ pub async fn main() {
         }
     };
 
+    // Initialize stats concentrator service and generator conditionally
+    let (stats_concentrator_handle, stats_generator) = if dd_serverless_stats_computation_enabled {
+        info!("Serverless stats computation enabled");
+        let (stats_concentrator_service, stats_concentrator_handle) =
+            stats_concentrator_service::StatsConcentratorService::new(config.clone());
+        tokio::spawn(stats_concentrator_service.run());
+        let stats_generator = Arc::new(stats_generator::StatsGenerator::new(
+            stats_concentrator_handle.clone(),
+        ));
+        (Some(stats_concentrator_handle), Some(stats_generator))
+    } else {
+        info!("Serverless stats computation disabled");
+        (None, None)
+    };
+
+    let trace_processor = Arc::new(trace_processor::ServerlessTraceProcessor {
+        stats_generator: stats_generator.clone(),
+    });
+
+    let stats_flusher = Arc::new(stats_flusher::ServerlessStatsFlusher {
+        stats_concentrator: stats_concentrator_handle.clone(),
+    });
+    let stats_processor = Arc::new(stats_processor::ServerlessStatsProcessor {});
+
     let trace_aggregator = Arc::new(TokioMutex::new(TraceAggregator::default()));
     let trace_flusher = Arc::new(trace_flusher::ServerlessTraceFlusher::new(
         trace_aggregator,
@@ -161,6 +187,8 @@ pub async fn main() {
         stats_processor,
         stats_flusher,
         proxy_flusher,
+        stats_concentrator: stats_concentrator_handle,
+        stats_generator,
     });
 
     tokio::spawn(async move {

crates/datadog-serverless-compat/src/main.rs

@@ -24,14 +24,19 @@ async-trait = "0.1.64"
 tracing = { version = "0.1", default-features = false }
 serde = { version = "1.0.145", features = ["derive"] }
 serde_json = "1.0"
+thiserror = { version = "1.0.58", default-features = false }
 libdd-common = { git = "https://github.com/DataDog/libdatadog", rev = "8c88979985154d6d97c0fc2ca9039682981eacad" }
 libdd-trace-protobuf = { git = "https://github.com/DataDog/libdatadog", rev = "8c88979985154d6d97c0fc2ca9039682981eacad" }
+libdd-trace-stats = { git = "https://github.com/DataDog/libdatadog", rev = "8c88979985154d6d97c0fc2ca9039682981eacad" }
 libdd-trace-utils = { git = "https://github.com/DataDog/libdatadog", rev = "8c88979985154d6d97c0fc2ca9039682981eacad", features = [
     "mini_agent",
 ] }
 libdd-trace-obfuscation = { git = "https://github.com/DataDog/libdatadog", rev = "8c88979985154d6d97c0fc2ca9039682981eacad" }
 datadog-fips = { path = "../datadog-fips" }
-reqwest = { version = "0.12.23", features = ["json", "http2"], default-features = false }
+reqwest = { version = "0.12.23", features = [
+    "json",
+    "http2",
+], default-features = false }
 bytes = "1.10.1"
 
 [dev-dependencies]

crates/datadog-trace-agent/Cargo.toml

@@ -109,6 +109,9 @@ pub struct Config {
     /// timeout for environment verification, in milliseconds
     pub verify_env_timeout_ms: u64,
     pub proxy_url: Option<String>,
+    pub service: Option<String>,
+    pub env: Option<String>,
+    pub version: Option<String>,
 }
 
 impl Config {
@@ -251,6 +254,9 @@ impl Config {
                 .or_else(|_| env::var("HTTPS_PROXY"))
                 .ok(),
             tags,
+            service: env::var("DD_SERVICE").ok(),
+            env: env::var("DD_ENV").ok(),
+            version: env::var("DD_VERSION").ok(),
         })
     }
 }
@@ -721,6 +727,9 @@ pub mod test_helpers {
             proxy_request_retry_backoff_base_ms: 100,
             verify_env_timeout_ms: 1000,
             proxy_url: None,
+            service: None,
+            env: None,
+            version: None,
         }
     }
 }

crates/datadog-trace-agent/src/config.rs

@@ -13,7 +13,9 @@ pub mod env_verifier;
 pub mod http_utils;
 pub mod mini_agent;
 pub mod proxy_flusher;
+pub mod stats_concentrator_service;
 pub mod stats_flusher;
+pub mod stats_generator;
 pub mod stats_processor;
 pub mod trace_flusher;
 pub mod trace_processor;

crates/datadog-trace-agent/src/lib.rs

@@ -19,7 +19,10 @@ use crate::proxy_flusher::{ProxyFlusher, ProxyRequest};
 #[cfg(all(windows, feature = "windows-pipes"))]
 use tokio::net::windows::named_pipe::ServerOptions;
 
-use crate::{config, env_verifier, stats_flusher, stats_processor, trace_flusher, trace_processor};
+use crate::{
+    config, env_verifier, stats_concentrator_service, stats_flusher, stats_generator,
+    stats_processor, trace_flusher, trace_processor,
+};
 use libdd_trace_protobuf::pb;
 use libdd_trace_utils::trace_utils;
 use libdd_trace_utils::trace_utils::SendData;
@@ -40,6 +43,8 @@ pub struct MiniAgent {
     pub stats_flusher: Arc<dyn stats_flusher::StatsFlusher + Send + Sync>,
     pub env_verifier: Arc<dyn env_verifier::EnvVerifier + Send + Sync>,
     pub proxy_flusher: Arc<ProxyFlusher>,
+    pub stats_concentrator: Option<stats_concentrator_service::StatsConcentratorHandle>,
+    pub stats_generator: Option<Arc<stats_generator::StatsGenerator>>,
 }
 
 impl MiniAgent {