Fixed #36714 -- Fixed context sharing among async signal handlers.

2 files changed, 273 insertions, 20 deletions

diff --git a/django/dispatch/dispatcher.py b/django/dispatch/dispatcher.py
index 63fb75285e..21d77bd884 100644
--- a/django/dispatch/dispatcher.py
+++ b/django/dispatch/dispatcher.py
@@ -1,4 +1,5 @@
 import asyncio
+import contextvars
 import logging
 import threading
 import weakref
@@ -22,13 +23,30 @@ NONE_ID = _make_id(None)
 NO_RECEIVERS = object()
 
 
-async def _gather(*coros):
+def _restore_context(context):
+    """
+    Check for changes in contextvars, and set them to the current
+    context for downstream consumers.
+    """
+    for cvar in context:
+        cvalue = context.get(cvar)
+        try:
+            if cvar.get() != cvalue:
+                cvar.set(cvalue)
+        except LookupError:
+            cvar.set(cvalue)
+
+
+async def _run_parallel(*coros):
+    """
+    Execute multiple asynchronous coroutines in parallel,
+    sharing the current context between them.
+    """
+    context = contextvars.copy_context()
+
     if len(coros) == 0:
         return []
 
-    if len(coros) == 1:
-        return [await coros[0]]
-
     async def run(i, coro):
         results[i] = await coro
 
@@ -36,12 +54,14 @@ async def _gather(*coros):
         async with asyncio.TaskGroup() as tg:
             results = [None] * len(coros)
             for i, coro in enumerate(coros):
-                tg.create_task(run(i, coro))
+                tg.create_task(run(i, coro), context=context)
         return results
     except BaseExceptionGroup as exception_group:
         if len(exception_group.exceptions) == 1:
             raise exception_group.exceptions[0]
         raise
+    finally:
+        _restore_context(context=context)
 
 
 class Signal:
@@ -233,7 +253,7 @@ class Signal:
         if async_receivers:
 
             async def asend():
-                async_responses = await _gather(
+                async_responses = await _run_parallel(
                     *(
                         receiver(signal=self, sender=sender, **named)
                         for receiver in async_receivers
@@ -275,6 +295,7 @@ class Signal:
         ):
             return []
         sync_receivers, async_receivers = self._live_receivers(sender)
+
         if sync_receivers:
 
             @sync_to_async
@@ -290,14 +311,12 @@ class Signal:
             async def sync_send():
                 return []
 
-        responses, async_responses = await _gather(
-            sync_send(),
-            _gather(
-                *(
-                    receiver(signal=self, sender=sender, **named)
-                    for receiver in async_receivers
-                )
-            ),
+        responses = await sync_send()
+        async_responses = await _run_parallel(
+            *(
+                receiver(signal=self, sender=sender, **named)
+                for receiver in async_receivers
+            )
         )
         responses.extend(zip(async_receivers, async_responses))
         return responses
@@ -362,7 +381,7 @@ class Signal:
                 return response
 
             async def asend():
-                async_responses = await _gather(
+                async_responses = await _run_parallel(
                     *(
                         asend_and_wrap_exception(receiver)
                         for receiver in async_receivers
@@ -436,11 +455,9 @@ class Signal:
                 return err
             return response
 
-        responses, async_responses = await _gather(
-            sync_send(),
-            _gather(
-                *(asend_and_wrap_exception(receiver) for receiver in async_receivers),
-            ),
+        responses = await sync_send()
+        async_responses = await _run_parallel(
+            *(asend_and_wrap_exception(receiver) for receiver in async_receivers),
         )
         responses.extend(zip(async_receivers, async_responses))
         return responses
diff --git a/tests/signals/tests.py b/tests/signals/tests.py
index 7cb64f6e05..907612459a 100644
--- a/tests/signals/tests.py
+++ b/tests/signals/tests.py
@@ -1,3 +1,4 @@
+import contextvars
 from unittest import mock
 
 from asgiref.sync import markcoroutinefunction
@@ -645,3 +646,238 @@ class AsyncReceiversTests(SimpleTestCase):
 
         result = await signal.asend_robust(self.__class__)
         self.assertEqual(result, [(async_handler, 1)])
+
+
+class TestReceiversContextVarsSharing(SimpleTestCase):
+    def setUp(self):
+        self.ctx_var = contextvars.ContextVar("test_var", default=0)
+
+        class CtxSyncHandler:
+            def __init__(self, ctx_var):
+                self.ctx_var = ctx_var
+                self.values = []
+
+            def __call__(self, **kwargs):
+                val = self.ctx_var.get()
+                self.ctx_var.set(val + 1)
+                self.values.append(self.ctx_var.get())
+                return self.ctx_var.get()
+
+        class CtxAsyncHandler:
+            def __init__(self, ctx_var):
+                self.ctx_var = ctx_var
+                self.values = []
+                markcoroutinefunction(self)
+
+            async def __call__(self, **kwargs):
+                val = self.ctx_var.get()
+                self.ctx_var.set(val + 1)
+                self.values.append(self.ctx_var.get())
+                return self.ctx_var.get()
+
+        self.CtxSyncHandler = CtxSyncHandler
+        self.CtxAsyncHandler = CtxAsyncHandler
+
+    async def test_asend_correct_contextvars_sharing_async_receivers(self):
+        handler1 = self.CtxAsyncHandler(self.ctx_var)
+        handler2 = self.CtxAsyncHandler(self.ctx_var)
+        signal = dispatch.Signal()
+        signal.connect(handler1)
+        signal.connect(handler2)
+
+        # set custom value outer signal
+        self.ctx_var.set(1)
+
+        await signal.asend(self.__class__)
+
+        self.assertEqual(len(handler1.values), 1)
+        self.assertEqual(len(handler2.values), 1)
+        self.assertEqual(sorted([*handler1.values, *handler2.values]), [2, 3])
+        self.assertEqual(self.ctx_var.get(), 3)
+
+    async def test_asend_correct_contextvars_sharing_sync_receivers(self):
+        handler1 = self.CtxSyncHandler(self.ctx_var)
+        handler2 = self.CtxSyncHandler(self.ctx_var)
+        signal = dispatch.Signal()
+        signal.connect(handler1)
+        signal.connect(handler2)
+
+        # set custom value outer signal
+        self.ctx_var.set(1)
+
+        await signal.asend(self.__class__)
+
+        self.assertEqual(len(handler1.values), 1)
+        self.assertEqual(len(handler2.values), 1)
+        self.assertEqual(sorted([*handler1.values, *handler2.values]), [2, 3])
+        self.assertEqual(self.ctx_var.get(), 3)
+
+    async def test_asend_correct_contextvars_sharing_mix_receivers(self):
+        handler1 = self.CtxSyncHandler(self.ctx_var)
+        handler2 = self.CtxAsyncHandler(self.ctx_var)
+        signal = dispatch.Signal()
+        signal.connect(handler1)
+        signal.connect(handler2)
+
+        # set custom value outer signal
+        self.ctx_var.set(1)
+
+        await signal.asend(self.__class__)
+
+        self.assertEqual(len(handler1.values), 1)
+        self.assertEqual(len(handler2.values), 1)
+        self.assertEqual(sorted([*handler1.values, *handler2.values]), [2, 3])
+        self.assertEqual(self.ctx_var.get(), 3)
+
+    async def test_asend_robust_correct_contextvars_sharing_async_receivers(self):
+        handler1 = self.CtxAsyncHandler(self.ctx_var)
+        handler2 = self.CtxAsyncHandler(self.ctx_var)
+        signal = dispatch.Signal()
+        signal.connect(handler1)
+        signal.connect(handler2)
+
+        # set custom value outer signal
+        self.ctx_var.set(1)
+
+        await signal.asend_robust(self.__class__)
+
+        self.assertEqual(len(handler1.values), 1)
+        self.assertEqual(len(handler2.values), 1)
+        self.assertEqual(sorted([*handler1.values, *handler2.values]), [2, 3])
+        self.assertEqual(self.ctx_var.get(), 3)
+
+    async def test_asend_robust_correct_contextvars_sharing_sync_receivers(self):
+        handler1 = self.CtxSyncHandler(self.ctx_var)
+        handler2 = self.CtxSyncHandler(self.ctx_var)
+        signal = dispatch.Signal()
+        signal.connect(handler1)
+        signal.connect(handler2)
+
+        # set custom value outer signal
+        self.ctx_var.set(1)
+
+        await signal.asend_robust(self.__class__)
+
+        self.assertEqual(len(handler1.values), 1)
+        self.assertEqual(len(handler2.values), 1)
+        self.assertEqual(sorted([*handler1.values, *handler2.values]), [2, 3])
+        self.assertEqual(self.ctx_var.get(), 3)
+
+    async def test_asend_robust_correct_contextvars_sharing_mix_receivers(self):
+        handler1 = self.CtxSyncHandler(self.ctx_var)
+        handler2 = self.CtxAsyncHandler(self.ctx_var)
+        signal = dispatch.Signal()
+        signal.connect(handler1)
+        signal.connect(handler2)
+
+        # set custom value outer signal
+        self.ctx_var.set(1)
+
+        await signal.asend_robust(self.__class__)
+
+        self.assertEqual(len(handler1.values), 1)
+        self.assertEqual(len(handler2.values), 1)
+        self.assertEqual(sorted([*handler1.values, *handler2.values]), [2, 3])
+        self.assertEqual(self.ctx_var.get(), 3)
+
+    def test_send_correct_contextvars_sharing_async_receivers(self):
+        handler1 = self.CtxAsyncHandler(self.ctx_var)
+        handler2 = self.CtxAsyncHandler(self.ctx_var)
+        signal = dispatch.Signal()
+        signal.connect(handler1)
+        signal.connect(handler2)
+
+        # set custom value outer signal
+        self.ctx_var.set(1)
+
+        signal.send(self.__class__)
+
+        self.assertEqual(len(handler1.values), 1)
+        self.assertEqual(len(handler2.values), 1)
+        self.assertEqual(sorted([*handler1.values, *handler2.values]), [2, 3])
+        self.assertEqual(self.ctx_var.get(), 3)
+
+    def test_send_correct_contextvars_sharing_sync_receivers(self):
+        handler1 = self.CtxSyncHandler(self.ctx_var)
+        handler2 = self.CtxSyncHandler(self.ctx_var)
+        signal = dispatch.Signal()
+        signal.connect(handler1)
+        signal.connect(handler2)
+
+        # set custom value outer signal
+        self.ctx_var.set(1)
+
+        signal.send(self.__class__)
+
+        self.assertEqual(len(handler1.values), 1)
+        self.assertEqual(len(handler2.values), 1)
+        self.assertEqual(sorted([*handler1.values, *handler2.values]), [2, 3])
+        self.assertEqual(self.ctx_var.get(), 3)
+
+    def test_send_correct_contextvars_sharing_mix_receivers(self):
+        handler1 = self.CtxSyncHandler(self.ctx_var)
+        handler2 = self.CtxAsyncHandler(self.ctx_var)
+        signal = dispatch.Signal()
+        signal.connect(handler1)
+        signal.connect(handler2)
+
+        # set custom value outer signal
+        self.ctx_var.set(1)
+
+        signal.send(self.__class__)
+
+        self.assertEqual(len(handler1.values), 1)
+        self.assertEqual(len(handler2.values), 1)
+        self.assertEqual(sorted([*handler1.values, *handler2.values]), [2, 3])
+        self.assertEqual(self.ctx_var.get(), 3)
+
+    def test_send_robust_correct_contextvars_sharing_async_receivers(self):
+        handler1 = self.CtxAsyncHandler(self.ctx_var)
+        handler2 = self.CtxAsyncHandler(self.ctx_var)
+        signal = dispatch.Signal()
+        signal.connect(handler1)
+        signal.connect(handler2)
+
+        # set custom value outer signal
+        self.ctx_var.set(1)
+
+        signal.send_robust(self.__class__)
+
+        self.assertEqual(len(handler1.values), 1)
+        self.assertEqual(len(handler2.values), 1)
+        self.assertEqual(sorted([*handler1.values, *handler2.values]), [2, 3])
+        self.assertEqual(self.ctx_var.get(), 3)
+
+    def test_send_robust_correct_contextvars_sharing_sync_receivers(self):
+        handler1 = self.CtxSyncHandler(self.ctx_var)
+        handler2 = self.CtxSyncHandler(self.ctx_var)
+        signal = dispatch.Signal()
+        signal.connect(handler1)
+        signal.connect(handler2)
+
+        # set custom value outer signal
+        self.ctx_var.set(1)
+
+        signal.send_robust(self.__class__)
+
+        self.assertEqual(len(handler1.values), 1)
+        self.assertEqual(len(handler2.values), 1)
+        self.assertEqual(sorted([*handler1.values, *handler2.values]), [2, 3])
+        self.assertEqual(self.ctx_var.get(), 3)
+
+    def test_send_robust_correct_contextvars_sharing_mix_receivers(self):
+        handler1 = self.CtxSyncHandler(self.ctx_var)
+        handler2 = self.CtxAsyncHandler(self.ctx_var)
+        signal = dispatch.Signal()
+        signal.connect(handler1)
+        signal.connect(handler2)
+
+        # set custom value outer signal
+        self.ctx_var.set(1)
+
+        signal.send_robust(self.__class__)
+
+        self.assertEqual(len(handler1.values), 1)
+        self.assertEqual(len(handler2.values), 1)
+        self.assertEqual(sorted([*handler1.values, *handler2.values]), [2, 3])
+        self.assertEqual(self.ctx_var.get(), 3)