Last Updated on December 8, 2023

You can have task local storage in asyncio programs using context variables in the contextvars module.

This provides thread-local-like storage for arbitrary contexts within one thread, such as within tasks in an asyncio program.

In this tutorial, you will discover how to use context variables for private and local state within each asyncio task.

Let’s get started.

Need Task-Specific State in Asyncio

It is common to execute the same tasks in our program with different data.

We can pass the data down to each task, then sub-tasks and sub-sub-tasks as arguments. This is fine for task-specific data.

But what about data that is unique to the task but not required to execute the task? Data that we might use to uniquely identify the task, such as during logging.

Data such as:

• Task id
• Session id
• Request id
• Error id
• etc.

This data is generally cumbersome to pass down the call graph using arguments or up the call graph using return values or exception messages.

One approach is to store this data in global variables and have each task access this shared data.

Generally, global variables are discouraged, and there is an opportunity for tasks and their child tasks to access the wrong state by mistake.

A common solution when using thread workers to execute tasks is to use thread-local storage. This is a private context available only to the thread. Each task accesses the same variable name, but the data stored and retrieved is unique and private.

Thread-local data is data whose values are thread specific.

— threading — Thread-based parallelism

You can learn more about thread local storage in the tutorial:

• Thread-Local Data in Python

What is needed is a version of thread local storage for asyncio tasks and coroutines.

This would allow each task and subtasks to access the same variable name, although the data accessed would be unique and private for the task and subtasks scheduled by that task.

How can we have thread-local storage in asyncio for tasks coroutines?

How can we have task-specific local storage?

Use Context Variables For Shared Local State

We can can maintain task-specific state using context variables.

Context variables or contextvars are a reasonably new API in the Python standard library.

They provide a way to have a shared local state that is private within a thread, such as within one given call graph.

This module provides APIs to manage, store, and access context-local state.

— contextvars — Context Variables

This capability was described in PEP 567, specifically designed as a way of offering thread-local-like behavior for tasks within an asyncio program. The solution allows thread-local behavior generally within a thread and private to each call graph.

This is opposed to thread-local memory which is a capability provided by the operating system native threads and offered within Python to have variables across threads and private to each thread.

Thread-local variables are insufficient for asynchronous tasks that execute concurrently in the same OS thread. Any context manager that saves and restores a context value using threading.local() will have its context values bleed to other code unexpectedly when used in async/await code.

— PEP 567 – Context Variables

The contextvars module was introduced in Python in version 3.7.

The new contextvars module and a set of new C APIs introduce support for context variables. Context variables are conceptually similar to thread-local variables. Unlike TLS, context variables support asynchronous code correctly.

— What’s New In Python 3.7

This approach to shared local state is now generally recommended in Python instead of using thread local, providing a more flexible approach that can operate at the preferred scope.

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.

— contextvars — Context Variables

It can be used for arbitrary call graphs within a Python program, and works with asyncio tasks and coroutines out of the box, without any additional configuration

Context variables are natively supported in asyncio and are ready to be used without any extra configuration.

— contextvars — Context Variables

Now that we know what contextvars are, let’s look at how we might use them in asyncio.

How to Use Context Variables in Asyncio

A context variable should be first created at the top level of the program, e.g. at the module level as a global variable.

Important: Context Variables should be created at the top module level and never in closures. Context objects hold strong references to context variables which prevents context variables from being properly garbage collected.

— contextvars — Context Variables

A context variable is created via the contextvars.ContextVar() object and a unique name for the variable are provided as an argument.

For example:

Once created, a value for the variable can be set via the set() method.

This will be private to the asyncio coroutine and any and all coroutines and tasks created by that task, as well as their descendants.

For example:

A value from a contextvar can be retrieved via the get() method.

The value will be unique and private for the caller’s context, e.g. call graph or block.

For example:

Importantly, the scope of the variable is the caller’s context.

I think the easiest way to think about this is in terms of call graphs, but you can also think of blocks or similar organizations of code.

For example, tasks that create tasks that create tasks. All base-level tasks and their children may share state using a context variable.

For example:

• Base Task: state.set(1)
  • Child Task: state.get() # 1
    • Grand Child Task: state.get() # 1
• Base Task: state.set(2)
  • Child Task: state.get() # 2
    • Grand Child Task: state.get() # 2
• Base Task: state.set(3)
  • Child Task: state.get() # 3
    • Grand Child Task: state.get() # 3
• …

All tasks access the same “variable” but “see” different values based on their “context” e.g. call graph.

To make this clearer, let’s consider some use cases:

A common use case for context variables in an asyncio program would be:

1. Create a context variable at the module level
2. Assign values (e.g. at the root level in a call graph)
3. Access values (e.g. at the leaf level of the call graph)

That being said, values can be read/written at any level.

For example, these are common patterns:

1. Write state at the root level, read at leaf level of call graph.
2. Write state at the leaf-level, read at the root level of the call graph.

Now that we know how to use context variables in asyncio, let’s look at some worked examples.

Example of Context Variables in Asyncio

We can explore an example of sharing private state from the root level of the task with all descendant tasks.

Each separate root-level task will have a different task identifier which might be used for logging, reporting, error handling, and so on.

In this case, the identifier is generated as a random number and stored by each task when it is started. The task then retrieves and reports this id and starts a sub-task. The sub-task retrieves and reports the id and starts its own sub-task which also reports the id.

Multiple root-level tasks are then created, each with its own private version of the id, although accessed via the shared global context variable.

Firstly, we can create the grandchild task that takes a name for the task, retrieves the shared local id, and reports the value.

Next, we can define the child task that takes the task name, retrieves and reports the id, and starts and awaits the grandchild task.

Next, we can define the root-level task.

The task first generates a unique task id using random.random().

This is first used to sleep the task, to ensure that each task and subtasks that are started store and accesses its context variable at different times. This is intentional, to demonstrate that reads and writes of the context variable are private to each call graph.

The task then stores the id, retrieves it, and starts a child task.

Finally, the main() coroutine starts 3 root level tasks using an asyncio.TaskGroup via the context manager interface and waits for them to complete.

If you are new to the asyncio.TaskGroup, a preferred way to manage groups of tasks. You can learn more about it in the tutorial:

• How to use asyncio.TaskGroup

Finally, before starting the event loop we can declare and define our shared context variable.

Tying this together, the complete example is listed below.

Running the example first declares and defines a context variable named “id” at the module level.

This will be accessible throughout the program as a global variable, except the context under which it is accessed will define what values are written and read.

Next, the asyncio event loop is started and the main() coroutine is run. The TaskGroup is created and 3 tasks named ‘A’, ‘B’, and ‘C’ are created and scheduled. The main() coroutine then blocks until all 3 tasks are done.

Each task runs, first generating a random (and likely unique) id and sleeping for some fraction of 5 seconds.

Each task then resumes in time and stores its id in the shared context global variable. Each task stores its own private version of value in the context variable.

The variable is not overwritten by each task.

The tasks retrieve their value and report it before starting a child task and awaiting it.

Child tasks run, retrieve their id, and report it before starting and awaiting a grandchild task. The process repeats

Reviewing the output we can see that Tasks in the “Task A” call graph all report the same id from the shared context variable.

If the shared_ids global variable was a regular global variable, this would not be the case, as each task() would overwrite the same value. Instead, because it is a context variable, each context, in this case, “Task A” has its own private version of the variable.

We see the same with “Task B” which maintains its own unique id, and “Task C”.

This highlights how we can use context variables to maintain private local state for tasks and their descendants in asyncio programs.

Further Reading

This section provides additional resources that you may find helpful.

Python Asyncio Books

• Python Asyncio Mastery, Jason Brownlee (my book!)
• Python Asyncio Jump-Start, Jason Brownlee.
• Python Asyncio Interview Questions, Jason Brownlee.
• Asyncio Module API Cheat Sheet

I also recommend the following books:

• Python Concurrency with asyncio, Matthew Fowler, 2022.
• Using Asyncio in Python, Caleb Hattingh, 2020.
• asyncio Recipes, Mohamed Mustapha Tahrioui, 2019.

Guides

• Python Asyncio: The Complete Guide
• Python Asynchronous Programming

APIs

• asyncio — Asynchronous I/O
• Asyncio Coroutines and Tasks
• Asyncio Streams
• Asyncio Subprocesses
• Asyncio Queues
• Asyncio Synchronization Primitives

References

• Asynchronous I/O, Wikipedia.
• Coroutine, Wikipedia.

Takeaways

You now know how to use context variables for private and local state within each asyncio task.

Did I make a mistake? See a typo?
I’m a simple humble human. Correct me, please!

Do you have any additional tips?
I’d love to hear about them!

Do you have any questions?
Ask your questions in the comments below and I will do my best to answer.

Photo by Itadaki on Unsplash