feat(logs): [SVLS-8582] Hold logs and add durable context to durable function logs

.gitlab/datasources/test-suites.yaml

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

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},
@@ -298,15 +299,20 @@ async fn extension_loop_active(
     // 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 +331,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 +1046,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 +1076,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/lifecycle/invocation/processor.rs

@@ -9,7 +9,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},
@@ -24,6 +26,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,
@@ -88,6 +91,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 {
@@ -98,6 +105,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()
@@ -127,6 +135,7 @@ impl Processor {
             dynamic_tags: HashMap::new(),
             active_invocations: 0,
             awaiting_first_invocation: false,
+            durable_context_tx,
         }
     }
 
@@ -1342,6 +1351,23 @@ impl Processor {
             self.context_buffer.add_tracer_span(request_id, span);
         }
     }
+
+    /// 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 fn forward_durable_context(
+        &mut self,
+        request_id: &str,
+        execution_id: &str,
+        execution_name: &str,
+    ) {
+        if let Err(e) = self.durable_context_tx.try_send(DurableContextUpdate {
+            request_id: request_id.to_owned(),
+            execution_id: execution_id.to_owned(),
+            execution_name: execution_name.to_owned(),
+        }) {
+            error!("Invocation Processor | Failed to forward durable context to logs agent: {e}");
+        }
+    }
 }
 
 #[cfg(test)]
@@ -1387,7 +1413,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]
@@ -1924,7 +1958,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(),

bottlecap/src/lifecycle/invocation/processor_service.rs

@@ -17,6 +17,7 @@ use crate::{
         context::{Context, ReparentingInfo},
         processor::Processor,
     },
+    logs::lambda::DurableContextUpdate,
     tags::provider,
     traces::{
         context::SpanContext, propagation::DatadogCompositePropagator,
@@ -111,6 +112,11 @@ pub enum ProcessorCommand {
     AddTracerSpan {
         span: Box<Span>,
     },
+    ForwardDurableContext {
+        request_id: String,
+        execution_id: String,
+        execution_name: String,
+    },
     OnOutOfMemoryError {
         timestamp: i64,
     },
@@ -380,6 +386,21 @@ impl InvocationProcessorHandle {
             .await
     }
 
+    pub async fn forward_durable_context(
+        &self,
+        request_id: String,
+        execution_id: String,
+        execution_name: String,
+    ) -> Result<(), mpsc::error::SendError<ProcessorCommand>> {
+        self.sender
+            .send(ProcessorCommand::ForwardDurableContext {
+                request_id,
+                execution_id,
+                execution_name,
+            })
+            .await
+    }
+
     pub async fn on_out_of_memory_error(
         &self,
         timestamp: i64,
@@ -430,6 +451,7 @@ impl InvocationProcessorService {
         aws_config: Arc<AwsConfig>,
         metrics_aggregator_handle: AggregatorHandle,
         propagator: Arc<DatadogCompositePropagator>,
+        durable_context_tx: mpsc::Sender<DurableContextUpdate>,
     ) -> (InvocationProcessorHandle, Self) {
         let (sender, receiver) = mpsc::channel(1000);
 
@@ -439,6 +461,7 @@ impl InvocationProcessorService {
             aws_config,
             metrics_aggregator_handle,
             propagator,
+            durable_context_tx,
         );
 
         let handle = InvocationProcessorHandle { sender };
@@ -588,6 +611,17 @@ impl InvocationProcessorService {
                 ProcessorCommand::AddTracerSpan { span } => {
                     self.processor.add_tracer_span(&span);
                 }
+                ProcessorCommand::ForwardDurableContext {
+                    request_id,
+                    execution_id,
+                    execution_name,
+                } => {
+                    self.processor.forward_durable_context(
+                        &request_id,
+                        &execution_id,
+                        &execution_name,
+                    );
+                }
                 ProcessorCommand::OnOutOfMemoryError { timestamp } => {
                     self.processor.on_out_of_memory_error(timestamp);
                 }

bottlecap/src/logs/agent.rs

@@ -12,9 +12,12 @@ use crate::{LAMBDA_RUNTIME_SLUG, config};
 
 const DRAIN_LOG_INTERVAL: Duration = Duration::from_millis(100);
 
+use crate::logs::lambda::DurableContextUpdate;
+
 #[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 +31,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 +41,18 @@ impl LogsAgent {
         );
 
         let (tx, rx) = mpsc::channel::<TelemetryEvent>(1000);
+        let (durable_context_tx, durable_context_rx) = mpsc::channel::<DurableContextUpdate>(500);
         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 +61,9 @@ impl LogsAgent {
                 Some(event) = self.rx.recv() => {
                     self.processor.process(event, &self.aggregator_handle).await;
                 }
+                Some(update) = self.durable_context_rx.recv() => {
+                    self.processor.process_durable_context_update(update, &self.aggregator_handle);
+                }
                 () = self.cancel_token.cancelled() => {
                     debug!("LOGS_AGENT | Received shutdown signal, draining remaining events");
 

bottlecap/src/logs/aggregator.rs

@@ -106,6 +106,7 @@ mod tests {
                 lambda: Lambda {
                     arn: "arn".to_string(),
                     request_id: Some("request_id".to_string()),
+                    ..Lambda::default()
                 },
                 timestamp: 0,
                 status: "status".to_string(),
@@ -130,6 +131,7 @@ mod tests {
                 lambda: Lambda {
                     arn: "arn".to_string(),
                     request_id: Some("request_id".to_string()),
+                    ..Lambda::default()
                 },
                 timestamp: 0,
                 status: "status".to_string(),
@@ -156,6 +158,7 @@ mod tests {
                 lambda: Lambda {
                     arn: "arn".to_string(),
                     request_id: Some("request_id".to_string()),
+                    ..Lambda::default()
                 },
                 timestamp: 0,
                 status: "status".to_string(),
@@ -196,6 +199,7 @@ mod tests {
                 lambda: Lambda {
                     arn: "arn".to_string(),
                     request_id: Some("request_id".to_string()),
+                    ..Lambda::default()
                 },
                 timestamp: 0,
                 status: "status".to_string(),

bottlecap/src/logs/aggregator_service.rs

@@ -122,6 +122,7 @@ mod tests {
                 lambda: Lambda {
                     arn: "arn".to_string(),
                     request_id: Some("request_id".to_string()),
+                    ..Lambda::default()
                 },
                 timestamp: 0,
                 status: "status".to_string(),