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

@@ -46,7 +46,7 @@ use bottlecap::{
     },
     logger,
     logs::{
-        agent::LogsAgent,
+        agent::{DurableContextUpdate, LogsAgent},
         aggregator_service::{
             AggregatorHandle as LogsAggregatorHandle, AggregatorService as LogsAggregatorService,
         },
@@ -298,15 +298,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(),
@@ -357,6 +362,7 @@ async fn extension_loop_active(
         &tags_provider,
         invocation_processor_handle.clone(),
         appsec_processor.clone(),
+        durable_context_tx,
         &shared_client,
     );
 
@@ -1039,14 +1045,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 +1075,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)]
@@ -1078,6 +1091,7 @@ fn start_trace_agent(
     tags_provider: &Arc<TagProvider>,
     invocation_processor_handle: InvocationProcessorHandle,
     appsec_processor: Option<Arc<TokioMutex<AppSecProcessor>>>,
+    durable_context_tx: Sender<DurableContextUpdate>,
     client: &Client,
 ) -> (
     Sender<SendDataBuilderInfo>,
@@ -1163,6 +1177,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

@@ -2,7 +2,7 @@ use std::sync::Arc;
 use std::time::{Duration, Instant};
 use tokio::sync::mpsc::{self, Sender};
 use tokio_util::sync::CancellationToken;
-use tracing::debug;
+use tracing::{debug, error};
 
 use crate::event_bus::Event;
 use crate::extension::telemetry::events::TelemetryEvent;
@@ -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.insert_to_durable_map(&request_id, &execution_id, &execution_name);
+                    let ready_logs = self.processor.take_ready_logs();
+                    if !ready_logs.is_empty() && let Err(e) = self.aggregator_handle.insert_batch(ready_logs) {
+                        error!("LOGS_AGENT | Failed to insert batch: {}", e);
+                    }
+                }
                 () = 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(),

bottlecap/src/logs/lambda/mod.rs

@@ -30,10 +30,14 @@ pub struct Message {
     pub status: String,
 }
 
-#[derive(Serialize, Debug, Clone, PartialEq)]
+#[derive(Serialize, Debug, Clone, Default, PartialEq)]
 pub struct Lambda {
     pub arn: String,
     pub request_id: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub durable_execution_id: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub durable_execution_name: Option<String>,
 }
 
 impl Message {
@@ -50,6 +54,7 @@ impl Message {
             lambda: Lambda {
                 arn: function_arn,
                 request_id,
+                ..Lambda::default()
             },
             timestamp,
             status: status.unwrap_or("info".to_string()),