Send proxy payload through mpsc channel to then get aggregated and flushed in proxy flusher

Author: kathiehuang

Cargo.lock

@@ -17,10 +17,12 @@ use tracing_subscriber::EnvFilter;
 use zstd::zstd_safe::CompressionLevel;
 
 use datadog_trace_agent::{
-    aggregator::TraceAggregator,
+    trace_aggregator::TraceAggregator,
     config, env_verifier, mini_agent, stats_flusher, stats_processor,
     trace_flusher::{self, TraceFlusher},
     trace_processor,
+    proxy_aggregator,
+    proxy_flusher,
 };
 
 use datadog_trace_utils::{config_utils::read_cloud_env, trace_utils::EnvironmentType};
@@ -120,13 +122,20 @@ pub async fn main() {
         Arc::clone(&config),
     ));
 
+    let proxy_aggregator = Arc::new(TokioMutex::new(proxy_aggregator::ProxyAggregator::default()));
+    let proxy_flusher = Arc::new(proxy_flusher::ProxyFlusher::new(
+        proxy_aggregator,
+        Arc::clone(&config),
+    ));
+
     let mini_agent = Box::new(mini_agent::MiniAgent {
         config: Arc::clone(&config),
         env_verifier,
         trace_processor,
         trace_flusher,
         stats_processor,
         stats_flusher,
+        proxy_flusher,
     });
 
     tokio::spawn(async move {

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

@@ -23,6 +23,7 @@ datadog-trace-protobuf = { git = "https://github.com/DataDog/libdatadog/", rev =
 datadog-trace-utils = { git = "https://github.com/DataDog/libdatadog/", rev = "4eb2b8673354f974591c61bab3f7d485b4c119e0", features = ["mini_agent"] }
 datadog-trace-normalization = { git = "https://github.com/DataDog/libdatadog/", rev = "4eb2b8673354f974591c61bab3f7d485b4c119e0" }
 datadog-trace-obfuscation = { git = "https://github.com/DataDog/libdatadog/", rev = "4eb2b8673354f974591c61bab3f7d485b4c119e0" }
+datadog-fips = { path = "../datadog-fips", default-features = false }
 reqwest = { version = "0.12.23", features = ["json"] }
 bytes = "1.10.1"
 

crates/datadog-trace-agent/Cargo.toml

@@ -11,7 +11,7 @@ use std::sync::OnceLock;
 use datadog_trace_obfuscation::obfuscation_config;
 use datadog_trace_utils::config_utils::{
     read_cloud_env, trace_intake_url, trace_intake_url_prefixed, trace_stats_url,
-    trace_stats_url_prefixed,
+    trace_stats_url_prefixed
 };
 use datadog_trace_utils::trace_utils;
 
@@ -86,6 +86,9 @@ pub struct Config {
     pub trace_flush_interval: u64,
     pub trace_intake: Endpoint,
     pub trace_stats_intake: Endpoint,
+    /// how often to flush proxy requests, in seconds
+    pub proxy_flush_interval: u64,
+    pub proxy_intake: Endpoint,
     /// timeout for environment verification, in milliseconds
     pub verify_env_timeout: u64,
     pub proxy_url: Option<String>,
@@ -111,6 +114,7 @@ impl Config {
         // trace stats to)
         let mut trace_intake_url = trace_intake_url(&dd_site);
         let mut trace_stats_intake_url = trace_stats_url(&dd_site);
+        let proxy_intake_url = format!("https://intake.profile.{}/api/v2/profile", dd_site);
 
         // DD_APM_DD_URL env var will primarily be used for integration tests
         // overrides the entire trace/trace stats intake url prefix
@@ -139,6 +143,7 @@ impl Config {
             max_request_content_length: 10 * 1024 * 1024, // 10MB in Bytes
             trace_flush_interval: 3,
             stats_flush_interval: 3,
+            proxy_flush_interval: 3,
             verify_env_timeout: 100,
             dd_dogstatsd_port,
             dd_site,
@@ -149,6 +154,11 @@ impl Config {
             },
             trace_stats_intake: Endpoint {
                 url: hyper::Uri::from_str(&trace_stats_intake_url).unwrap(),
+                api_key: Some(api_key.clone()),
+                ..Default::default()
+            },
+            proxy_intake: Endpoint {
+                url: hyper::Uri::from_str(&proxy_intake_url).unwrap(),
                 api_key: Some(api_key),
                 ..Default::default()
             },

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

@@ -9,6 +9,9 @@ use hyper::{
 };
 use serde_json::json;
 use tracing::{debug, error};
+use datadog_fips::reqwest_adapter::create_reqwest_client_builder;
+use core::time::Duration;
+use std::error::Error;
 
 /// Does two things:
 /// 1. Logs the given message. A success status code (within 200-299) will cause an info log to be
@@ -111,6 +114,19 @@ pub fn verify_request_content_length(
     None
 }
 
+/// Builds a reqwest client with optional proxy configuration and timeout.
+/// Uses FIPS-compliant TLS when the fips feature is enabled.
+pub fn build_client(
+    proxy_url: Option<&str>,
+    timeout: Duration,
+) -> Result<reqwest::Client, Box<dyn Error>> {
+    let mut builder = create_reqwest_client_builder()?.timeout(timeout);
+    if let Some(proxy) = proxy_url {
+        builder = builder.proxy(reqwest::Proxy::all(proxy)?);
+    }
+    Ok(builder.build()?)
+}
+
 #[cfg(test)]
 mod tests {
     use ddcommon::hyper_migration;

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

@@ -7,12 +7,13 @@
 #![cfg_attr(not(test), deny(clippy::todo))]
 #![cfg_attr(not(test), deny(clippy::unimplemented))]
 
-pub mod aggregator;
+pub mod trace_aggregator;
 pub mod config;
 pub mod env_verifier;
 pub mod http_utils;
 pub mod mini_agent;
 pub mod proxy_aggregator;
+pub mod proxy_flusher;
 pub mod stats_flusher;
 pub mod stats_processor;
 pub mod trace_flusher;

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

@@ -10,14 +10,11 @@ use std::io;
 use std::net::SocketAddr;
 use std::sync::Arc;
 use std::time::Instant;
-use tokio::sync::{
-    Mutex,
-    mpsc::{self, Receiver, Sender},
-};
+use tokio::sync::mpsc::{self, Receiver, Sender};
 use tracing::{debug, error};
 
 use crate::http_utils::{self, log_and_create_http_response};
-use crate::{config, env_verifier, stats_flusher, stats_processor, trace_flusher, trace_processor, proxy_aggregator::{self, ProxyRequest}};
+use crate::{config, env_verifier, stats_flusher, stats_processor, trace_flusher, trace_processor, proxy_aggregator::ProxyRequest, proxy_flusher};
 use datadog_trace_protobuf::pb;
 use datadog_trace_utils::trace_utils;
 use datadog_trace_utils::trace_utils::SendData;
@@ -29,6 +26,7 @@ const INFO_ENDPOINT_PATH: &str = "/info";
 const PROFILING_ENDPOINT_PATH: &str = "/profiling/v1/input";
 const TRACER_PAYLOAD_CHANNEL_BUFFER_SIZE: usize = 10;
 const STATS_PAYLOAD_CHANNEL_BUFFER_SIZE: usize = 10;
+const PROXY_PAYLOAD_CHANNEL_BUFFER_SIZE: usize = 10;
 
 pub struct MiniAgent {
     pub config: Arc<config::Config>,
@@ -37,7 +35,7 @@ pub struct MiniAgent {
     pub stats_processor: Arc<dyn stats_processor::StatsProcessor + Send + Sync>,
     pub stats_flusher: Arc<dyn stats_flusher::StatsFlusher + Send + Sync>,
     pub env_verifier: Arc<dyn env_verifier::EnvVerifier + Send + Sync>,
-    pub proxy_aggregator: Arc<Mutex<proxy_aggregator::Aggregator>>,
+    pub proxy_flusher: Arc<proxy_flusher::ProxyFlusher>,
 }
 
 impl MiniAgent {
@@ -89,18 +87,23 @@ impl MiniAgent {
                 .start_stats_flusher(stats_config, stats_rx)
                 .await;
         });
+        // channels to send processed profiling requests to our proxy flusher.
+        let (proxy_tx, proxy_rx): (
+            Sender<ProxyRequest>,
+            Receiver<ProxyRequest>,
+        ) = mpsc::channel(PROXY_PAYLOAD_CHANNEL_BUFFER_SIZE);
 
         // start our proxy flusher for profiling requests
-        let proxy_aggregator_for_flusher = self.proxy_aggregator.clone();
+        let proxy_flusher = self.proxy_flusher.clone();
         tokio::spawn(async move {
-            let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(10));
+            let proxy_flusher = proxy_flusher.clone();
+            proxy_flusher.start_proxy_flusher(proxy_rx).await;
         });
 
         // setup our hyper http server, where the endpoint_handler handles incoming requests
         let trace_processor = self.trace_processor.clone();
         let stats_processor = self.stats_processor.clone();
         let endpoint_config = self.config.clone();
-        let proxy_aggregator = self.proxy_aggregator.clone();
 
         let service = service_fn(move |req| {
             let trace_processor = trace_processor.clone();
@@ -111,7 +114,8 @@ impl MiniAgent {
 
             let endpoint_config = endpoint_config.clone();
             let mini_agent_metadata = Arc::clone(&mini_agent_metadata);
-            let proxy_aggregator = proxy_aggregator.clone();
+
+            let proxy_tx = proxy_tx.clone();
 
             MiniAgent::trace_endpoint_handler(
                 endpoint_config.clone(),
@@ -121,7 +125,7 @@ impl MiniAgent {
                 stats_processor.clone(),
                 stats_tx.clone(),
                 Arc::clone(&mini_agent_metadata),
-                proxy_aggregator.clone(),
+                proxy_tx.clone(),
             )
         });
 
@@ -185,7 +189,7 @@ impl MiniAgent {
         stats_processor: Arc<dyn stats_processor::StatsProcessor + Send + Sync>,
         stats_tx: Sender<pb::ClientStatsPayload>,
         mini_agent_metadata: Arc<trace_utils::MiniAgentMetadata>,
-        proxy_aggregator: Arc<Mutex<proxy_aggregator::Aggregator>>,
+        proxy_tx: Sender<ProxyRequest>,
     ) -> http::Result<hyper_migration::HttpResponse> {
         match (req.method(), req.uri().path()) {
             (&Method::PUT | &Method::POST, TRACE_ENDPOINT_PATH) => {
@@ -210,7 +214,7 @@ impl MiniAgent {
                 }
             }
             (&Method::POST, PROFILING_ENDPOINT_PATH) => {
-                match Self::profiling_proxy_handler(config, req, proxy_aggregator).await {
+                match Self::profiling_proxy_handler(config, req, proxy_tx).await {
                     Ok(res) => Ok(res),
                     Err(err) => log_and_create_http_response(
                         &format!("Error processing profiling request: {err}"),
@@ -236,8 +240,8 @@ impl MiniAgent {
     async fn profiling_proxy_handler(
         config: Arc<config::Config>,
         request: hyper_migration::HttpRequest,
-        proxy_aggregator: Arc<Mutex<proxy_aggregator::Aggregator>>,
-    ) -> Result<hyper_migration::HttpResponse, Box<dyn std::error::Error + Send + Sync>> {
+        proxy_tx: Sender<ProxyRequest>,
+    ) -> http::Result<hyper_migration::HttpResponse> {
         debug!("Trace Agent | Proxied request for profiling");
 
         // Extract headers and body
@@ -247,25 +251,37 @@ impl MiniAgent {
             config.max_request_content_length,
             "Error processing profiling request",
         ) {
-            return response.map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>);
+            return response;
         }
 
-        let body_bytes = body.collect().await?.to_bytes();
+        let body_bytes = match body.collect().await {
+            Ok(collected) => collected.to_bytes(),
+            Err(e) => {
+                return log_and_create_http_response(
+                    &format!("Error reading profiling request body: {e}"),
+                    StatusCode::BAD_REQUEST,
+                );
+            }
+        };
 
         // Create proxy request
         let proxy_request = ProxyRequest {
             headers: parts.headers,
             body: body_bytes,
-            target_url: format!("https://intake.profile.{}/api/v2/profile", config.dd_site),
+            target_url: config.proxy_intake.url.to_string(),
         };
 
-        let mut proxy_aggregator = proxy_aggregator.lock().await;
-        proxy_aggregator.add(proxy_request);
-
-        Response::builder()
-            .status(200)
-            .body(hyper_migration::Body::from("Acknowledged profiling request"))
-            .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)
+        // Send to channel - flusher will aggregate and send
+        match proxy_tx.send(proxy_request).await {
+            Ok(_) => log_and_create_http_response(
+                "Successfully buffered profiling request to be flushed",
+                StatusCode::OK,
+            ),
+            Err(err) => log_and_create_http_response(
+                &format!("Error sending profiling request to the proxy flusher: {err}"),
+                StatusCode::INTERNAL_SERVER_ERROR,
+            ),
+        }
     }
 
     fn info_handler(dd_dogstatsd_port: u16) -> http::Result<hyper_migration::HttpResponse> {

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

@@ -1,26 +1,27 @@
-use bytes::Bytes; // TODO: Do we use bytes?
-use reqwest::header::HeaderMap; // TODO: Do we use reqwest?
+use bytes::Bytes;
+use reqwest::header::HeaderMap;
 
+#[derive(Clone)]
 pub struct ProxyRequest {
     pub headers: HeaderMap,
     pub body: Bytes,
     pub target_url: String,
 }
 
 /// Takes in individual proxy requests and aggregates them into batches to be flushed to Datadog.
-pub struct Aggregator {
+pub struct ProxyAggregator {
     queue: Vec<ProxyRequest>,
 }
 
-impl Default for Aggregator {
+impl Default for ProxyAggregator {
     fn default() -> Self {
-        Aggregator {
+        ProxyAggregator {
             queue: Vec::with_capacity(128), // arbitrary capacity for request queue
         }
     }
 }
 
-impl Aggregator {
+impl ProxyAggregator {
     /// Takes in an individual proxy request.
     pub fn add(&mut self, request: ProxyRequest) {
         self.queue.push(request);