Handle proxy payload

Author: kathiehuang

crates/datadog-trace-agent/Cargo.toml

@@ -23,6 +23,8 @@ 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" }
+reqwest = { version = "0.12.23", features = ["json"] }
+bytes = "1.10.1"
 
 [dev-dependencies]
 rmp-serde = "1.1.1"

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

@@ -12,6 +12,7 @@ pub mod config;
 pub mod env_verifier;
 pub mod http_utils;
 pub mod mini_agent;
+pub mod proxy_aggregator;
 pub mod stats_flusher;
 pub mod stats_processor;
 pub mod trace_flusher;

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

@@ -2,18 +2,22 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use ddcommon::hyper_migration;
+use http_body_util::BodyExt;
 use hyper::service::service_fn;
 use hyper::{http, Method, Response, StatusCode};
 use serde_json::json;
 use std::io;
 use std::net::SocketAddr;
 use std::sync::Arc;
 use std::time::Instant;
-use tokio::sync::mpsc::{self, Receiver, Sender};
+use tokio::sync::{
+    Mutex,
+    mpsc::{self, Receiver, Sender},
+};
 use tracing::{debug, error};
 
-use crate::http_utils::log_and_create_http_response;
-use crate::{config, env_verifier, stats_flusher, stats_processor, trace_flusher, trace_processor};
+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 datadog_trace_protobuf::pb;
 use datadog_trace_utils::trace_utils;
 use datadog_trace_utils::trace_utils::SendData;
@@ -22,6 +26,7 @@ const MINI_AGENT_PORT: usize = 8126;
 const TRACE_ENDPOINT_PATH: &str = "/v0.4/traces";
 const STATS_ENDPOINT_PATH: &str = "/v0.6/stats";
 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;
 
@@ -32,6 +37,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>>,
 }
 
 impl MiniAgent {
@@ -84,10 +90,17 @@ impl MiniAgent {
                 .await;
         });
 
+        // start our proxy flusher for profiling requests
+        let proxy_aggregator_for_flusher = self.proxy_aggregator.clone();
+        tokio::spawn(async move {
+            let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(10));
+        });
+
         // 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();
@@ -98,6 +111,7 @@ impl MiniAgent {
 
             let endpoint_config = endpoint_config.clone();
             let mini_agent_metadata = Arc::clone(&mini_agent_metadata);
+            let proxy_aggregator = proxy_aggregator.clone();
 
             MiniAgent::trace_endpoint_handler(
                 endpoint_config.clone(),
@@ -107,6 +121,7 @@ impl MiniAgent {
                 stats_processor.clone(),
                 stats_tx.clone(),
                 Arc::clone(&mini_agent_metadata),
+                proxy_aggregator.clone(),
             )
         });
 
@@ -170,6 +185,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>>,
     ) -> http::Result<hyper_migration::HttpResponse> {
         match (req.method(), req.uri().path()) {
             (&Method::PUT | &Method::POST, TRACE_ENDPOINT_PATH) => {
@@ -193,6 +209,15 @@ impl MiniAgent {
                     ),
                 }
             }
+            (&Method::POST, PROFILING_ENDPOINT_PATH) => {
+                match Self::profiling_proxy_handler(config, req, proxy_aggregator).await {
+                    Ok(res) => Ok(res),
+                    Err(err) => log_and_create_http_response(
+                        &format!("Error processing profiling request: {err}"),
+                        StatusCode::INTERNAL_SERVER_ERROR,
+                    ),
+                }
+            }
             (_, INFO_ENDPOINT_PATH) => match Self::info_handler(config.dd_dogstatsd_port) {
                 Ok(res) => Ok(res),
                 Err(err) => log_and_create_http_response(
@@ -208,13 +233,49 @@ 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>> {
+        debug!("Trace Agent | Proxied request for profiling");
+
+        // Extract headers and body
+        let (parts, body) = request.into_parts();
+        if let Some(response) = http_utils::verify_request_content_length(
+            &parts.headers,
+            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>);
+        }
+
+        let body_bytes = body.collect().await?.to_bytes();
+
+        // 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),
+        };
+
+        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>)
+    }
+
     fn info_handler(dd_dogstatsd_port: u16) -> http::Result<hyper_migration::HttpResponse> {
         let response_json = json!(
             {
                 "endpoints": [
                     TRACE_ENDPOINT_PATH,
                     STATS_ENDPOINT_PATH,
-                    INFO_ENDPOINT_PATH
+                    INFO_ENDPOINT_PATH,
+                    PROFILING_ENDPOINT_PATH
                 ],
                 "client_drop_p0s": true,
                 "config": {
@@ -226,4 +287,5 @@ impl MiniAgent {
             .status(200)
             .body(hyper_migration::Body::from(response_json.to_string()))
     }
+
 }

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

@@ -0,0 +1,38 @@
+use bytes::Bytes; // TODO: Do we use bytes?
+use reqwest::header::HeaderMap; // TODO: Do we use reqwest?
+
+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 {
+    queue: Vec<ProxyRequest>,
+}
+
+impl Default for Aggregator {
+    fn default() -> Self {
+        Aggregator {
+            queue: Vec::with_capacity(128), // arbitrary capacity for request queue
+        }
+    }
+}
+
+impl Aggregator {
+    /// Takes in an individual proxy request.
+    pub fn add(&mut self, request: ProxyRequest) {
+        self.queue.push(request);
+    }
+
+    /// Returns a batch of proxy requests.
+    pub fn get_batch(&mut self) -> Vec<ProxyRequest> {
+        std::mem::take(&mut self.queue)
+    }
+
+    /// Flush the queue.
+    pub fn clear(&mut self) {
+        self.queue.clear();
+    }
+}