Why does getting a contextvar here not work?

Question

So here's the basic code (sorry it's long)

import argparse
import asyncio
from contextvars import ContextVar
import sys

# This thing is the offender
message_var = ContextVar("message")


class ServerProtocol(asyncio.Protocol):
    def connection_made(self, transport):
        peername = transport.get_extra_info("peername")
        print("Server: Connection from {}".format(peername))
        self.transport = transport

    def data_received(self, data):
        message = data.decode()
        print("Server: Data received: {!r}".format(message))

        print("Server: Send: {!r}".format(message))
        self.transport.write(data)

        print("Server: Close the client socket")
        self.transport.close()


class ClientProtocol(asyncio.Protocol):
    def __init__(self, on_conn_lost):
        self.on_conn_lost = on_conn_lost
        self.transport = None
        self.is_connected: bool = False

    def connection_made(self, transport):
        self.transport = transport
        self.is_connected = True

    def data_received(self, data):
        # reading back supposed contextvar
        message = message_var.get()

        print(f"{message} : {data.decode()}")

    def connection_lost(self, exc):
        print("The server closed the connection")
        self.is_connected = False
        self.on_conn_lost.set_result(True)

    def send(self, message: str):
        # Setting context var
        message_var.set(message)
        if self.transport:
            self.transport.write(message.encode())

    def close(self):
        self.transport.close()
        self.is_connected = False
        if not self.on_conn_lost.done():
            self.on_conn_lost.set_result(True)


async def get_input(client: ClientProtocol):
    loop = asyncio.get_running_loop()
    while client.is_connected:
        message = await loop.run_in_executor(None, input, ">>>")
        if message == "q":
            client.close()
            return
        client.send(message)


async def main(args):
    host = "127.0.0.1"
    port = 5001

    loop = asyncio.get_running_loop()

    if args.server:
        server = await loop.create_server(lambda: ServerProtocol(), host, port)
        async with server:
            await server.serve_forever()
        return

    on_conn_lost = loop.create_future()
    client = ClientProtocol(on_conn_lost)
    await loop.create_connection(lambda: client, host, port)
    await get_input(client)


if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--server", "-s", default=False, action="store_true", help="Start server"
    )
    arguments = parser.parse_args(sys.argv[1:])
    asyncio.run(main(args=arguments))

This crashes with the following exception:

Exception in callback _ProactorReadPipeTransport._loop_reading(<_OverlappedF...shed result=4>)
handle: <Handle _ProactorReadPipeTransport._loop_reading(<_OverlappedF...shed result=4>)>
Traceback (most recent call last):
  File "C:\Users\brent\AppData\Local\Programs\Python\Python310\lib\asyncio\events.py", line 80, in _run
    self._context.run(self._callback, *self._args)
  File "C:\Users\brent\AppData\Local\Programs\Python\Python310\lib\asyncio\proactor_events.py", line 320, in _loop_reading
    self._data_received(data, length)
  File "C:\Users\brent\AppData\Local\Programs\Python\Python310\lib\asyncio\proactor_events.py", line 270, in _data_received
    self._protocol.data_received(data)
  File "E:\Development\Python\ibcs2023\_prep\experimental\asyncio_context.py", line 40, in data_received
    message = message_var.get()
LookupError: <ContextVar name='message' at 0x0000023F30A54FE0>
The server closed the connection

Why does calling message = message_var.get() cause a crash? Why can't Python find the context var? Why is data_received not in the same context as send ? How can I keep them in the same context?

I'm working on a larger project with the main branch of Textual and it uses a contextvar that loses context every time a message is received using a modified version of the code above.

Answer 1

Keeping a separated "context" for each task is exactly what contextvars are about. You could only assert that the send and data_received methods were called within the same context if you had control over the "uperlying" (as opposed to 'underlying') driver of your Protocol class - that is not the case, and both are called in different contexts. I mean, the answer to "How can I keep them in the same context?" is: you can't unless you write your own implementation of the code which makes this work inside asyncio.

There is no way you can keep track of metadata from a message, and retrieve this metadata on getting the reply, unless there is a marker on the message itself, that will survive the round-trip. That is: your networking/communication protocol itself have to spec a way to identify messages It might be as simple as a sequential integer number prefixing every string, for example - or, in this case where you simply echo the message back, it could be the message itself. Once you have that, a simple dictionary having these message IDs as keys, will work for what you seem to intend in this example.