Repository: awslabs/aws-lambda-rust-runtime (language: rust)

Issue:
`_X_AMZN_TRACE_ID` environment variable disappeared
I'm trying [the layering of runtime feature](https://github.com/awslabs/aws-lambda-rust-runtime/pull/845)

Below code failed with `"_X_AMZN_TRACE_ID not set"` error.  It seems the environment variable is unset [here](https://github.com/awslabs/aws-lambda-rust-runtime/blob/lambda-runtime-0.11.0/lambda-runtime/src/layers/trace.rs#L64). How can I create a span using OpenTelemetry SDK (not tracing crate) whose trace id is retrieved from `_X_AMZN_TRACE_ID` environment variable?

```rust
fn init_observability() -> Result<TracerProvider, Error> {
    let export_config = ExportConfig {
        endpoint: "http://0.0.0.0:4317".to_string(),
        timeout: std::time::Duration::from_secs(3),
        protocol: Protocol::Grpc,
    };
    let exporter = opentelemetry_otlp::new_exporter()
        .tonic()
        .with_export_config(export_config)
        .build_span_exporter()
        .expect("building exporter failed");
    let tracer_provider = TracerProvider::builder()
        .with_batch_exporter(exporter, opentelemetry_sdk::runtime::Tokio)
        .with_config(
            trace::config()
                .with_resource(Resource::new(vec![KeyValue::new("service.name", "lambda")])),
        )
        .build();
    let fmt_layer = fmt::layer().event_format(format().json());
    let filter_layer = EnvFilter::from_default_env();
    let telemetry_layer =
        tracing_opentelemetry::OpenTelemetryLayer::new(tracer_provider.tracer("my-app"));

    tracing_subscriber::registry()
        .with(filter_layer)
        .with(fmt_layer)
        .with(telemetry_layer)
        .init();

    Ok(tracer_provider)
}

fn get_span_context_from_environment_var() -> Result<SpanContext, &'static str> {
    let xray_trace_id =
        std::env::var("_X_AMZN_TRACE_ID").map_err(|_| "_X_AMZN_TRACE_ID not set")?;
    span_context_from_str(&xray_trace_id).ok_or("Invalid _X_AMZN_TRACE_ID")
}

#[tokio::main]
async fn main() -> Result<(), Error> {
    let tracer_provider = init_observability()?;
    let func = service_fn(handler);
    // Initialize the Lambda runtime and add OpenTelemetry tracing
    let runtime = Runtime::new(func).layer(OpenTelemetryLayer::new(|| {
        // Make sure that the trace is exported before the Lambda runtime is frozen
        tracer_provider.force_flush();
    }));
    runtime.run().await?;
    Ok(())
}

async fn handler(event: LambdaEvent<S3Event>) -> Result<Value, Error> {
    info!("Handler is called");
    let parent_context = match get_span_context_from_environment_var() {
        Ok(context) => context,
        Err(e) => {
            error!(
                "get_span_context_from_environment_var raised error: {:?}",
                e
            );
            return Err(e.into());
        }
    };
    // ...
    let labels = detect_labels(client, argument)
        .with_context(Context::new().with_remote_span_context(parent_context))
        .await?;
    Ok(json!({ "message": format!("Labels is {:?}!", labels) }))
}
```


Discussion / hints:
It's explained here: https://docs.rs/lambda_runtime/latest/lambda_runtime/struct.Runtime.html#method.new

> Note that manually creating a [Runtime](https://docs.rs/lambda_runtime/latest/lambda_runtime/struct.Runtime.html) does not add tracing to the executed handler as is done by [super::run](https://docs.rs/lambda_runtime/latest/lambda_runtime/fn.run.html). If you want to add the default tracing functionality, call [Runtime::layer](https://docs.rs/lambda_runtime/latest/lambda_runtime/struct.Runtime.html#method.layer) with a [super::layers::TracingLayer](https://docs.rs/lambda_runtime/latest/lambda_runtime/layers/struct.TracingLayer.html).

Maybe we need to move that variable somewhere else that doesn't depend in the Tracing layer.
You can also take the information from the information context in your own layer. See how that's extracted in these two lines of code: 

https://github.com/awslabs/aws-lambda-rust-runtime/blob/lambda-runtime-0.11.0/lambda-runtime/src/layers/trace.rs#L46
https://github.com/awslabs/aws-lambda-rust-runtime/blob/lambda-runtime-0.11.0/lambda-runtime/src/layers/trace.rs#L54

If you want to get the trace id in your handler, you can use `event.context.xray_trace_id`: https://docs.rs/lambda_runtime/latest/lambda_runtime/struct.Context.html
I'm going to change the layers to make that env variable always present. I'll probably release the change this weekend.

Fix patch:
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                                             */
 /* --------------------------------------------------------------------------------------------- */


Test patch: