add end invocation call

alexgallotta · alexgallotta · commit d33d1648b3bc · 2025-03-27T15:47:50.000-04:00
diff --git a/bottlecap/src/lifecycle/invocation/processor.rs b/bottlecap/src/lifecycle/invocation/processor.rs
@@ -458,6 +458,7 @@ impl Processor {
 
         // Set the extracted trace context to the spans
         if let Some(sc) = &self.extracted_span_context {
+            debug!("LIP: found span context in headers {:?}", sc);
             self.span.trace_id = sc.trace_id;
             self.span.parent_id = sc.span_id;
 
@@ -476,6 +477,9 @@ impl Processor {
         if let Some(inferred_span) = &self.inferrer.inferred_span {
             self.span.parent_id = inferred_span.span_id;
         }
+
+        debug!("LIP: trace context found trace_id:{trace_id} span_id:{span_id} parent_id:{parent_id}",
+            trace_id = self.span.trace_id, span_id = self.span.span_id, parent_id = self.span.parent_id);
         (self.span.span_id, self.span.trace_id, self.span.parent_id)
     }
 
diff --git a/bottlecap/src/lifecycle/listener.rs b/bottlecap/src/lifecycle/listener.rs
@@ -7,14 +7,13 @@ use std::net::SocketAddr;
 use std::sync::Arc;
 
 use hyper::body::HttpBody;
-use hyper::http::request::Parts;
 use hyper::service::{make_service_fn, service_fn};
 use hyper::{http, Body, HeaderMap, Method, Request, Response, StatusCode};
 use serde_json::json;
 use tokio::sync::Mutex;
 use tracing::{debug, error, warn};
 
-use crate::lifecycle::invocation::processor::{Processor as InvocationProcessor, Processor};
+use crate::lifecycle::invocation::processor::{Processor as InvocationProcessor};
 use crate::traces::propagation::text_map_propagator::{
     DATADOG_HIGHER_ORDER_TRACE_ID_BITS_KEY, DATADOG_SAMPLING_PRIORITY_KEY, DATADOG_TAGS_KEY,
     DATADOG_TRACE_ID_KEY,
@@ -66,7 +65,8 @@ impl Listener {
                 Self::start_invocation_handler(parts.headers, body, invocation_processor).await
             }
             (&Method::POST, END_INVOCATION_PATH) => {
-                match Self::end_invocation_handler(req, invocation_processor).await {
+                let (parts, body) = req.into_parts();
+                match Self::end_invocation_handler(parts.headers, body, invocation_processor).await {
                     Ok(response) => Ok(response),
                     Err(e) => {
                         error!("Failed to end invocation {e}");
@@ -139,17 +139,17 @@ impl Listener {
     }
 
     pub async fn end_invocation_handler(
-        req: Request<Body>,
+        headers: HeaderMap,
+        body: Body,
         invocation_processor: Arc<Mutex<InvocationProcessor>>,
     ) -> http::Result<Response<Body>> {
         debug!("Received end invocation request");
-        let (parts, body) = req.into_parts();
         match body.collect().await {
             Ok(b) => {
                 let body = b.to_bytes().to_vec();
                 let mut processor = invocation_processor.lock().await;
 
-                let headers = Self::headers_to_map(parts.headers);
+                let headers = Self::headers_to_map(headers);
                 processor.on_invocation_end(headers, body);
                 drop(processor);
 
diff --git a/bottlecap/src/lwa/proxy.rs b/bottlecap/src/lwa/proxy.rs
@@ -1,7 +1,6 @@
 use crate::{
     lifecycle::invocation::processor::Processor,
     lifecycle::listener::Listener,
-    traces::trace_processor::{ServerlessTraceProcessor, TraceProcessor},
 };
 use hyper::http::request::Parts;
 use hyper::{
@@ -11,9 +10,8 @@ use hyper::{
     Body, Client, Error, Request, Response, Server, Uri,
 };
 use hyper_proxy::{Intercept, Proxy, ProxyConnector};
-use rand::random;
 use serde_json::Value;
-use std::{collections::HashMap, net::SocketAddr, sync::Arc};
+use std::{net::SocketAddr, sync::Arc};
 use tokio::{sync::Mutex, task::JoinHandle};
 use tracing::{debug, error};
 
@@ -87,7 +85,7 @@ fn parse_env_addresses() -> Option<(SocketAddr, Uri)> {
                 if let (Some(host), Some(port)) = (host, port) {
                     if host == "localhost" {
                         error!("LWA: Cannot use localhost as host in AWS_LWA_PROXY_LAMBDA_RUNTIME_API, use 127.0.0.1 instead");
-                        None
+                        return None;
                     }
                     format!("{host}:{port}")
                         .parse::<SocketAddr>()
@@ -113,7 +111,6 @@ fn parse_env_addresses() -> Option<(SocketAddr, Uri)> {
             }
         },
         Err(e) => {
-            error!("LWA: Error parsing AWS_LWA_PROXY_LAMBDA_RUNTIME_API: {}", e);
             None
         }
     };
@@ -148,8 +145,9 @@ async fn intercept_payload(
     let (intercepted_parts, intercepted_body) = intercepted.into_parts();
     debug!("LWA: Intercepted request: {:?}", intercepted_parts);
 
+    let waited_intercepted_body = intercepted_body.collect().await?.to_bytes();
     let forward_intercepted =
-        forward_request(aws_runtime_addr, &intercepted_parts, intercepted_body).await?;
+        forward_request(aws_runtime_addr, &intercepted_parts, waited_intercepted_body.clone().into()).await?;
 
     // response after forwarding to AWS runtime API
     let response_to_intercepted_req = client.request(forward_intercepted).await?;
@@ -190,7 +188,6 @@ async fn intercept_payload(
         {
             // intercepted response to runtime/invocation. The *request* contains the returned
             // values and headers from lambda handler
-            // let parsed_body = serde_json::from_slice::<Value>(&request_body_waited);
             // let _ = Listener::start_invocation_handler(
             //     req,
             //     processor,
@@ -201,6 +198,12 @@ async fn intercept_payload(
             //     parsed_body,
             // )
             // .await;
+            let _ = Listener::end_invocation_handler(
+                intercepted_parts.headers,
+                waited_intercepted_body.into(),
+                Arc::clone(&processor),
+            )
+                .await;
             // only parsing of the original request (handler -> runtime API) is needed so
             // the original response can be used
             Ok(response_to_intercepted_req)
@@ -269,18 +272,18 @@ async fn forward_request(
 // }
 // }
 
-fn deserialize_json(response: Result<Bytes, Error>) -> Option<Value> {
-    match response {
-        Ok(bytes) => serde_json::from_slice(bytes.as_ref()).unwrap_or_else(|e| {
-            error!("Error deserializing response body: {}", e);
-            None
-        }),
-        Err(e) => {
-            error!("Error reading response body: {}", e);
-            None
-        }
-    }
-}
+// fn deserialize_json(response: Result<Bytes, Error>) -> Option<Value> {
+//     match response {
+//         Ok(bytes) => serde_json::from_slice(bytes.as_ref()).unwrap_or_else(|e| {
+//             error!("Error deserializing response body: {}", e);
+//             None
+//         }),
+//         Err(e) => {
+//             error!("Error reading response body: {}", e);
+//             None
+//         }
+//     }
+// }
 
 #[cfg(test)]
 mod tests {

Original file line number	Diff line number	Diff line change
`@@ -458,6 +458,7 @@ impl Processor {`
`458`	`458`
`459`	`459`	`// Set the extracted trace context to the spans`
`460`	`460`	`if let Some(sc) = &self.extracted_span_context {`
	`461`	`+ debug!("LIP: found span context in headers {:?}", sc);`
`461`	`462`	`self.span.trace_id = sc.trace_id;`
`462`	`463`	`self.span.parent_id = sc.span_id;`
`463`	`464`
`@@ -476,6 +477,9 @@ impl Processor {`
`476`	`477`	`if let Some(inferred_span) = &self.inferrer.inferred_span {`
`477`	`478`	`self.span.parent_id = inferred_span.span_id;`
`478`	`479`	`}`
	`480`	`+`
	`481`	`+ debug!("LIP: trace context found trace_id:{trace_id} span_id:{span_id} parent_id:{parent_id}",`
	`482`	`+ trace_id = self.span.trace_id, span_id = self.span.span_id, parent_id = self.span.parent_id);`
`479`	`483`	`(self.span.span_id, self.span.trace_id, self.span.parent_id)`
`480`	`484`	`}`
`481`	`485`