From b39c2136460dc2cfbd9d8142191c6fc11f1a974b Mon Sep 17 00:00:00 2001 From: Leonard Gerard Date: Thu, 15 Dec 2022 10:56:55 -0800 Subject: [PATCH 1/3] Add tracing span with request id to the handler --- lambda-runtime/src/lib.rs | 120 +++++++++++++++++++++----------------- 1 file changed, 68 insertions(+), 52 deletions(-) diff --git a/lambda-runtime/src/lib.rs b/lambda-runtime/src/lib.rs index 7664a499..6bcaf22b 100644 --- a/lambda-runtime/src/lib.rs +++ b/lambda-runtime/src/lib.rs @@ -26,7 +26,7 @@ use tokio::io::{AsyncRead, AsyncWrite}; use tokio_stream::{Stream, StreamExt}; pub use tower::{self, service_fn, Service}; use tower::{util::ServiceFn, ServiceExt}; -use tracing::{error, trace}; +use tracing::{error, span, trace, Instrument}; mod requests; #[cfg(test)] @@ -120,15 +120,6 @@ where continue; } - let body = hyper::body::to_bytes(body).await?; - trace!("response body - {}", std::str::from_utf8(&body)?); - - #[cfg(debug_assertions)] - if parts.status.is_server_error() { - error!("Lambda Runtime server returned an unexpected error"); - return Err(parts.status.to_string().into()); - } - let ctx: Context = Context::try_from(parts.headers)?; let ctx: Context = ctx.with_config(config); let request_id = &ctx.request_id.clone(); @@ -138,55 +129,80 @@ where Some(trace_id) => env::set_var("_X_AMZN_TRACE_ID", trace_id), None => env::remove_var("_X_AMZN_TRACE_ID"), } - let body = match serde_json::from_slice(&body) { - Ok(body) => body, - Err(err) => { - let req = build_event_error_request(request_id, err)?; - client.call(req).await.expect("Unable to send response to Runtime APIs"); - return Ok(()); - } + let request_span = match xray_trace_id { + Some(trace_id) => tracing::span!( + tracing::Level::TRACE, + "Lambda request", + requestId = request_id, + xrayTraceId = trace_id + ), + None => tracing::span!(tracing::Level::TRACE, "Lambda request", requestId = request_id), }; - let req = match handler.ready().await { - Ok(handler) => { - // Catches panics outside of a `Future` - let task = - panic::catch_unwind(panic::AssertUnwindSafe(|| handler.call(LambdaEvent::new(body, ctx)))); - - let task = match task { - // Catches panics inside of the `Future` - Ok(task) => panic::AssertUnwindSafe(task).catch_unwind().await, - Err(err) => Err(err), - }; - - match task { - Ok(response) => match response { - Ok(response) => { - trace!("Ok response from handler (run loop)"); - EventCompletionRequest { - request_id, - body: response, + async { + // Group the handling in one future and instrument it with the span + let body = hyper::body::to_bytes(body).await?; + trace!("response body - {}", std::str::from_utf8(&body)?); + + #[cfg(debug_assertions)] + if parts.status.is_server_error() { + error!("Lambda Runtime server returned an unexpected error"); + return Err(parts.status.to_string().into()); + } + + let body = match serde_json::from_slice(&body) { + Ok(body) => body, + Err(err) => { + let req = build_event_error_request(request_id, err)?; + client.call(req).await.expect("Unable to send response to Runtime APIs"); + return Ok(()); + } + }; + + let req = match handler.ready().await { + Ok(handler) => { + // Catches panics outside of a `Future` + let task = + panic::catch_unwind(panic::AssertUnwindSafe(|| handler.call(LambdaEvent::new(body, ctx)))); + + let task = match task { + // Catches panics inside of the `Future` + Ok(task) => panic::AssertUnwindSafe(task).catch_unwind().await, + Err(err) => Err(err), + }; + + match task { + Ok(response) => match response { + Ok(response) => { + trace!("Ok response from handler (run loop)"); + EventCompletionRequest { + request_id, + body: response, + } + .into_req() } - .into_req() + Err(err) => build_event_error_request(request_id, err), + }, + Err(err) => { + error!("{:?}", err); + let error_type = type_name_of_val(&err); + let msg = if let Some(msg) = err.downcast_ref::<&str>() { + format!("Lambda panicked: {}", msg) + } else { + "Lambda panicked".to_string() + }; + EventErrorRequest::new(request_id, error_type, &msg).into_req() } - Err(err) => build_event_error_request(request_id, err), - }, - Err(err) => { - error!("{:?}", err); - let error_type = type_name_of_val(&err); - let msg = if let Some(msg) = err.downcast_ref::<&str>() { - format!("Lambda panicked: {}", msg) - } else { - "Lambda panicked".to_string() - }; - EventErrorRequest::new(request_id, error_type, &msg).into_req() } } - } - Err(err) => build_event_error_request(request_id, err), - }?; + Err(err) => build_event_error_request(request_id, err), + }?; - client.call(req).await.expect("Unable to send response to Runtime APIs"); + client.call(req).await.expect("Unable to send response to Runtime APIs"); + Ok::<(), Error>(()) + } + .instrument(request_span) + .await?; } Ok(()) } From 98ba6d7aa6b9989a8a38b0601fd8e780c608e961 Mon Sep 17 00:00:00 2001 From: llgerard <87908777+llgerard@users.noreply.github.com> Date: Thu, 15 Dec 2022 15:18:29 -0800 Subject: [PATCH 2/3] Bump span logging to INFO Tracing filters by default trim events and *spans* according to the level. This means that entering a span with a level below the threshold will be a no-op and the current span is not updated. The lambda request span is a meaningful info one, so bump to INFO level. --- lambda-runtime/src/lib.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/lambda-runtime/src/lib.rs b/lambda-runtime/src/lib.rs index 6bcaf22b..e6228d0c 100644 --- a/lambda-runtime/src/lib.rs +++ b/lambda-runtime/src/lib.rs @@ -131,16 +131,16 @@ where } let request_span = match xray_trace_id { Some(trace_id) => tracing::span!( - tracing::Level::TRACE, + tracing::Level::INFO, "Lambda request", requestId = request_id, xrayTraceId = trace_id ), - None => tracing::span!(tracing::Level::TRACE, "Lambda request", requestId = request_id), + None => tracing::span!(tracing::Level::INFO, "Lambda request", requestId = request_id), }; + // Group the handling in one future and instrument it with the span async { - // Group the handling in one future and instrument it with the span let body = hyper::body::to_bytes(body).await?; trace!("response body - {}", std::str::from_utf8(&body)?); From 86d47ba1bd0472fabe5d5acd220b2f83f0c47f84 Mon Sep 17 00:00:00 2001 From: Nina Wang Date: Fri, 16 Dec 2022 17:45:25 +0000 Subject: [PATCH 3/3] fix formatting issue --- lambda-runtime/src/lib.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lambda-runtime/src/lib.rs b/lambda-runtime/src/lib.rs index e6228d0c..cad99de4 100644 --- a/lambda-runtime/src/lib.rs +++ b/lambda-runtime/src/lib.rs @@ -26,7 +26,7 @@ use tokio::io::{AsyncRead, AsyncWrite}; use tokio_stream::{Stream, StreamExt}; pub use tower::{self, service_fn, Service}; use tower::{util::ServiceFn, ServiceExt}; -use tracing::{error, span, trace, Instrument}; +use tracing::{error, trace, Instrument}; mod requests; #[cfg(test)]