pythonpython-3.xasynchronouspython-asynciopython-contextvars

understanding ContextVar object in Python


Here is a simple echo server that works correctly as expected - by "correctly" I mean, in the output we see that every user gets its own unique address when echoing inside listen_for_messages.

import asyncio
from contextvars import ContextVar


class Server:
    user_address = ContextVar("user_address")

    def __init__(self, host: str, port: int):
        self.host = host
        self.port = port

    async def start_server(self):
        server = await asyncio.start_server(self._on_connected, self.host, self.port)
        await server.serve_forever()

    def _on_connected(self, reader, writer):
        self.user_address.set(writer.get_extra_info("peername"))
        asyncio.create_task(self.listen_for_messages(reader))

    async def listen_for_messages(self, reader):
        while data := await reader.readline():
            print(f"Got message {data} from {self.user_address.get()}")

async def main():
    server = Server("127.0.0.1", 9000)
    await server.start_server()

asyncio.run(main())

Now look at where we set the context var object's value. My question is why it is working correctly?

To my understanding, every "Task" will have its own copy of the ContextVar. If the line:

self.user_address.set(writer.get_extra_info("peername"))

was inside listen_for_messages coroutine, I would not ask this question, since we explicitly created a new task for it. Even if _on_connected was a coroutine function, I wouldn't ask too. because:

client_connected_cb can be a plain callable or a coroutine function; if it is a coroutine function, it will be automatically scheduled as a Task.

But it's inside the _on_connected (plain) function and it's not a new task. Before creating the new listen_for_messages task, we have a parent Task of main() (which .run() function created). I guessed it shouldn't work and the function simply just overrides the value (because it's just part of a single parent Task)

I would appreciate if you tell me which part of my understanding is wrong.


Solution

  • The asyncio framework in this case calls the _on_connected callback with a new context copy - that is the expected behavior for ContextVars anyway - a callback chain started by the asyncio loop is counted as a "new task" as far as COntextvars are concerned.

    The callback itself, either for step over a co-routine, or an sync callback is executed through a call to context.run, which creates a new context - the relevant "under the hood" code is at https://github.com/python/cpython/blob/9c73a9acec095c05a178e7dff638f7d9769318f3/Lib/asyncio/events.py#L84