implemented sharing the same event loop with AsyncSingleThreadContext

Author: Arfey

asgiref/sync.py

@@ -8,6 +8,7 @@
 import threading
 import warnings
 import weakref
+from asyncio import constants
 from concurrent.futures import Future, ThreadPoolExecutor
 from typing import (
     TYPE_CHECKING,
@@ -36,6 +37,10 @@
     # This is not available to import at runtime
     from _typeshed import OptExcInfo
 
+    AsyncSingleThreadContextMapType = weakref.WeakKeyDictionary[
+        "AsyncSingleThreadContext", tuple[ThreadPoolExecutor, asyncio.AbstractEventLoop]
+    ]
+
 _F = TypeVar("_F", bound=Callable[..., Any])
 _P = ParamSpec("_P")
 _R = TypeVar("_R")
@@ -97,11 +102,45 @@ def __enter__(self):
 
         return self
 
+    def _cancel_all_tasks(self, loop):
+        to_cancel = asyncio.all_tasks(loop)
+        if not to_cancel:
+            return
+
+        for task in to_cancel:
+            task.cancel()
+
+        loop.run_until_complete(asyncio.gather(*to_cancel, return_exceptions=True))
+
+        for task in to_cancel:
+            if task.cancelled():
+                continue
+            if task.exception() is not None:
+                loop.call_exception_handler(
+                    {
+                        "message": "unhandled exception during AsyncSingleThreadContext shutdown",
+                        "exception": task.exception(),
+                        "task": task,
+                    }
+                )
+
     def __exit__(self, exc, value, tb):
         if not self.token:
             return
 
-        executor = AsyncToSync.context_to_thread_executor.pop(self, None)
+        executor, loop = AsyncToSync.async_single_thread_context_map.pop(
+            self, (None, None)
+        )
+        if loop:
+            try:
+                self._cancel_all_tasks(loop)
+                loop.run_until_complete(loop.shutdown_asyncgens())
+                loop.run_until_complete(
+                    loop.shutdown_default_executor(constants.THREAD_JOIN_TIMEOUT)
+                )
+            finally:
+                loop.close()
+
         if executor:
             executor.shutdown()
 
@@ -174,7 +213,7 @@ class AsyncToSync(Generic[_P, _R]):
         contextvars.ContextVar("async_single_thread_context")
     )
 
-    context_to_thread_executor: "weakref.WeakKeyDictionary[AsyncSingleThreadContext, ThreadPoolExecutor]" = (
+    async_single_thread_context_map: "AsyncSingleThreadContextMapType" = (
         weakref.WeakKeyDictionary()
     )
 
@@ -293,25 +332,29 @@ async def new_loop_wrap() -> None:
                 running_in_main_event_loop = False
 
             if not running_in_main_event_loop:
-                loop_executor = None
-
                 if self.async_single_thread_context.get(None):
                     single_thread_context = self.async_single_thread_context.get()
 
-                    if single_thread_context in self.context_to_thread_executor:
-                        loop_executor = self.context_to_thread_executor[
-                            single_thread_context
-                        ]
+                    if single_thread_context in self.async_single_thread_context_map:
+                        (
+                            loop_executor,
+                            context_loop,
+                        ) = self.async_single_thread_context_map[single_thread_context]
                     else:
+                        context_loop = asyncio.new_event_loop()
                         loop_executor = ThreadPoolExecutor(max_workers=1)
-                        self.context_to_thread_executor[
-                            single_thread_context
-                        ] = loop_executor
+                        self.async_single_thread_context_map[single_thread_context] = (
+                            loop_executor,
+                            context_loop,
+                        )
+
+                    loop_future = loop_executor.submit(
+                        context_loop.run_until_complete, new_loop_wrap()
+                    )
                 else:
                     # Make our own event loop - in a new thread - and run inside that.
                     loop_executor = ThreadPoolExecutor(max_workers=1)
-
-                loop_future = loop_executor.submit(asyncio.run, new_loop_wrap())
+                    loop_future = loop_executor.submit(asyncio.run, new_loop_wrap())
                 # Run the CurrentThreadExecutor until the future is done.
                 current_executor.run_until_future(loop_future)
                 # Wait for future and/or allow for exception propagation

tests/test_sync.py

@@ -15,6 +15,7 @@
 
 from asgiref.sync import (
     AsyncSingleThreadContext,
+    AsyncToSync,
     ThreadSensitiveContext,
     async_to_sync,
     iscoroutinefunction,
@@ -549,21 +550,24 @@ def inner(result):
 def test_async_single_thread_context_matches():
     """
     Tests that functions wrapped with async_to_sync and executed within an
-    AsyncSingleThreadContext run on the same thread, even without a main_event_loop.
+    AsyncSingleThreadContext run on the same thread/loop, even without a
+    main_event_loop.
     """
     result_1 = {}
     result_2 = {}
 
-    async def store_thread_async(result):
+    async def store_thread_info_async(result):
         result["thread"] = threading.current_thread()
+        result["loop"] = asyncio.get_running_loop()
 
     with AsyncSingleThreadContext():
-        async_to_sync(store_thread_async)(result_1)
-        async_to_sync(store_thread_async)(result_2)
+        async_to_sync(store_thread_info_async)(result_1)
+        async_to_sync(store_thread_info_async)(result_2)
 
     # They should not have run in the main thread, and on the same threads
     assert result_1["thread"] != threading.current_thread()
     assert result_1["thread"] == result_2["thread"]
+    assert result_1["loop"] == result_2["loop"]
 
 
 def test_async_single_thread_nested_context():
@@ -572,20 +576,28 @@ def test_async_single_thread_nested_context():
     """
     result_1 = {}
     result_2 = {}
+    result_3 = {}
 
     @async_to_sync
-    async def store_thread(result):
+    async def store_thread_info(result):
         result["thread"] = threading.current_thread()
+        result["loop"] = asyncio.get_running_loop()
 
     with AsyncSingleThreadContext():
-        store_thread(result_1)
+        store_thread_info(result_1)
 
         with AsyncSingleThreadContext():
-            store_thread(result_2)
+            store_thread_info(result_2)
+
+        # verify that the context does not shut down the executor or event loop
+        store_thread_info(result_3)
 
     # They should not have run in the main thread, and on the same threads
     assert result_1["thread"] != threading.current_thread()
     assert result_1["thread"] == result_2["thread"]
+    assert result_1["thread"] == result_3["thread"]
+    assert result_1["loop"] == result_2["loop"]
+    assert result_1["loop"] == result_3["loop"]
 
 
 def test_async_single_thread_context_without_async_work():
@@ -621,21 +633,54 @@ async def test_async_single_thread_context_matches_from_async_thread():
     """
     result_1 = {}
     result_2 = {}
+    result_3 = {}
 
     @async_to_sync
-    async def store_thread_async(result):
+    async def store_thread_info_async(result):
         result["thread"] = threading.current_thread()
+        result["loop"] = asyncio.get_running_loop()
 
     def inner():
         with AsyncSingleThreadContext():
-            store_thread_async(result_1)
-            store_thread_async(result_2)
+            store_thread_info_async(result_1)
+            store_thread_info_async(result_2)
+
+        # verify that the context does not shut down the executor or event loop
+        store_thread_info_async(result_3)
 
     await sync_to_async(inner)()
 
     # They should both have run in the current thread.
     assert result_1["thread"] == threading.current_thread()
     assert result_1["thread"] == result_2["thread"]
+    assert result_1["thread"] == result_3["thread"]
+    assert result_1["loop"] == result_2["loop"]
+    assert result_1["loop"] == result_3["loop"]
+
+
+def test_async_single_thread_context_shutdown():
+    """
+    Verifies the shutdown process works as expected.
+    """
+    result = {}
+
+    async def inner_long_running_task():
+        await asyncio.sleep(5)
+        result["inner_long_running_task"] = True
+
+    async def handler():
+        # start but do not await, so it runs in the background
+        asyncio.create_task(inner_long_running_task())
+        return True
+
+    with AsyncSingleThreadContext() as ctx:
+        assert async_to_sync(handler)()
+
+        executor, loop = AsyncToSync.async_single_thread_context_map[ctx]
+
+    assert loop.is_closed(), "Event loop was not closed"
+    assert executor._shutdown, "Executor was not shutdown"
+    assert not result, "Background tasks should have been canceled"
 
 
 @pytest.mark.asyncio