Don't hold span.enter() across an .await: a wrong-trace-id bug in async Rust

#backend #debugging #rust

While investigating our logs I found a classic async Rust trap where logs contained both resuls from polling by internal service and also messages as result of incoming API calls to our system. Logs from one request could carry the trace id of a different request. The cause was one line in a axum middleware.

TLDR

Do not hold the guard from span.enter() across an .await.
Use .instrument(span) on the future instead. It enters the span on every poll and leaves it on every yield.
Test it with two requests that interleave on one thread, and check that each log line sits in exactly one span.

The code that looked fine

An axum middleware that gives every request a trace id:

pub async fn trace_id_middleware(req: Request<Body>, next: Next) -> Response {
    let trace_id = Uuid::new_v4().to_string();
    let span = info_span!("api_request", trace_id = %trace_id);
    let _enter = span.enter();          // <- the problem

    let mut req = req;
    req.extensions_mut().insert(trace_id);

    next.run(req).await
}

It compiles, and with one request at a time it works. Under concurrent load the logs got mixed up.

Why it breaks

span.enter() returns a guard. While the guard is alive, the span is "current" on that thread. That is a thread-local setting, not a property of your future.

When the handler hits an .await and yields, the guard is still alive inside the suspended future, so the thread still has this request's span as current. The runtime then polls another request on the same thread. That request's logs, and any span it creates, now sit under the first request's span.

The tracing docs warn about this. It is easy to miss because nothing fails: you only see logs with the wrong context.

The fix

use tracing::Instrument;

pub async fn trace_id_middleware(req: Request<Body>, next: Next) -> Response {
    let trace_id = Uuid::new_v4().to_string();
    let span = info_span!("api_request", trace_id = %trace_id);

    let mut req = req;
    req.extensions_mut().insert(trace_id);

    next.run(req).instrument(span).await
}

.instrument(span) wraps the future. Each time the future is polled it enters the span, and when it yields it leaves it. The span is current only while this request's code is actually running.

No more:

let _guard = span.enter();

Rather:

some_future.instrument(span).await

Test it

Add a test with two requests that interleave on one thread, and check that every log line sits inside exactly one api_request span, the right one. Then put span.enter() back once and confirm the test fails. A test that passes on the buggy code proves nothing.

The same idea in a polling loop

We have loops that peek a message from a queue, store it, and then acknowledge it. Each message has a trace id minted at peek time. Passing it by hand into every log call is easy to forget, so each message now runs inside its own span:

let span = message_span("DDK", &collection, &peek_trace_id);
async {
    // ... parse, persist, acknowledge ...
    tracing::Span::current().record("message_id", msg_id.as_str());
    // ...
}
.instrument(span)
.await;

Two details:

Declare the field as tracing::field::Empty when you create the span, and fill it in later with Span::current().record(...) once the value is known.
Inside the async block, continue is not allowed. The early exits became return, which ends this message and lets the loop move on to the next tick.

Takeaways

In async code, never keep a span.enter() guard alive across an .await.
Put the span on the future with .instrument(span).
Test with interleaved tasks on a single thread, and check the full span scope of each event, not just the innermost span.
Break the code on purpose once to see the test fail.

DEV Community