Always export the _X_AMZN_TRACE_ID env variable. (#850)

This variable is expected across runtimes regardless of other features.

Signed-off-by: David Calavera <[email protected]>

Author: calavera

lambda-runtime/src/layers/trace.rs

@@ -1,4 +1,3 @@
-use std::env;
 use tower::{Layer, Service};
 use tracing::{instrument::Instrumented, Instrument};
 
@@ -53,15 +52,13 @@ where
 fn request_span(ctx: &Context) -> tracing::Span {
     match &ctx.xray_trace_id {
         Some(trace_id) => {
-            env::set_var("_X_AMZN_TRACE_ID", trace_id);
             tracing::info_span!(
                 "Lambda runtime invoke",
                 requestId = &ctx.request_id,
                 xrayTraceId = trace_id
             )
         }
         None => {
-            env::remove_var("_X_AMZN_TRACE_ID");
             tracing::info_span!("Lambda runtime invoke", requestId = &ctx.request_id)
         }
     }

lambda-runtime/src/runtime.rs

@@ -6,6 +6,7 @@ use http_body_util::BodyExt;
 use lambda_runtime_api_client::BoxError;
 use lambda_runtime_api_client::Client as ApiClient;
 use serde::{Deserialize, Serialize};
+use std::env;
 use std::fmt::Debug;
 use std::future::Future;
 use std::sync::Arc;
@@ -176,6 +177,9 @@ where
             let context = Context::new(invoke_request_id(&parts.headers)?, config.clone(), &parts.headers)?;
             let invocation = LambdaInvocation { parts, body, context };
 
+            // Setup Amazon's default tracing data
+            amzn_trace_env(&invocation.context);
+
             // Wait for service to be ready
             let ready = service.ready().await?;
 
@@ -232,6 +236,13 @@ fn incoming(
     }
 }
 
+fn amzn_trace_env(ctx: &Context) {
+    match &ctx.xray_trace_id {
+        Some(trace_id) => env::set_var("_X_AMZN_TRACE_ID", trace_id),
+        None => env::remove_var("_X_AMZN_TRACE_ID"),
+    }
+}
+
 /* --------------------------------------------------------------------------------------------- */
 /*                                             TESTS                                             */
 /* --------------------------------------------------------------------------------------------- */