Accept explicit contextvars.Context in asyncio create_task() API #91150

Activity

1. asvetlov commented on Mar 12, 2022

   

   asvetlov

   on Mar 12, 2022

   ContributorAuthor

   More actions

   Now asyncio creates a new context copy on task creation.

   It is the perfect behavior *by default* and should stay as is.

   However, sometimes passing an explicit context into a task and using back the context modified by task code is desired.

   The most obvious use-case is testing: both unittest and pytest have multi-phase test initialization (asyncSetUp() methods and async fixtures correspondingly). If asyncSetUp() updates context variable -- test code expects to see this change.

   Currently, unittest runs the setup-test-cleanup chain in a single task and uses an internal queue to push a new coroutine for execution and get results back. It works but is cumbersome.

   Another solution is to create a task per test execution step and wrap the task creation with Context.run(). The problem is in getting the updated context back. A task creates a context copy on starting, thus the modified context is stored in the task internal attribute only. To get it back a trampoline async function should be used, e.g.

   async def wrapper(coro):
       try:
           return await coro
       finally:
           context = contextvars.copy_context()
           # store the context copy somewhere

   Again, it looks more complicated as it should be.

   The proposal is:

   1. Add 'context' keyword-only argument to asyncio.create_task() and loop.create_task().
   2. Use this context if explicitly passed, create a copy of the current context otherwise.

   The proposal is backward-compatible. Low-level API (call_soon(), call_later() etc.) already accept 'context' argument.

   The attached PR demonstrates how the proposed API simplifies unittest.IsolatedAsyncioTestCase internals.
2. asvetlov

   added

   3.11only security fixesonly security fixes

   topic-asyncio

   type-featureA feature request or enhancementA feature request or enhancement

   on Mar 12, 2022
3. 1st1 commented on Mar 13, 2022

   

   1st1

   on Mar 13, 2022

   Member

   More actions

   Yeah, +1 to add a parameter. Fwiw it was on my idea list when i was working on the pep
4. asvetlov

   closed this as completedon Mar 14, 2022
5. asvetlov

   closed this as completedon Mar 14, 2022
6. asvetlov commented on Mar 14, 2022

   

   asvetlov

   on Mar 14, 2022

   ContributorAuthor

   More actions

   New changeset 9523c0d by Andrew Svetlov in branch 'main':
   bpo-46994: Accept explicit contextvars.Context in asyncio create_task() API (GH-31837)
   9523c0d
7. vstinner commented on Mar 18, 2022

   

   vstinner

   on Mar 18, 2022

   Member

   More actions

   Nice enhancement!
8. ezio-melotti

   transferred this issue fromon Apr 10, 2022
9. nzig commented on May 25, 2022

   

   nzig

   on May 25, 2022

   Last edited by nzig

   More actions

   Any idea how this could be backported to older Python versions? I'm using 3.8. The best I've come up with is

   lock = threading.Lock()
   
   def create_task_with_context(coro, *, name= None, context = None):
       if context is None:
           return asyncio.create_task(coro, name=name)
   
       this_thread = threading.get_ident()
   
       def patched():
           if threading.get_ident() == this_thread:
               return context
           else:
               return original_copy_context()
   
       with lock:
           original_copy_context = contextvars.copy_context
           contextvars.copy_context = patched
           try:
               task = asyncio.tasks._PyTask(coro, name=name)
           finally:
               contextvars.copy_context = original_copy_context
   
       return task

   Which is a horrible hack. Any better ideas?

   Edit: Added thread safety check.

   👍React with 👍1Reacted by Adrian Garcia Badaracco
10. agronholm commented on May 25, 2022

    

    agronholm

    on May 25, 2022

    Contributor

    More actions

    Which is a horrible hack. Any better ideas?

    Never mind the hackiness, this is outright dangerous. What if there's another event loop running in another thread, doing the same exact thing?
11. nzig commented on May 25, 2022

    

    nzig

    on May 25, 2022

    Last edited by nzig

    More actions

    @agronholm I think the GIL protects us here. We don't do anything that releases the GIL between patching and unpatching the module, so it's thread safe.
    Never mind, I was confused and this is not thread safe.
12. agronholm commented on May 25, 2022

    

    agronholm

    on May 25, 2022

    Contributor

    More actions

    The GIL can switch threads any time between byte code instructions. Here's the disassembled byte code for that function of yours:

    Disassembly of <code object create_task_with_context at 0x7fc7974bf470, file "test.py", line 1>:
      2           0 LOAD_DEREF               0 (context)
                  2 LOAD_CONST               0 (None)
                  4 IS_OP                    0
                  6 POP_JUMP_IF_FALSE       11 (to 22)
    
      3           8 LOAD_GLOBAL              0 (asyncio)
                 10 LOAD_ATTR                1 (create_task)
                 12 LOAD_FAST                0 (coro)
                 14 LOAD_FAST                1 (name)
                 16 LOAD_CONST               1 (('name',))
                 18 CALL_FUNCTION_KW         2
                 20 RETURN_VALUE
    
      5     >>   22 LOAD_GLOBAL              2 (contextvars)
                 24 LOAD_ATTR                3 (copy_context)
                 26 STORE_FAST               3 (original_copy_context)
    
      6          28 LOAD_CLOSURE             0 (context)
                 30 BUILD_TUPLE              1
                 32 LOAD_CONST               2 (<code object <lambda> at 0x7fc7974bf3c0, file "test.py", line 6>)
                 34 LOAD_CONST               3 ('create_task_with_context.<locals>.<lambda>')
                 36 MAKE_FUNCTION            8 (closure)
                 38 LOAD_GLOBAL              2 (contextvars)
                 40 STORE_ATTR               3 (copy_context)
    
      7          42 SETUP_FINALLY           14 (to 72)
    
      8          44 LOAD_GLOBAL              0 (asyncio)
                 46 LOAD_ATTR                4 (tasks)
                 48 LOAD_ATTR                5 (_PyTask)
                 50 LOAD_FAST                0 (coro)
                 52 LOAD_FAST                1 (name)
                 54 LOAD_CONST               1 (('name',))
                 56 CALL_FUNCTION_KW         2
                 58 STORE_FAST               4 (task)
                 60 POP_BLOCK
    
     10          62 LOAD_FAST                3 (original_copy_context)
                 64 LOAD_GLOBAL              2 (contextvars)
                 66 STORE_ATTR               3 (copy_context)
    
     12          68 LOAD_FAST                4 (task)
                 70 RETURN_VALUE
    
     10     >>   72 LOAD_FAST                3 (original_copy_context)
                 74 LOAD_GLOBAL              2 (contextvars)
                 76 STORE_ATTR               3 (copy_context)
                 78 RERAISE                  0
    
13. agronholm commented on May 25, 2022

    

    agronholm

    on May 25, 2022

    Contributor

    More actions

    The proposal is backward-compatible. Low-level API (call_soon(), call_later() etc.) already accept 'context' argument.

    Depends on the definition of backwards compatible, I guess? This particular change broke AnyIO's tests with uvloop: agronholm/anyio#439 (not that I'm against this addition, but...).