Merge branch 'main' into tianning.li/SVLS-8755-arn-secret-key

Author: litianningdatadog

.gitignore

@@ -30,3 +30,4 @@ integration-tests/cdk.context.json
 
 .gitlab/pipeline*
 /CLAUDE.md
+/AGENTS.md

.gitlab/datasources/test-suites.yaml

@@ -3,3 +3,4 @@ test_suites:
   - name: otlp
   - name: snapstart
   - name: lmi
+  - name: auth

README.md

@@ -21,7 +21,7 @@ For `v67` through `v87`, you can opt out of the next-generation Lambda Extension
 
 Today, all workloads using Logs and Metrics are supported.
 
-APM Tracing is supported for Python, NodeJS, Go, Java, and .NET runtimes.
+APM Tracing is supported for Python, NodeJS, Go, Java, .NET, and Ruby runtimes.
 
 ### Feedback
 

bottlecap/src/bin/bottlecap/main.rs

@@ -51,6 +51,7 @@ use bottlecap::{
             AggregatorHandle as LogsAggregatorHandle, AggregatorService as LogsAggregatorService,
         },
         flusher::LogsFlusher,
+        lambda::DurableContextUpdate,
     },
     otlp::{agent::Agent as OtlpAgent, should_enable_otlp_agent},
     proxy::{interceptor, should_start_proxy},
@@ -149,7 +150,11 @@ async fn main() -> anyhow::Result<()> {
     let config = Arc::new(config::get_config(Path::new(&lambda_directory)));
 
     let aws_config = Arc::new(aws_config);
-    let api_key_factory = create_api_key_factory(&config, &aws_config);
+    // Build one shared reqwest::Client for metrics, logs, trace proxy flushing, and calls to
+    // Datadog APIs (e.g. delegated auth). reqwest::Client is Arc-based internally, so cloning
+    // just increments a refcount and shares the connection pool.
+    let shared_client = bottlecap::http::get_client(&config);
+    let api_key_factory = create_api_key_factory(&config, &aws_config, &shared_client);
 
     let r = response
         .await
@@ -160,6 +165,7 @@ async fn main() -> anyhow::Result<()> {
         Arc::clone(&aws_config),
         &config,
         &client,
+        shared_client,
         &r,
         Arc::clone(&api_key_factory),
         start_time,
@@ -245,17 +251,23 @@ fn get_flush_strategy_for_mode(
     }
 }
 
-fn create_api_key_factory(config: &Arc<Config>, aws_config: &Arc<AwsConfig>) -> Arc<ApiKeyFactory> {
+fn create_api_key_factory(
+    config: &Arc<Config>,
+    aws_config: &Arc<AwsConfig>,
+    client: &reqwest::Client,
+) -> Arc<ApiKeyFactory> {
     let config = Arc::clone(config);
     let aws_config = Arc::clone(aws_config);
+    let client = client.clone();
     let api_key_secret_reload_interval = config.api_key_secret_reload_interval;
 
     Arc::new(ApiKeyFactory::new_from_resolver(
         Arc::new(move || {
             let config = Arc::clone(&config);
             let aws_config = Arc::clone(&aws_config);
+            let client = client.clone();
 
-            Box::pin(async move { resolve_secrets(config, aws_config).await })
+            Box::pin(async move { resolve_secrets(config, aws_config, client).await })
         }),
         api_key_secret_reload_interval,
     ))
@@ -284,6 +296,7 @@ async fn extension_loop_active(
     aws_config: Arc<AwsConfig>,
     config: &Arc<Config>,
     client: &Client,
+    shared_client: reqwest::Client,
     r: &RegisterResponse,
     api_key_factory: Arc<ApiKeyFactory>,
     start_time: Instant,
@@ -293,20 +306,20 @@ async fn extension_loop_active(
     let account_id = r.account_id.as_ref().unwrap_or(&"none".to_string()).clone();
     let tags_provider = setup_tag_provider(&Arc::clone(&aws_config), config, &account_id);
 
-    // Build one shared reqwest::Client for metrics, logs, and trace proxy flushing.
-    // reqwest::Client is Arc-based internally, so cloning just increments a refcount
-    // and shares the connection pool.
-    let shared_client = bottlecap::http::get_client(config);
-
-    let (logs_agent_channel, logs_flusher, logs_agent_cancel_token, logs_aggregator_handle) =
-        start_logs_agent(
-            config,
-            Arc::clone(&api_key_factory),
-            &tags_provider,
-            event_bus_tx.clone(),
-            aws_config.is_managed_instance_mode(),
-            &shared_client,
-        );
+    let (
+        logs_agent_channel,
+        logs_flusher,
+        logs_agent_cancel_token,
+        logs_aggregator_handle,
+        durable_context_tx,
+    ) = start_logs_agent(
+        config,
+        Arc::clone(&api_key_factory),
+        &tags_provider,
+        event_bus_tx.clone(),
+        aws_config.is_managed_instance_mode(),
+        &shared_client,
+    );
 
     let (metrics_flushers, metrics_aggregator_handle, dogstatsd_cancel_token) = start_dogstatsd(
         tags_provider.clone(),
@@ -325,6 +338,7 @@ async fn extension_loop_active(
             Arc::clone(&aws_config),
             metrics_aggregator_handle.clone(),
             Arc::clone(&propagator),
+            durable_context_tx,
         );
     tokio::spawn(async move {
         invocation_processor_service.run().await;
@@ -1039,14 +1053,15 @@ fn start_logs_agent(
     LogsFlusher,
     CancellationToken,
     LogsAggregatorHandle,
+    Sender<DurableContextUpdate>,
 ) {
     let (aggregator_service, aggregator_handle) = LogsAggregatorService::default();
     // Start service in background
     tokio::spawn(async move {
         aggregator_service.run().await;
     });
 
-    let (mut agent, tx) = LogsAgent::new(
+    let (mut agent, tx, durable_context_tx) = LogsAgent::new(
         Arc::clone(tags_provider),
         Arc::clone(config),
         event_bus,
@@ -1068,7 +1083,13 @@ fn start_logs_agent(
         config.clone(),
         client.clone(),
     );
-    (tx, flusher, cancel_token, aggregator_handle)
+    (
+        tx,
+        flusher,
+        cancel_token,
+        aggregator_handle,
+        durable_context_tx,
+    )
 }
 
 #[allow(clippy::type_complexity)]

bottlecap/src/config/env.rs

@@ -482,6 +482,12 @@ pub struct EnvConfig {
     /// The delay between two samples of the API Security schema collection, in seconds.
     #[serde(deserialize_with = "deserialize_optional_duration_from_seconds")]
     pub api_security_sample_delay: Option<Duration>,
+
+    /// @env `DD_ORG_UUID`
+    ///
+    /// The Datadog organization UUID. When set, delegated auth is auto-enabled.
+    #[serde(deserialize_with = "deserialize_string_or_int")]
+    pub org_uuid: Option<String>,
 }
 
 #[allow(clippy::too_many_lines)]
@@ -684,6 +690,8 @@ fn merge_config(config: &mut Config, env_config: &EnvConfig) {
     merge_option_to_value!(config, env_config, appsec_waf_timeout);
     merge_option_to_value!(config, env_config, api_security_enabled);
     merge_option_to_value!(config, env_config, api_security_sample_delay);
+
+    merge_string!(config, dd_org_uuid, env_config, org_uuid);
 }
 
 #[derive(Debug, PartialEq, Clone, Copy)]
@@ -1044,6 +1052,8 @@ mod tests {
                 appsec_waf_timeout: Duration::from_secs(1),
                 api_security_enabled: false,
                 api_security_sample_delay: Duration::from_secs(60),
+
+                dd_org_uuid: String::default(),
             };
 
             assert_eq!(config, expected_config);

bottlecap/src/config/mod.rs

@@ -364,6 +364,8 @@ pub struct Config {
     pub span_dedup_timeout: Option<Duration>,
     pub api_key_secret_reload_interval: Option<Duration>,
 
+    pub dd_org_uuid: String,
+
     pub serverless_appsec_enabled: bool,
     pub appsec_rules: Option<String>,
     pub appsec_waf_timeout: Duration,
@@ -479,6 +481,8 @@ impl Default for Config {
             span_dedup_timeout: None,
             api_key_secret_reload_interval: None,
 
+            dd_org_uuid: String::default(),
+
             serverless_appsec_enabled: false,
             appsec_rules: None,
             appsec_waf_timeout: Duration::from_millis(5),

bottlecap/src/config/yaml.rs

@@ -1036,6 +1036,8 @@ api_security_sample_delay: 60 # Seconds
                 dogstatsd_so_rcvbuf: Some(1_048_576),
                 dogstatsd_buffer_size: Some(65507),
                 dogstatsd_queue_size: Some(2048),
+
+                dd_org_uuid: String::default(),
             };
 
             // Assert that

bottlecap/src/lifecycle/invocation/processor.rs

@@ -16,7 +16,9 @@ use libdd_trace_protobuf::pb::Span;
 use libdd_trace_utils::tracer_header_tags;
 use serde_json::Value;
 use tokio::time::Instant;
-use tracing::{debug, trace, warn};
+use tracing::{debug, error, trace, warn};
+
+use tokio::sync::mpsc;
 
 use crate::{
     config::{self, aws::AwsConfig},
@@ -31,6 +33,7 @@ use crate::{
         span_inferrer::{self, SpanInferrer},
         triggers::get_default_service_name,
     },
+    logs::lambda::DurableContextUpdate,
     metrics::enhanced::lambda::{EnhancedMetricData, Lambda as EnhancedMetrics},
     proc::{
         self, CPUData, NetworkData,
@@ -89,6 +92,10 @@ pub struct Processor {
     /// Tracks whether if first invocation after init has been received in Managed Instance mode.
     /// Used to determine if we should search for the empty context on an invocation.
     awaiting_first_invocation: bool,
+    /// Sender used to forward durable execution context extracted from `aws.lambda` spans to the
+    /// logs agent. Decouples the trace agent from the logs agent: the trace agent sends spans
+    /// to the lifecycle processor, which extracts durable context and relays it here.
+    durable_context_tx: mpsc::Sender<DurableContextUpdate>,
 }
 
 impl Processor {
@@ -99,6 +106,7 @@ impl Processor {
         aws_config: Arc<AwsConfig>,
         metrics_aggregator: dogstatsd::aggregator::AggregatorHandle,
         propagator: Arc<DatadogCompositePropagator>,
+        durable_context_tx: mpsc::Sender<DurableContextUpdate>,
     ) -> Self {
         let resource = tags_provider
             .get_canonical_resource_name()
@@ -128,6 +136,7 @@ impl Processor {
             dynamic_tags: HashMap::new(),
             active_invocations: 0,
             awaiting_first_invocation: false,
+            durable_context_tx,
         }
     }
 
@@ -1358,6 +1367,29 @@ impl Processor {
                 .add_tracer_span(request_id, span, client_computed_stats);
         }
     }
+
+    /// Forwards durable execution context extracted from an `aws.lambda` span to the logs
+    /// pipeline so it can release held logs and tag them with durable execution metadata.
+    pub async fn forward_durable_context(
+        &mut self,
+        request_id: &str,
+        execution_id: &str,
+        execution_name: &str,
+        first_invocation: Option<bool>,
+    ) {
+        if let Err(e) = self
+            .durable_context_tx
+            .send(DurableContextUpdate {
+                request_id: request_id.to_owned(),
+                execution_id: execution_id.to_owned(),
+                execution_name: execution_name.to_owned(),
+                first_invocation,
+            })
+            .await
+        {
+            error!("Invocation Processor | Failed to forward durable context to logs agent: {e}");
+        }
+    }
 }
 
 #[cfg(test)]
@@ -1403,7 +1435,15 @@ mod tests {
         tokio::spawn(service.run());
 
         let propagator = Arc::new(DatadogCompositePropagator::new(Arc::clone(&config)));
-        Processor::new(tags_provider, config, aws_config, handle, propagator)
+        let (durable_context_tx, _) = tokio::sync::mpsc::channel(1);
+        Processor::new(
+            tags_provider,
+            config,
+            aws_config,
+            handle,
+            propagator,
+            durable_context_tx,
+        )
     }
 
     #[test]
@@ -1940,7 +1980,15 @@ mod tests {
 
         let propagator = Arc::new(DatadogCompositePropagator::new(Arc::clone(&config)));
 
-        let processor = Processor::new(tags_provider, config, aws_config, handle, propagator);
+        let (durable_context_tx, _) = tokio::sync::mpsc::channel(1);
+        let processor = Processor::new(
+            tags_provider,
+            config,
+            aws_config,
+            handle,
+            propagator,
+            durable_context_tx,
+        );
 
         assert!(
             processor.is_managed_instance_mode(),
@@ -2160,12 +2208,14 @@ mod tests {
             AggregatorService::new(EMPTY_TAGS, 1024).expect("failed to create aggregator service");
         tokio::spawn(aggregator_service.run());
         let propagator = Arc::new(DatadogCompositePropagator::new(Arc::clone(&config)));
+        let (durable_context_tx, _) = tokio::sync::mpsc::channel(1);
         let mut p = Processor::new(
             Arc::clone(&tags_provider),
             Arc::clone(&config),
             aws_config,
             aggregator_handle,
             propagator,
+            durable_context_tx,
         );
 
         let (trace_tx, mut trace_rx) = mpsc::channel(10);