Trace Amounts of Rust

I feel like I start all these posts with “recently at work…”, but the truth is, I learn a lot at work, and not documenting my work for my future self to find it seems criminal. It’s inevitable the next time I have these same questions, and if I can do 35 year old me a solid (hopefully Google indexes my pages by then), why shouldn’t I?

Anyway, recently at work, I was trying to improve the observability of one of our services. It’s a pretty high-QPS service that receives a lot of messages over the wire. Every once in a while, one of these messages is slow to process, and I wanted to set up traces that we could look at when this happened.

The canonical way to set up tracing in Rust is to use the aptly named tracing package. This gives you a simple way to instrument functions:

#[tracing::instrument(name="function_a")]
fn a(val: i32) {
    println!("a executed {}", val);
}

The instrument call will automagically generate a trace that tracks the name, the arguments of the function and some other metadata. Now remember, we have a websocket based protocol that process a lot of messages, and only some of these are slow. We don’t want to record every trace because we don’t like setting money on fire (aka giving it to Datadog), we only want to record these slow ones.

A future (present?) version of me would tell you that this was easy, you just slap Opentelemetry Collector in between your service and Datadog, set up tail-sampling and go home to the expensive mattress you bought with the money you didn’t spend on traces. But let’s pretend we aren’t that wise, we’re still learning about this whole space and do not realize that the next two weeks of work we do was in service of nothing. Sigh..

Ok - so we want to tail sample in the application layer! Let’s set up some boilerplate. This post assumes some prior knowledge of tracing and is meant as a showcase of how to do something more complicated than what the docs explicitly explain. Bryan had a fantastic post about tracing that helped my get all the tracing abstractions right in my head.If you want a starter intro, I’d recommend going over there, and coming back here after.

If you want to tail sample in a distributed system and not a single application (like we’re doing here) you still have to use some kind of centralized collector to implement this logic.

Now that we have simple traces, let’s immediately remove it in favor of our own Layer. The output we’re aiming for here is to only print traces for a when it’s slow. Before we do that, there are two important things to know - tracing made a choice to make you store data you need for it’s spans, and the main Layer functionality we care about here is the span related stuff. We’re not going to implement much of the logging, or event related stuff here.

Ok - that worked! So what did we do here? We used a CustomLayer to record the StartTime when a Span started, and if it took longer than a certain amount of time, we log it. We’re done here! We can all go home! Except…

Now the only the trace for b shows up, but we’d really like all the child spans to show up as well in a case like this. If we’re looking into why b is slow, we really need to know that a ran 100 times, otherwise just looking at the trace of b would be fairly useless. Ok, how do we do that? The child spans close before the parent, so we need some way to store the parent-child information, and emit all the spans once we make the determination at the root level that a trace has taken too long.

A root level span is a span that has no parent.

Let’s do the simplest thing possible. We’ll keep around all the data we need to log for a child span until the parent closes. The obvious tradeoff here is memory: before we could emit all these spans immediately, and never had to store any data about them after they exited. Now, all the data is kept around until the parent span itself closes, at which point we still throw away most of these traces.

And there we have it!! Of course, a real tracing solution would make sure to store all the data in the child spans, along with their respective parent IDs so we can build nice visualizations, but these are the essential building blocks for how we might implement tail-sampling in the application layer with tracing. Happy debugging!