Is there any way to access 'request-local' state from a macro, e.g. with tokio::task_local!

[thomasmost]

I am trying to write a macro to replace println! for logging in my Rocket application. This macro would ideally prepend a request id no matter where it was called from. I know I can pass a RequestId guard into any route handler, but I don't want to have to pipe it around through every function that logs... Any advice on how to do this?

So far I have something like this:

use std::sync::atomic::{AtomicU64, Ordering};

pub static REQUEST_ID: AtomicU64 = AtomicU64::new(0);

// This is set in the RequestGuard,but gets overwritten by concurrent requests
pub fn set(val: u64) {
  REQUEST_ID.store(val, Ordering::Relaxed)
}

pub fn get() -> u64 {
  REQUEST_ID.load(Ordering::Relaxed)
}

macro_rules! printlog {
    ($($arg:tt)*) => {{
      let request_id = crate::printlog::get();
      std::println!("RequestId={:?}: {}", request_id, std::format_args!($($arg)*));
    }};
  }

...but obviously with a server handling concurrent requests, this will only ever log the most recently initialized request ID.

There seems to be something promising about using tokio::task_local! —but I can't figure out where to scope the "task-local storage value" around Rocket's request handler. E.g.

rocket::tokio::task_local! {
  static TOKIO_REQUEST_ID: RefCell<Option<String>>;
}

// This causes a panic with the message: 'cannot access a task-local storage value without setting it first'
pub fn log (line: String) -> () {
  TOKIO_REQUEST_ID.with(|request_id_ref| {
    let request_id = request_id_ref.borrow_mut();
    let thread_id = std::thread::current().id();
    std::println!("RequestId={:?}; {:?}: {}", request_id, thread_id, line);
});
}

macro_rules! printlog {
  ($($arg:tt)*) => {{
    crate::printlog::log(std::format_args!($($arg)*).to_string());
  }};
}

Seems like this may relate to #2160

[edgarogh]

The clean way to do things (and likely the fastest) would be to pass around your RequestId as you mentioned in your first paragraph. That's how Rust works and should work.

If that's really a problem to you, you can use a task_local! slot, but as far as I know, there is no simple way to define it for a request handler.

---

Normally, you'd have to call LocalKey::scope and the future you give it will have the thread-local value set inside it. But Rocket fairings don't let you call the handler yourself, so you wouldn't be able to make a fairing that sets your request id. The only solution I can think of is to create the scope yourself in each and every handler:

tokio::task_local! {
  static REQUEST_ID: RequestId;
}

#[get("/")]
async fn index(rid: RequestId) {
  REQUEST_ID.scope(rid, async move {
    // ...
  })
}

#[get("/mypage")]
async fn my_page(rid: RequestId) {
  REQUEST_ID.scope(rid, async move {
    // ...
  })
}

If that becomes to ugly or hard to maintain, you could also make your own #[with_rid] procedural macro to transform any request handler and add the boilerplate LocalKey::scope call as well as the guard/parameter.

[thomasmost]

The clean way to do things (and likely the fastest) would be to pass around your RequestId as you mentioned in your first paragraph. That's how Rust works and should work.

The implication that figuring out how to implement 'context' for this is not "how Rust works" (or should work) is strange to me. A fixed-length request ID is a perfectly safe data structure to store in scoped memory, and the developer experience benefits of automatically getting the request ID logged seem to me to more than justify the effort (if it's doable).

Yes, as you observe, implementing this for each request handler is not ideal, though it's still perhaps better than piping an argument through every helper function. I'll look into the macro suggestion, which is again better but not ideal. I confess I'm tempted to fork the project to get around this...

But Rocket fairings don't let you call the handler yourself

...but I recognize that's probably not worth the maintenance headache I'm inviting down the line.

Applying logic to or running wrapper code around all requests in a given mount method call seems like a common enough use case I'm a bit surprised Rocket doesn't support it. I imagine @SergioBenitez has philosophical reasons for excluding it, much like he does about middleware (and I largely agree with him on that point) but still, having NO way of repeating functionality across all endpoints in a set automatically feels like a bit of a gap in the architecture...

[the10thWiz]

This is somewhat more complicated than you may think. Fairings can run before and after the request handler, which is more than enough for most use cases. However, in this case, we need the ability to wrap the code in a scope, which isn't common for literally anything else.

For a more traditional ThreadLocal value, fairings are good enough. All it would require is a call to update the thread local value before the handler, and a call to delete it afterwards. As you pointed out, Rocket is async, so ThreadLocal isn't good enough. Tokio task local values don't provide an API compatible with Rocket's model. As a side note, async-std does provide a task local with the required API, but we can't use it since we're on tokio.

To be honest, I'm not sure how this should be done. At first glance, it might be nice to support something like this, but Rocket would likely need to make some major changes to make this work (unless Rocket took control of the RequestID, and simply offered a function to get it. Implementing this is also complicated by the aspirational hope that Rocket can be runtime neutral, but that's not the highest priority.

The implication that figuring out how to implement 'context' for this is not "how Rust works" (or should work) is strange to me.

This is actually an interesting point. Rust usually doesn't make it easy to grab any kind of context (even a backtrace), and Rust typically likes to avoid implicitly passing around any kind of context (since doing so would require a more complex runtime). This is easy to see in the design of Futures in Rust, the context needed to execute futures is explicitly passed to the future when it's polled.

Overall, I'm not sure where I lie on this issue. To be honest, I think the solution might be for Rocket to provide a RequestId task local. This would solve this particular issue, and avoids trying to come up with a generic API for it. Rocket would then be able to take advantage of this by default, so at least any logs generated by Rocket would be marked with the request id.

[thomasmost]

I agree it's complicated—but I also believe it's incredibly useful, arguably to a universal degree. For the purpose of debugging production applications, what developer doesn't want a unique request ID logged automatically to trace the progression of logs executing a specific request across async threads?

Because of this, I think I agree: this should be in Rocket core.

[mgumbley-resilient]

I looked briefly at using https://github.com/kmdreko/task-local but it has been abandoned in favour of tokio's tokio::task::LocalKey. (which I couldn't see how to make work) Might this approach work?

[thomasmost]

FWIW, I migrated my project to Axum — its middleware layers allow configuring a tracing span with the request ID

[SergioBenitez]

Sorry to hear we lost you. We're planning on fully integrating tracing soon, which will enable this and other use cases.

[thomasmost]

Hey I'm still a fan of Rocket! Just needed something a little less on rails for a production-ready enterprise application. Will still be watching the project closely 🙂 Thank you for all your work!