Replace worker use of `threading.local()` with ContextVar

Context variables are now the recommended way to manage thread/context-local state:

Context managers that have state should use Context Variables instead of threading.local() to prevent their state from bleeding to other code unexpectedly, when used in concurrent code. https://docs.python.org/3/library/contextvars.html

They work for threads, but as a bonus, also work for async code.

In many places in the worker (and even more in distributed in general), we use threading.local() variables currently. However, workers also support being used in an async mode (used in most tests), and increasingly support running async functions (https://github.com/dask/distributed/pull/5151). The fact that thread-local variables can have the wrong value when used in an async context means there are likely lots of bugs around async mode. https://github.com/dask/distributed/issues/4959 is just one example.

If we want to support async mode properly, switching Worker use of thread-local variables to ContextVars would be a good place to start. These sorts of bugs are generally very confusing and tedious to track down. It might be worth getting ahead of these bugs instead of repeatedly hunting them down.

As a bonus, we’re inconsistent about ContextVars vs thread-local currently (for example, get_client checks a ContextVar, but get_worker checks a thread-local var). Consolidating on one interface might make future development smoother.

cc @jcrist @crusaderky @fjetter

[crusaderky]

+1 for replacing thread-local variables with contetxvars wherever possible.

[gjoseph92]

I’m surprised to hear that. I thought I found the opposite was the case in https://github.com/dask/distributed/pull/5517.