feat(logs): detect durable function context from aws.lambda span tags

When an aws.lambda span (resource = dd-tracer-serverless-span) carries
durable_function_execution_id and durable_function_execution_name meta
tags, forward the context to the logs agent via a dedicated mpsc channel
so held logs can be tagged and released.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: lym953

bottlecap/src/bin/bottlecap/main.rs

@@ -46,7 +46,7 @@ use bottlecap::{
     },
     logger,
     logs::{
-        agent::LogsAgent,
+        agent::{DurableContextUpdate, LogsAgent},
         aggregator_service::{
             AggregatorHandle as LogsAggregatorHandle, AggregatorService as LogsAggregatorService,
         },
@@ -291,14 +291,19 @@ 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);
 
-    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(),
-        );
+    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(),
+    );
 
     let (metrics_flushers, metrics_aggregator_handle, dogstatsd_cancel_token) =
         start_dogstatsd(tags_provider.clone(), Arc::clone(&api_key_factory), config).await;
@@ -344,6 +349,7 @@ async fn extension_loop_active(
         &tags_provider,
         invocation_processor_handle.clone(),
         appsec_processor.clone(),
+        durable_context_tx,
     );
 
     let api_runtime_proxy_shutdown_signal = start_api_runtime_proxy(
@@ -1024,14 +1030,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,
@@ -1048,7 +1055,7 @@ fn start_logs_agent(
     });
 
     let flusher = LogsFlusher::new(api_key_factory, aggregator_handle.clone(), config.clone());
-    (tx, flusher, cancel_token, aggregator_handle)
+    (tx, flusher, cancel_token, aggregator_handle, durable_context_tx)
 }
 
 #[allow(clippy::type_complexity)]
@@ -1058,6 +1065,7 @@ fn start_trace_agent(
     tags_provider: &Arc<TagProvider>,
     invocation_processor_handle: InvocationProcessorHandle,
     appsec_processor: Option<Arc<TokioMutex<AppSecProcessor>>>,
+    durable_context_tx: Sender<DurableContextUpdate>,
 ) -> (
     Sender<SendDataBuilderInfo>,
     Arc<trace_flusher::TraceFlusher>,
@@ -1130,6 +1138,7 @@ fn start_trace_agent(
         Arc::clone(tags_provider),
         stats_concentrator_handle.clone(),
         span_dedup_handle,
+        durable_context_tx,
     );
     let trace_agent_channel = trace_agent.get_sender_copy();
     let shutdown_token = trace_agent.shutdown_token();

bottlecap/src/logs/agent.rs

@@ -12,9 +12,13 @@ use crate::{LAMBDA_RUNTIME_SLUG, config};
 
 const DRAIN_LOG_INTERVAL: Duration = Duration::from_millis(100);
 
+/// `(request_id, execution_id, execution_name)` extracted from an `aws.lambda` span.
+pub type DurableContextUpdate = (String, String, String);
+
 #[allow(clippy::module_name_repetitions)]
 pub struct LogsAgent {
     rx: mpsc::Receiver<TelemetryEvent>,
+    durable_context_rx: mpsc::Receiver<DurableContextUpdate>,
     processor: LogsProcessor,
     aggregator_handle: AggregatorHandle,
     cancel_token: CancellationToken,
@@ -28,7 +32,7 @@ impl LogsAgent {
         event_bus: Sender<Event>,
         aggregator_handle: AggregatorHandle,
         is_managed_instance_mode: bool,
-    ) -> (Self, Sender<TelemetryEvent>) {
+    ) -> (Self, Sender<TelemetryEvent>, Sender<DurableContextUpdate>) {
         let processor = LogsProcessor::new(
             Arc::clone(&datadog_config),
             tags_provider,
@@ -38,16 +42,18 @@ impl LogsAgent {
         );
 
         let (tx, rx) = mpsc::channel::<TelemetryEvent>(1000);
+        let (durable_context_tx, durable_context_rx) = mpsc::channel::<DurableContextUpdate>(100);
         let cancel_token = CancellationToken::new();
 
         let agent = Self {
             rx,
+            durable_context_rx,
             processor,
             aggregator_handle,
             cancel_token,
         };
 
-        (agent, tx)
+        (agent, tx, durable_context_tx)
     }
 
     pub async fn spin(&mut self) {
@@ -56,6 +62,13 @@ impl LogsAgent {
                 Some(event) = self.rx.recv() => {
                     self.processor.process(event, &self.aggregator_handle).await;
                 }
+                Some((request_id, execution_id, execution_name)) = self.durable_context_rx.recv() => {
+                    self.processor.update_durable_map(&request_id, &execution_id, &execution_name);
+                    let ready = self.processor.take_ready_logs();
+                    if !ready.is_empty() {
+                        let _ = self.aggregator_handle.insert_batch(ready);
+                    }
+                }
                 () = self.cancel_token.cancelled() => {
                     debug!("LOGS_AGENT | Received shutdown signal, draining remaining events");
 

bottlecap/src/logs/lambda/processor.rs

@@ -208,7 +208,8 @@ impl LambdaProcessor {
                 let rv = runtime_version.unwrap_or("?".to_string()); // TODO: check what does containers display
                 let rv_arn = runtime_version_arn.unwrap_or("?".to_string()); // TODO: check what do containers display
 
-                self.is_durable_function = Some(rv.contains("DurableFunction"));
+                let is_durable = rv.contains("DurableFunction");
+                self.is_durable_function = Some(is_durable);
                 self.resolve_held_logs_on_durable_function_set();
 
                 Ok(Message::new(
@@ -493,37 +494,35 @@ impl LambdaProcessor {
         }
     }
 
-    /// Parses `message` as JSON and, if it contains both `durable_execution_id` and
-    /// `durable_execution_name` fields, inserts or updates the entry for `request_id` in the
-    /// durable ID map (evicting the oldest entry when the map is at capacity 5).
-    ///
-    /// Returns `true` if a brand-new entry was added (the caller may then drain `held_logs`
-    /// for that `request_id`).
-    fn try_update_durable_map(&mut self, request_id: &str, message: &str) -> bool {
-        let Ok(serde_json::Value::Object(obj)) = serde_json::from_str(message) else {
-            return false;
-        };
-        let execution_id = obj
-            .get("durable_execution_id")
-            .and_then(serde_json::Value::as_str);
-        let execution_name = obj
-            .get("durable_execution_name")
-            .and_then(serde_json::Value::as_str);
-        if let (Some(id), Some(name)) = (execution_id, execution_name) {
-            let is_new = !self.durable_id_map.contains_key(request_id);
-            if is_new {
-                if self.durable_id_order.len() >= DURABLE_ID_MAP_CAPACITY
-                    && let Some(oldest) = self.durable_id_order.pop_front()
-                {
-                    self.durable_id_map.remove(&oldest);
-                }
-                self.durable_id_order.push_back(request_id.to_string());
+    /// Records the durable execution context for a `request_id`, received from the `aws.lambda`
+    /// span. Evicts the oldest entry when the map is at capacity. If a new entry is added and
+    /// `is_durable_function` is already `Some(true)`, drains any held logs for that `request_id`.
+    pub fn update_durable_map(
+        &mut self,
+        request_id: &str,
+        execution_id: &str,
+        execution_name: &str,
+    ) {
+        let is_new = !self.durable_id_map.contains_key(request_id);
+        if is_new {
+            if self.durable_id_order.len() >= DURABLE_ID_MAP_CAPACITY
+                && let Some(oldest) = self.durable_id_order.pop_front()
+            {
+                self.durable_id_map.remove(&oldest);
             }
-            self.durable_id_map
-                .insert(request_id.to_string(), (id.to_string(), name.to_string()));
-            return is_new;
+            self.durable_id_order.push_back(request_id.to_string());
+        }
+        self.durable_id_map.insert(
+            request_id.to_string(),
+            (execution_id.to_string(), execution_name.to_string()),
+        );
+        if is_new && self.is_durable_function == Some(true) {
+            self.drain_held_for_request_id(request_id);
         }
-        false
+    }
+
+    pub fn take_ready_logs(&mut self) -> Vec<String> {
+        std::mem::take(&mut self.ready_logs)
     }
 
     /// Moves all logs held for `request_id` into `ready_logs`, tagging each with the
@@ -567,28 +566,23 @@ impl LambdaProcessor {
                 }
             }
             Some(true) => {
+                // Durable context is populated by span processing, not from log messages.
+                // Flush any held logs whose request_id is already in the map; keep the rest.
                 for (request_id, logs) in held {
-                    // Try to discover durable context from the held logs themselves.
-                    for log in &logs {
-                        self.try_update_durable_map(&request_id, &log.message.message);
-                    }
                     let tags_suffix = self.durable_id_map.get(&request_id).map(|(id, name)| {
                         format!(",durable_execution_id:{id},durable_execution_name:{name}")
                     });
                     // Borrow of durable_id_map released here.
-                    match tags_suffix {
-                        Some(suffix) => {
-                            for mut log in logs {
-                                log.tags.push_str(&suffix);
-                                if let Ok(s) = serde_json::to_string(&log) {
-                                    self.ready_logs.push(s);
-                                }
+                    if let Some(suffix) = tags_suffix {
+                        for mut log in logs {
+                            log.tags.push_str(&suffix);
+                            if let Ok(s) = serde_json::to_string(&log) {
+                                self.ready_logs.push(s);
                             }
                         }
-                        None => {
-                            // No context yet — put back and wait for it to arrive.
-                            self.held_logs.insert(request_id, logs);
-                        }
+                    } else {
+                        // No context yet — wait for the aws.lambda span to arrive.
+                        self.held_logs.insert(request_id, logs);
                     }
                 }
             }
@@ -611,22 +605,18 @@ impl LambdaProcessor {
     /// - `None`        → stash in `held_logs[request_id]`; logs without a `request_id` are
     ///   flushed immediately since they cannot carry durable context.
     /// - `Some(false)` → serialize and push straight to `ready_logs`.
-    /// - `Some(true)`  → try to update `durable_id_map` from the log; if a new entry was
-    ///   added, drain `held_logs` for that `request_id`; then flush this log if
-    ///   its `request_id` is in the map, otherwise stash it in `held_logs`.
+    /// - `Some(true)`  → flush if this log's `request_id` is already in `durable_id_map`
+    ///   (context was populated by an `aws.lambda` span); otherwise stash in `held_logs`.
     fn queue_log_after_rules(&mut self, mut log: IntakeLog) {
         match self.is_durable_function {
             None => {
-                match log.message.lambda.request_id.clone() {
-                    Some(rid) => {
-                        self.held_logs.entry(rid).or_default().push(log);
-                    }
-                    None => {
-                        // No request_id — cannot associate with durable context; flush now.
-                        if let Ok(s) = serde_json::to_string(&log) {
-                            drop(log);
-                            self.ready_logs.push(s);
-                        }
+                if let Some(rid) = log.message.lambda.request_id.clone() {
+                    self.held_logs.entry(rid).or_default().push(log);
+                } else {
+                    // No request_id — cannot associate with durable context; flush now.
+                    if let Ok(s) = serde_json::to_string(&log) {
+                        drop(log);
+                        self.ready_logs.push(s);
                     }
                 }
             }
@@ -639,14 +629,8 @@ impl LambdaProcessor {
                 }
             }
             Some(true) => {
-                if let Some(rid) = log.message.lambda.request_id.clone()
-                    && self.try_update_durable_map(&rid, &log.message.message)
-                {
-                    // New durable context just discovered — drain previously held logs.
-                    self.drain_held_for_request_id(&rid);
-                }
-
-                // Flush this log if its request_id now has durable context; otherwise hold.
+                // Durable context is populated by span processing (update_durable_map).
+                // Flush this log if its request_id already has context; otherwise hold.
                 let durable_tags = log
                     .message
                     .lambda
@@ -671,7 +655,7 @@ impl LambdaProcessor {
                         if let Some(rid) = log.message.lambda.request_id.clone() {
                             self.held_logs.entry(rid).or_default().push(log);
                         }
-                        // Logs without a request_id cannot match the durable map; drop them.
+                        // No request_id in durable mode: drop the log.
                     }
                 }
             }
@@ -1932,7 +1916,7 @@ mod tests {
             },
         };
         let start_msg = processor.get_message(start_event).await.unwrap();
-        processor.get_intake_log(start_msg).unwrap();
+        processor.get_intake_log(start_msg, false).unwrap();
 
         // Test WARN level
         let event = TelemetryEvent {
@@ -2017,7 +2001,7 @@ mod tests {
             },
         };
         let start_msg = processor.get_message(start_event).await.unwrap();
-        processor.get_intake_log(start_msg).unwrap();
+        processor.get_intake_log(start_msg, false).unwrap();
 
         // JSON without level field
         let event = TelemetryEvent {
@@ -2058,7 +2042,7 @@ mod tests {
             },
         };
         let start_msg = processor.get_message(start_event).await.unwrap();
-        processor.get_intake_log(start_msg).unwrap();
+        processor.get_intake_log(start_msg, false).unwrap();
 
         // JSON with unrecognized level
         let event = TelemetryEvent {
@@ -2189,7 +2173,7 @@ mod tests {
             },
         };
         let start_msg = processor.get_message(start_event).await.unwrap();
-        processor.get_intake_log(start_msg).unwrap();
+        processor.get_intake_log(start_msg, false).unwrap();
 
         // level as a number
         let event = TelemetryEvent {
@@ -2252,7 +2236,7 @@ mod tests {
             },
         };
         let start_msg = processor.get_message(start_event).await.unwrap();
-        processor.get_intake_log(start_msg).unwrap();
+        processor.get_intake_log(start_msg, false).unwrap();
 
         // JSON log with both ddtags and level
         let event = TelemetryEvent {
@@ -2301,7 +2285,7 @@ mod tests {
             },
         };
         let start_msg = processor.get_message(start_event).await.unwrap();
-        processor.get_intake_log(start_msg).unwrap();
+        processor.get_intake_log(start_msg, false).unwrap();
 
         // JSON log with "status" field instead of "level" (Datadog convention)
         let event = TelemetryEvent {

bottlecap/src/logs/processor.rs

@@ -41,6 +41,25 @@ impl LogsProcessor {
             }
         }
     }
+
+    pub fn update_durable_map(
+        &mut self,
+        request_id: &str,
+        execution_id: &str,
+        execution_name: &str,
+    ) {
+        match self {
+            LogsProcessor::Lambda(p) => {
+                p.update_durable_map(request_id, execution_id, execution_name);
+            }
+        }
+    }
+
+    pub fn take_ready_logs(&mut self) -> Vec<String> {
+        match self {
+            LogsProcessor::Lambda(p) => p.take_ready_logs(),
+        }
+    }
 }
 
 #[allow(clippy::module_name_repetitions)]