search

tokio-rs / tokio

A runtime for writing reliable asynchronous applications with Rust. Provides I/O, networking, scheduling, timers, ...
https://tokio.rs
MIT License
26.99k stars 2.49k forks source link

Non-Scoped Task-Local Storage #5037

Open azdle opened 2 years ago

azdle commented 2 years ago

I'm working on a tracing layer that we're trying to integrate with tokio and I would like to be able to store metadata about each task somehow on each task itself.

We use thread-local storage for the same metadata when tracing os-thread based programs already, but I believe I'm unable to use the current scope-style task-local storage unless I get users of my library to wrap any task futures with a task-local scope. I would like to avoid this if at all possible because it would be both a stumbling block and I feel it would be likely to be very error prone.

Proposed Solution

I would like to see the current task-local storage grow what is essentially a default task-scoped task-local storage that would look, from the outside, essentially how @LucioFranco describes the task-local storage of futures 0.1 in the original commit for the current implementation: 619d730d610bbfd0a13285178d6bf3d89a29d6a3

Said another way, I would like to see task-local storage also work like std's thread-local storage with dynamic initialization of a default value, but only when there isn't an active scope with a value.

Context

The basic metadata I'm hoping to store on the task is a unique ID for the task and what amounts to a boolean that tracks whether I've seen this task before (or really if I've declared this task in my trace data).

For the unique ID I have the workaround of hashing the task_id from #4721 into something, but would still like to cache that something.

For the boolean state, I can track that locally in a map by the unique ID, but I believe I have no way to tell when a task has finished to be able to stop tracking that task and not grow my map forever.

Alternatives

I could require all tasks in a user's program to be wrapped in a thread-local scope. Even with something like a custom attribute macro I find this unacceptable because it would be very intrusive to user code and would require users to always remember to wrap any future that will become a task.

I could require all tasks be wrapped in a span. This has the same set of problems as above, but is probably overall less objectionable since it's something people might be expected to do anyway to trace custom values.

A new tokio API to be notified about the start and end of tasks could be added to let me accurately track active IDs. This seems weird and like something you wouldn't want to add, but it would technically solve my immediate problem.

Implementation Musing

After doing a bit of poking, I'm assuming this could look a lot like #4721, but with a pointer to a linked list of values instead of the Id, perhaps even consuming the task_id as a default task-local value.

Thoughts?

LucioFranco commented 2 years ago

@azdle hi sorry for the delay! I would imagine you might want something like a LazyTaskLocal where when you init it in the macro you pass some closure that defines the default state of that task local. This could allow you to set a random id for each task and then have that id available via the task local value.