From 3c4941a88c906dd9622db320a591154c71a2ec54 Mon Sep 17 00:00:00 2001
From: David Calavera <david.calavera@gmail.com>
Date: Fri, 29 Mar 2024 17:39:49 -0700
Subject: [PATCH] Always export the _X_AMZN_TRACE_ID env variable.

This variable is expected across runtimes regardless of other features.

Signed-off-by: David Calavera <david.calavera@gmail.com>
---
 lambda-runtime/src/layers/trace.rs |  3 ---
 lambda-runtime/src/runtime.rs      | 11 +++++++++++
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/lambda-runtime/src/layers/trace.rs b/lambda-runtime/src/layers/trace.rs
index 0d635154..7a9f8370 100644
--- a/lambda-runtime/src/layers/trace.rs
+++ b/lambda-runtime/src/layers/trace.rs
@@ -1,4 +1,3 @@
-use std::env;
 use tower::{Layer, Service};
 use tracing::{instrument::Instrumented, Instrument};
 
@@ -53,7 +52,6 @@ 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,
@@ -61,7 +59,6 @@ fn request_span(ctx: &Context) -> tracing::Span {
             )
         }
         None => {
-            env::remove_var("_X_AMZN_TRACE_ID");
             tracing::info_span!("Lambda runtime invoke", requestId = &ctx.request_id)
         }
     }
diff --git a/lambda-runtime/src/runtime.rs b/lambda-runtime/src/runtime.rs
index 0fc328cf..5ded610c 100644
--- a/lambda-runtime/src/runtime.rs
+++ b/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                                             */
 /* --------------------------------------------------------------------------------------------- */