X-Ray tracing for AWS Rust SDK in Lambda runtime

Rust becomes more and more widely used in many places and AWS cloud is no exception. With AWS SDK for Rust (currently in preview) and Rust runtime for Lambda it's possible to build regular Lambda functions. Rust gives those serverless components safety and runtime speed but it's still quite new and ecosystem around Rust for AWS Lambda does not cover all features from other SDKs. One of such missing features is good X-Ray integration. There is no fully-flagged instrumentation for Rust SDK and if you search for solutions how to include such tracing in Rust projects using SDK… you will find planty of different solutions and there is high chance that none of them is what you are looking for.

Current state

One solution that you may find (in different variants) relies on OpenTelemetry implementation for Rust. But it's not really feasible in Lambda runtime for two reasons:

• to send X-Ray segments into same trace you need to include trace ID that is exposed by Lambda runtime, but at leasts from my investigation OpenTelemetry implementation does not handle it (probably it's possible to somehow bind it, but even then…);
• in Lambda runtime environment there is an X-Ray UDP daemon running, not OpenTelemetry TCP one.

So most of the solutions I've found on the internet were fine for containerized workflows or other environments, like running apps on EC2 or even outside of AWS. What I was missing is to have X-Ray tracing in my Lambda-based workflows running on native runtime.

X-Ray UDP client

That's right - in Lambda environment, by default the UDP daemon is running. This case, unfortunately, is not supported by OpenTelemetry integration. In simple words - there is no X-Ray integration/instrumentation for Rust SDK in native Lambda runtime. There are some crates that provide X-Ray UDP client - mainly originating from https://github.com/softprops/xray I think. There is however one issue still to solve: tracing spans from AWS SDK for Rust are in no way compatible with segments model in X-Ray and additionally no integration between these crates exists. One could just generate extra code to send trace segments via UDP every time SDK call is used in the code but this also isn't simple - outside of SDK we may have no access to HTTP-level response to record error state.

X-Ray subscriber for tracing

This is the piece I have done in my fork of the mentioned client - I created tracing subscriber that translates SDK spans into X-Ray segments. Since Rust can use code directly from GitHub you can include it in your project with:

xray = { git = "https://github.com/rafalwrzeszcz/xray", rev = "ca84d3228d917570ff1e1ce32ae7c7323358ed53" }

Disclaimer though! This is my very first project in Rust, made for my own purpose. It may lack some of the use-cases or even not work and break your program. Ok - you have been warned, go on.

It's completely transparent! Just bind X-Ray subscriber as tracing global dispatcher and it will work on all AWS SDK client (if you want you can bind it more granularily, but this is the shortest way):

use tracing_core::dispatcher::set_global_default;
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::Registry;
use xray::XRaySubscriber;

set_global_default(
    Registry::default()
         .with(XRaySubscriber::default())
         .into()
)?;

let response = client
    .get_item()
    .table_name(table_name.as_str())
    .key("customerId", S(customer_id.to_string()))
    .send()
    .await?;

AWS metadata

One more catch. In some cases X-Ray needs more meta-data - for example, for DynamoDB calls it needs to know table name, otherwise all calls will be traced just on service-level (you won't see individual tables in traces). Current tracing data from SDK does not expose this data. To handle that, listener implementation allows special span named "aws_metadata" - you can instrument your operation with this span and pass two extra properties - region and table_name. Every SDK call that happens within such span will be enriched by specified values:

use xray::aws_metadata;

let response = client
    .get_item()
    .table_name(table_name.as_str())
    .key("customerId", S(customer_id.to_string()))
    .send()
    .instrument(
        aws_metadata(
            None, // region is not helpful - we are in same region
            Some(table_name.as_str())
        )
    )
    .await?;

X-Amzn-Trace-Id

Ok, so our Lambda emits now X-Ray (sub)segments to UDP side-car daemon. What about X-Amzn-Trace-Id header that is mentioned in every page you found in Google when looking for solution?

Firstly, it's handled by SDK - in recursion_detection middleware (part of default middleware stack).

And secondly, it's not the full story - this header allows propagation of tracking traces and allows active services to pick it up and send segments to the same trace. Some services do not actually send traces and in such cases X-Ray generates so-called inferred segments - but for that active (client in this case) side needs to send segment data. And this is the missing piece solved by UDP client integration.

Result

This is how it looks in AWS X-Ray console: