[SVLS-7018] Support profiling in Azure Functions (#40)

* Handle proxy payload

* Send proxy payload through mpsc channel and handle forwarding and retries in background task

* Add additional tags in profiles

* Respect DD_APM_PROFILING_DD_URL for profiling intake url override

* Detect environment when setting _dd.origin tag

Author: kathiehuang

Cargo.lock

@@ -18,7 +18,7 @@ use zstd::zstd_safe::CompressionLevel;
 
 use datadog_trace_agent::{
     aggregator::TraceAggregator,
-    config, env_verifier, mini_agent, stats_flusher, stats_processor,
+    config, env_verifier, mini_agent, proxy_flusher, stats_flusher, stats_processor,
     trace_flusher::{self, TraceFlusher},
     trace_processor,
 };
@@ -124,13 +124,16 @@ pub async fn main() {
         Arc::clone(&config),
     ));
 
+    let proxy_flusher = Arc::new(proxy_flusher::ProxyFlusher::new(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 {

LICENSE-3rdparty.csv

@@ -26,6 +26,9 @@ libdd-trace-protobuf = { git = "https://github.com/DataDog/libdatadog", rev = "4
 libdd-trace-utils = { git = "https://github.com/DataDog/libdatadog", rev = "435107c245112397914935c0f7148a18b91cafc6", features = [
     "mini_agent",
 ] }
+datadog-fips = { path = "../datadog-fips" }
+reqwest = { version = "0.12.23", features = ["json", "http2"], default-features = false }
+bytes = "1.10.1"
 
 [dev-dependencies]
 rmp-serde = "1.1.1"

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

@@ -81,13 +81,21 @@ pub struct Config {
     pub os: String,
     pub tags: Tags,
     /// how often to flush stats, in seconds
-    pub stats_flush_interval: u64,
+    pub stats_flush_interval_secs: u64,
     /// how often to flush traces, in seconds
-    pub trace_flush_interval: u64,
+    pub trace_flush_interval_secs: u64,
     pub trace_intake: Endpoint,
     pub trace_stats_intake: Endpoint,
+    /// Profiling intake endpoint (for proxying profiling data to Datadog)
+    pub profiling_intake: Endpoint,
+    /// Timeout for each proxy request, in seconds
+    pub proxy_request_timeout_secs: u64,
+    /// Maximum number of retry attempts for failed proxy requests
+    pub proxy_request_max_retries: u32,
+    /// Base backoff duration for proxy request retries, in milliseconds
+    pub proxy_request_retry_backoff_base_ms: u64,
     /// timeout for environment verification, in milliseconds
-    pub verify_env_timeout: u64,
+    pub verify_env_timeout_ms: u64,
     pub proxy_url: Option<String>,
 }
 
@@ -119,6 +127,14 @@ impl Config {
             trace_stats_intake_url = trace_stats_url_prefixed(&endpoint_prefix);
         };
 
+        // TODO: Create helper functions for this in libdatadog
+        let mut profiling_intake_url = format!("https://intake.profile.{}/api/v2/profile", dd_site);
+        // DD_APM_PROFILING_DD_URL env var will primarily be used for integration tests
+        // overrides the prefix of the profiling intake url
+        if let Ok(endpoint_prefix) = env::var("DD_APM_PROFILING_DD_URL") {
+            profiling_intake_url = format!("{endpoint_prefix}/api/v2/profile");
+        };
+
         let obfuscation_config = obfuscation_config::ObfuscationConfig::new().map_err(|err| {
             anyhow::anyhow!(
                 "Error creating obfuscation config, Mini Agent will not start. Error: {err}",
@@ -137,9 +153,12 @@ impl Config {
             env_type,
             os: env::consts::OS.to_string(),
             max_request_content_length: 10 * 1024 * 1024, // 10MB in Bytes
-            trace_flush_interval: 3,
-            stats_flush_interval: 3,
-            verify_env_timeout: 100,
+            trace_flush_interval_secs: 3,
+            stats_flush_interval_secs: 3,
+            proxy_request_timeout_secs: 30,
+            proxy_request_max_retries: 3,
+            proxy_request_retry_backoff_base_ms: 100,
+            verify_env_timeout_ms: 100,
             dd_dogstatsd_port,
             dd_site,
             trace_intake: Endpoint {
@@ -149,6 +168,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()
+            },
+            profiling_intake: Endpoint {
+                url: hyper::Uri::from_str(&profiling_intake_url).unwrap(),
                 api_key: Some(api_key),
                 ..Default::default()
             },

crates/datadog-trace-agent/Cargo.toml

@@ -1,13 +1,16 @@
 // Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/
 // SPDX-License-Identifier: Apache-2.0
 
+use core::time::Duration;
+use datadog_fips::reqwest_adapter::create_reqwest_client_builder;
 use hyper::{
     header,
     http::{self, HeaderMap},
     Response, StatusCode,
 };
 use libdd_common::hyper_migration;
 use serde_json::json;
+use std::error::Error;
 use tracing::{debug, error};
 
 /// Does two things:
@@ -111,6 +114,19 @@ pub fn verify_request_content_length(
     None
 }
 
+/// Builds a reqwest client with optional proxy configuration and timeout.
+/// Uses rustls TLS by default. FIPS-compliant TLS is available via the fips feature
+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::https(proxy)?);
+    }
+    Ok(builder.build()?)
+}
+
 #[cfg(test)]
 mod tests {
     use http_body_util::BodyExt;

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

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

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

@@ -1,6 +1,7 @@
 // Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/
 // SPDX-License-Identifier: Apache-2.0
 
+use http_body_util::BodyExt;
 use hyper::service::service_fn;
 use hyper::{http, Method, Response, StatusCode};
 use libdd_common::hyper_migration;
@@ -12,7 +13,8 @@ use std::time::Instant;
 use tokio::sync::mpsc::{self, Receiver, Sender};
 use tracing::{debug, error};
 
-use crate::http_utils::log_and_create_http_response;
+use crate::http_utils::{log_and_create_http_response, verify_request_content_length};
+use crate::proxy_flusher::{ProxyFlusher, ProxyRequest};
 use crate::{config, env_verifier, stats_flusher, stats_processor, trace_flusher, trace_processor};
 use libdd_trace_protobuf::pb;
 use libdd_trace_utils::trace_utils;
@@ -22,8 +24,10 @@ 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;
+const PROXY_PAYLOAD_CHANNEL_BUFFER_SIZE: usize = 10;
 
 pub struct MiniAgent {
     pub config: Arc<config::Config>,
@@ -32,6 +36,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_flusher: Arc<ProxyFlusher>,
 }
 
 impl MiniAgent {
@@ -42,7 +47,7 @@ impl MiniAgent {
         let mini_agent_metadata = Arc::new(
             self.env_verifier
                 .verify_environment(
-                    self.config.verify_env_timeout,
+                    self.config.verify_env_timeout_ms,
                     &self.config.env_type,
                     &self.config.os,
                 )
@@ -82,6 +87,17 @@ impl MiniAgent {
                 .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_flusher = self.proxy_flusher.clone();
+        tokio::spawn(async move {
+            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();
@@ -96,14 +112,17 @@ impl MiniAgent {
             let endpoint_config = endpoint_config.clone();
             let mini_agent_metadata = Arc::clone(&mini_agent_metadata);
 
+            let proxy_tx = proxy_tx.clone();
+
             MiniAgent::trace_endpoint_handler(
                 endpoint_config,
                 req.map(hyper_migration::Body::incoming),
-                trace_processor,
-                trace_tx,
-                stats_processor,
-                stats_tx,
-                mini_agent_metadata,
+                trace_processor.clone(),
+                trace_tx.clone(),
+                stats_processor.clone(),
+                stats_tx.clone(),
+                Arc::clone(&mini_agent_metadata),
+                proxy_tx.clone(),
             )
         });
 
@@ -167,6 +186,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_tx: Sender<ProxyRequest>,
     ) -> http::Result<hyper_migration::HttpResponse> {
         match (req.method(), req.uri().path()) {
             (&Method::PUT | &Method::POST, TRACE_ENDPOINT_PATH) => {
@@ -190,6 +210,15 @@ impl MiniAgent {
                     ),
                 }
             }
+            (&Method::POST, PROFILING_ENDPOINT_PATH) => {
+                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}"),
+                        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(
@@ -205,13 +234,67 @@ impl MiniAgent {
         }
     }
 
+    /// Handles incoming proxy requests for profiling - can be abstracted into a generic proxy handler for other proxy requests in the future
+    async fn profiling_proxy_handler(
+        config: Arc<config::Config>,
+        request: hyper_migration::HttpRequest,
+        proxy_tx: Sender<ProxyRequest>,
+    ) -> http::Result<hyper_migration::HttpResponse> {
+        debug!("Trace Agent | Received profiling request");
+
+        // Extract headers and body
+        let (parts, body) = request.into_parts();
+        if let Some(response) = verify_request_content_length(
+            &parts.headers,
+            config.max_request_content_length,
+            "Error processing profiling request",
+        ) {
+            return response;
+        }
+
+        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: config.profiling_intake.url.to_string(),
+        };
+
+        debug!(
+            "Trace Agent | Sending profiling request to channel, target: {}",
+            proxy_request.target_url
+        );
+
+        // Send to channel
+        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> {
         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": {