How to reset ContextVar in separate sync_to_async calls

[adamantike]

After filing https://code.djangoproject.com/ticket/32815, it was more evident that the issue I'm experiencing when trying to reset a ContextVar is related to asgiref. However, I'm not sure about this being a bug or an implementation issue, so any feedback is appreciated.

This issue has been raised in our efforts to add Django ASGI support to OpenTelemetry (open-telemetry/opentelemetry-python-contrib#391). OpenTelemetry uses a Django middleware class to implement both process_request and process_response methods. As both of them are synchronous, they are converted with sync_to_async and called separately (with thread_sensitive=True). The first call sets the ContextVar value, and persists the token to reset it. The second call uses the persisted token to reset the ContextVar to its previous value.

---

Tests added to tests/test_sync_contextvars.py, which reproduce the issue:

@pytest.mark.asyncio
async def test_sync_to_async_contextvar_reset():
    """
    Tests a ContextVar can be set and reset in different calls to sync_to_async.
    """
    # Define set and reset functions
    def sync_set_function():
        assert foo.get() == "bar"
        token = foo.set("baz")
        return token

    def sync_reset_function(token):
        assert foo.get() == "baz"
        foo.reset(token)
        return 42

    foo.set("bar")
    async_set_function = sync_to_async(sync_set_function)
    token = await async_set_function()
    async_reset_function = sync_to_async(sync_reset_function)
    assert await async_reset_function(token) == 42
    assert foo.get() == "bar"


def test_async_to_sync_contextvar_reset():
    """
    Tests a ContextVar can be set and reset in different calls to async_to_sync.
    """
    # Define set and reset functions
    async def async_set_function():
        await asyncio.sleep(1)
        assert foo.get() == "bar"
        token = foo.set("baz")
        return token

    async def async_reset_function(token):
        await asyncio.sleep(1)
        assert foo.get() == "baz"
        foo.reset(token)
        return 42

    foo.set("bar")
    sync_set_function = async_to_sync(async_set_function)
    token = sync_set_function()
    sync_reset_function = async_to_sync(async_reset_function)
    assert sync_reset_function(token) == 42
    assert foo.get() == "bar"

Test traceback:

tests/test_sync_contextvars.py:79:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
asgiref/sync.py:443: in __call__
    ret = await asyncio.wait_for(future, timeout=None)
/usr/lib/python3.9/asyncio/tasks.py:442: in wait_for
    return await fut
/usr/lib/python3.9/concurrent/futures/thread.py:52: in run
    result = self.fn(*self.args, **self.kwargs)
asgiref/sync.py:484: in thread_handler
    return func(*args, **kwargs)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

token = <Token var=<ContextVar name='foo' at 0x7f70119bf220> at 0x7f70119bbf40>

    def sync_reset_function(token):
        assert foo.get() == "baz"
>       foo.reset(token)
E       ValueError: <Token var=<ContextVar name='foo' at 0x7f70119bf220> at 0x7f70119bbf40> was created in a different Context

[andrewgodwin]

What would you expect to happen - that everything shares the same context (both in and outside the sync_to_async/async_to_sync calls)? I believe that is what the intended effect is, but I didn't author the contextvar code in asgiref so I'd have to go re-learn it to see where this might be.

[adamantike]

Thanks for following up, Andrew! :)

Yes, I think that's what I would expect. This is the simplest test I found to reproduce the issue:

import contextvars

from asgiref.sync import sync_to_async
import pytest

foo = contextvars.ContextVar("foo")

@pytest.mark.asyncio
async def test_sync_to_async_contextvar_reset():
    token = foo.set("bar")
    await sync_to_async(foo.reset)(token)

• sync_to_async copies the current context to run the wrapped function:

  asgiref/asgiref/sync.py

  Lines 421 to 426 in 13d0b82

  	if contextvars is not None:
  	context = contextvars.copy_context()
  	child = functools.partial(self.func, *args, **kwargs)
  	func = context.run
  	args = (child,)
  	kwargs = {}

• However, CPython expects the exact same context reference for ContextVar.reset to work: https://github.com/python/cpython/blob/8363ac8607eca7398e568e1336154e1262a995a0/Python/context.c#L314-L319
• When copying the context, the object in memory is now different:

> /mnt/data/Proyectos/third_party/asgiref/asgiref/sync.py(423)__call__()
    422             context = contextvars.copy_context()
--> 423             child = functools.partial(self.func, *args, **kwargs)
    424             func = context.run

ipdb> !context
<Context object at 0x7fcf31388f00>
ipdb> contextvars.copy_context()
<Context object at 0x7fcf30a76700>
ipdb> context2 = contextvars.copy_context()
ipdb> !context
<Context object at 0x7fcf31388f00>
ipdb> !context2
<Context object at 0x7fcf30a76700>
ipdb> contextvars.copy_context()
<Context object at 0x7fcf31394a40>

[andrewgodwin]

Hm, interesting. I suspect whoever initially wrote this didn't test the reset() usage, then, and instead were just trying to persist values inside.

Unless you want to dive into this and see if you can make it work "as intended", I'll have to see if I can get around to looking at this later in the month (I'm mostly taking a break from OSS at the moment) - at the very worst, I'd hope we can at least get better errors and document the intended behaviours.