fix(lifecycle): spawn tasks on listener (#796)

# What?

Make Lifecycle Listener for Universal Instrumentation to spawn tasks to
process data

# Motivation

There is no need to await for it to finish in order to continue, also
[SVLS-7464](https://datadoghq.atlassian.net/browse/SVLS-7464)

# Notes

Still trying to figure out how to handle invocation end

[SVLS-7464]:
https://datadoghq.atlassian.net/browse/SVLS-7464?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ

Author: duncanista

bottlecap/src/bin/bottlecap/main.rs

@@ -9,7 +9,6 @@
 #![deny(missing_copy_implementations)]
 #![deny(missing_debug_implementations)]
 
-use bottlecap::traces::trace_processor::SendingTraceProcessor;
 #[cfg(not(target_env = "msvc"))]
 use tikv_jemallocator::Jemalloc;
 
@@ -21,8 +20,9 @@ use bottlecap::{
     DOGSTATSD_PORT, EXTENSION_ACCEPT_FEATURE_HEADER, EXTENSION_FEATURES, EXTENSION_HOST,
     EXTENSION_HOST_IP, EXTENSION_ID_HEADER, EXTENSION_NAME, EXTENSION_NAME_HEADER, EXTENSION_ROUTE,
     LAMBDA_RUNTIME_SLUG, TELEMETRY_PORT,
-    appsec::processor::Error::FeatureDisabled as AppSecFeatureDisabled,
-    appsec::processor::Processor as AppSecProcessor,
+    appsec::processor::{
+        Error::FeatureDisabled as AppSecFeatureDisabled, Processor as AppSecProcessor,
+    },
     base_url,
     config::{
         self, Config,
@@ -50,14 +50,15 @@ use bottlecap::{
         listener::TelemetryListener,
     },
     traces::{
+        propagation::DatadogCompositePropagator,
         proxy_aggregator,
         proxy_flusher::Flusher as ProxyFlusher,
         stats_aggregator::StatsAggregator,
         stats_flusher::{self, StatsFlusher},
         stats_processor, trace_agent,
         trace_aggregator::{self, SendDataBuilderInfo},
         trace_flusher::{self, ServerlessTraceFlusher, TraceFlusher},
-        trace_processor,
+        trace_processor::{self, SendingTraceProcessor},
     },
 };
 use datadog_fips::reqwest_adapter::create_reqwest_client_builder;
@@ -514,12 +515,15 @@ async fn extension_loop_active(
         &metrics_aggr_handle,
         config,
     )));
+
+    let propagator = Arc::new(DatadogCompositePropagator::new(Arc::clone(config)));
     // Lifecycle Invocation Processor
     let invocation_processor = Arc::new(TokioMutex::new(InvocationProcessor::new(
         Arc::clone(&tags_provider),
         Arc::clone(config),
         Arc::clone(&aws_config),
         metrics_aggr_handle.clone(),
+        Arc::clone(&propagator),
     )));
     // AppSec processor (if enabled)
     let appsec_processor = match AppSecProcessor::new(config) {
@@ -555,11 +559,12 @@ async fn extension_loop_active(
         aws_config,
         &invocation_processor,
         appsec_processor.as_ref(),
+        Arc::clone(&propagator),
     );
 
-    let lifecycle_listener = LifecycleListener {
-        invocation_processor: Arc::clone(&invocation_processor),
-    };
+    let lifecycle_listener =
+        LifecycleListener::new(Arc::clone(&invocation_processor), Arc::clone(&propagator));
+    let lifecycle_listener_shutdown_token = lifecycle_listener.get_shutdown_token();
     // TODO(astuyve): deprioritize this task after the first request
     tokio::spawn(async move {
         let res = lifecycle_listener.start().await;
@@ -801,6 +806,7 @@ async fn extension_loop_active(
             trace_agent_shutdown_token.cancel();
             dogstatsd_cancel_token.cancel();
             telemetry_listener_cancel_token.cancel();
+            lifecycle_listener_shutdown_token.cancel();
 
             // gotta lock here
             let mut locked_metrics = metrics_flushers.lock().await;
@@ -1237,6 +1243,7 @@ fn start_api_runtime_proxy(
     aws_config: Arc<AwsConfig>,
     invocation_processor: &Arc<TokioMutex<InvocationProcessor>>,
     appsec_processor: Option<&Arc<TokioMutex<AppSecProcessor>>>,
+    propagator: Arc<DatadogCompositePropagator>,
 ) -> Option<CancellationToken> {
     if !should_start_proxy(config, Arc::clone(&aws_config)) {
         debug!("Skipping API runtime proxy, no LWA proxy or datadog wrapper found");
@@ -1245,5 +1252,11 @@ fn start_api_runtime_proxy(
 
     let invocation_processor = invocation_processor.clone();
     let appsec_processor = appsec_processor.map(Arc::clone);
-    interceptor::start(aws_config, invocation_processor, appsec_processor).ok()
+    interceptor::start(
+        aws_config,
+        invocation_processor,
+        appsec_processor,
+        propagator,
+    )
+    .ok()
 }

bottlecap/src/http.rs

@@ -1,14 +1,14 @@
 use crate::config;
 use axum::{
     extract::{FromRequest, Request},
-    http::{self, StatusCode},
+    http::{self, HeaderMap, StatusCode},
     response::{IntoResponse, Response},
 };
 use bytes::Bytes;
 use core::time::Duration;
 use datadog_fips::reqwest_adapter::create_reqwest_client_builder;
-use std::error::Error;
 use std::sync::Arc;
+use std::{collections::HashMap, error::Error};
 use tracing::error;
 
 #[must_use]
@@ -68,3 +68,16 @@ pub async fn extract_request_body(
 
     Ok((parts, bytes))
 }
+
+#[must_use]
+pub fn headers_to_map(headers: &HeaderMap) -> HashMap<String, String> {
+    headers
+        .iter()
+        .map(|(k, v)| {
+            (
+                k.as_str().to_string(),
+                v.to_str().unwrap_or_default().to_string(),
+            )
+        })
+        .collect()
+}

bottlecap/src/lifecycle/invocation/context.rs

@@ -8,6 +8,7 @@ use std::{
 };
 
 use datadog_trace_protobuf::pb::Span;
+use serde_json::Value;
 use tracing::debug;
 
 #[derive(Debug, Clone, PartialEq)]
@@ -127,7 +128,7 @@ pub struct ContextBuffer {
 
 struct UniversalInstrumentationData {
     headers: HashMap<String, String>,
-    payload: Vec<u8>,
+    payload_value: Value,
 }
 
 impl Default for ContextBuffer {
@@ -227,12 +228,14 @@ impl ContextBuffer {
     pub fn pair_invoke_event(
         &mut self,
         request_id: &str,
-    ) -> Option<(HashMap<String, String>, Vec<u8>)> {
-        if let Some(UniversalInstrumentationData { headers, payload }) =
-            self.universal_instrumentation_start_events.pop_front()
+    ) -> Option<(HashMap<String, String>, Value)> {
+        if let Some(UniversalInstrumentationData {
+            headers,
+            payload_value,
+        }) = self.universal_instrumentation_start_events.pop_front()
         {
             // Bad scenario, we found an `UniversalInstrumentationStart`
-            Some((headers, payload))
+            Some((headers, payload_value))
         } else {
             // `UniversalInstrumentationStart` event hasn't occurred yet, this is good,
             // push the Invoke event to the queue and return `None`
@@ -248,7 +251,7 @@ impl ContextBuffer {
     pub fn pair_universal_instrumentation_start_event(
         &mut self,
         headers: &HashMap<String, String>,
-        payload: &[u8],
+        payload_value: &Value,
     ) -> Option<String> {
         if let Some(request_id) = self.invoke_events_request_ids.pop_front() {
             // Bad scenario, we found an `UniversalInstrumentationStart`
@@ -259,7 +262,7 @@ impl ContextBuffer {
             self.universal_instrumentation_start_events
                 .push_back(UniversalInstrumentationData {
                     headers: headers.clone(),
-                    payload: payload.to_vec(),
+                    payload_value: payload_value.clone(),
                 });
             None
         }
@@ -271,7 +274,7 @@ impl ContextBuffer {
     pub fn pair_universal_instrumentation_end_event(
         &mut self,
         headers: &HashMap<String, String>,
-        payload: &[u8],
+        payload_value: &Value,
     ) -> Option<String> {
         if let Some(request_id) = self.platform_runtime_done_events_request_ids.pop_front() {
             // Bad scenario, we found a `PlatformRuntimeDone`
@@ -282,7 +285,7 @@ impl ContextBuffer {
             self.universal_instrumentation_end_events
                 .push_back(UniversalInstrumentationData {
                     headers: headers.clone(),
-                    payload: payload.to_vec(),
+                    payload_value: payload_value.clone(),
                 });
             None
         }
@@ -294,12 +297,14 @@ impl ContextBuffer {
     pub fn pair_platform_runtime_done_event(
         &mut self,
         request_id: &str,
-    ) -> Option<(HashMap<String, String>, Vec<u8>)> {
-        if let Some(UniversalInstrumentationData { headers, payload }) =
-            self.universal_instrumentation_end_events.pop_front()
+    ) -> Option<(HashMap<String, String>, Value)> {
+        if let Some(UniversalInstrumentationData {
+            headers,
+            payload_value,
+        }) = self.universal_instrumentation_end_events.pop_front()
         {
             // Good scenario, we found an `UniversalInstrumentationEnd`
-            Some((headers, payload))
+            Some((headers, payload_value))
         } else {
             // `UniversalInstrumentationEnd` hasn't occurred yet, this is bad,
             // push the `PlatformRuntimeDone` event to the queue and return `None`
@@ -425,6 +430,7 @@ impl ContextBuffer {
 #[allow(clippy::unwrap_used)]
 mod tests {
     use crate::proc::{CPUData, NetworkData};
+    use serde_json::json;
     use std::collections::HashMap;
     use tokio::sync::watch;
 
@@ -619,14 +625,14 @@ mod tests {
         let request_id = "test-request-1".to_string();
         let mut headers = HashMap::new();
         headers.insert("test-header".to_string(), "test-value".to_string());
-        let payload = vec![1, 2, 3];
+        let payload = json!({ "test": "payload" });
 
         // Test case 1: UniversalInstrumentationStart arrives before Invoke
         buffer
             .universal_instrumentation_start_events
             .push_back(UniversalInstrumentationData {
                 headers: headers.clone(),
-                payload: payload.clone(),
+                payload_value: payload.clone(),
             });
 
         // When Invoke arrives, it should pair with the existing UniversalInstrumentationStart
@@ -653,7 +659,7 @@ mod tests {
         let request_id = "test-request-1".to_string();
         let mut headers = HashMap::new();
         headers.insert("test-header".to_string(), "test-value".to_string());
-        let payload = vec![1, 2, 3];
+        let payload = json!({ "test": "payload" });
 
         // Test case 1: Invoke arrives before UniversalInstrumentationStart
         buffer
@@ -674,7 +680,7 @@ mod tests {
             .front()
             .unwrap();
         assert_eq!(stored_event.headers, headers);
-        assert_eq!(stored_event.payload, payload);
+        assert_eq!(stored_event.payload_value, payload);
     }
 
     #[test]
@@ -683,7 +689,7 @@ mod tests {
         let request_id = "test-request-1".to_string();
         let mut headers = HashMap::new();
         headers.insert("test-header".to_string(), "test-value".to_string());
-        let payload = vec![1, 2, 3];
+        let payload = json!({ "test": "payload" });
 
         // Test case 1: PlatformRuntimeDone arrives before UniversalInstrumentationEnd
         buffer
@@ -701,7 +707,7 @@ mod tests {
         assert_eq!(buffer.universal_instrumentation_end_events.len(), 1);
         let stored_event = buffer.universal_instrumentation_end_events.front().unwrap();
         assert_eq!(stored_event.headers, headers);
-        assert_eq!(stored_event.payload, payload);
+        assert_eq!(stored_event.payload_value, payload);
     }
 
     #[test]
@@ -710,14 +716,14 @@ mod tests {
         let request_id = "test-request-1".to_string();
         let mut headers = HashMap::new();
         headers.insert("test-header".to_string(), "test-value".to_string());
-        let payload = vec![1, 2, 3];
+        let payload = json!({ "test": "payload" });
 
         // Test case 1: UniversalInstrumentationEnd arrives before PlatformRuntimeDone
         buffer
             .universal_instrumentation_end_events
             .push_back(UniversalInstrumentationData {
                 headers: headers.clone(),
-                payload: payload.clone(),
+                payload_value: payload.clone(),
             });
 
         // When PlatformRuntimeDone arrives, it should pair with the existing UniversalInstrumentationEnd
@@ -747,15 +753,15 @@ mod tests {
         let request_id = "test-request-1".to_string();
         let mut headers = HashMap::new();
         headers.insert("test-header".to_string(), "test-value".to_string());
-        let payload = vec![1, 2, 3];
+        let payload = json!({ "test": "payload" });
 
         // Test the complete flow with events arriving in different orders
         // Case 1: Events arrive in reverse order
         buffer
             .universal_instrumentation_end_events
             .push_back(UniversalInstrumentationData {
                 headers: headers.clone(),
-                payload: payload.clone(),
+                payload_value: payload.clone(),
             });
         buffer
             .platform_runtime_done_events_request_ids