Repository: django/asgiref
Pull request #227 (2021-01-14 15:02:29)

Internalize thread sensitive context management in asgiref/sync.py
Rather than using an external context management function, this commit
updates the thread_sensitive mode to reference an internal contextvar.

This approach mainly allows for the thread executor to be shut down when
the context exits, rather than relying on the thread pool shutting down
when going out of scope.

This approach also has other advantages:
* Maintains the existing sync_to_async interfaces.

* Context variable management cannot be broken by omitting the
  current_context_func arg

* This approach opens the door to more complex management techniques
  such as enforcing a thread pool limit.

Django PR showing implementation: https://github.com/django/django/pull/13882

a-feld: Thoughts on this approach @andrewgodwin ? PR is in draft because there are a couple more tests to write. I was racking my brain with how to manage ThreadPoolExecutors with the previous approach - this seems to be a bit more workable 😄 
[Django PR](https://github.com/django/django/pull/13882)
andrewgodwin: I'll try to take a look at this next weekend; don't have much time right now

Patch:
diff --git a/MANIFEST.in b/MANIFEST.in
index 065cc87..0aa915a 100644
--- a/MANIFEST.in
+++ b/MANIFEST.in
@@ -1,2 +1,3 @@
 include LICENSE
+include asgiref/py.typed
 recursive-include tests *.py
diff --git a/README.rst b/README.rst
index 259f31a..a7ebc04 100644
--- a/README.rst
+++ b/README.rst
@@ -163,14 +163,9 @@ The first compromise you get to might be that ``thread_sensitive`` code should
 just run in the same thread and not spawn in a sub-thread, fulfilling the first
 restriction, but that immediately runs you into the second restriction.
 
-By default, a variant of ThreadPoolExecutor executes any ``thread_sensitive``
-code on the outermost synchronous thread - either the main thread, or a single
-spawned subthread.
-
-There is an option to override the default single-threaded behavior so that
-there is 1 synchronous thread per context. If ``current_context_func`` is
-specified, this function will be called to retrieve the current context. There
-will be exactly 1 synchronous thread per context in this case.
+The only real solution is to essentially have a variant of ThreadPoolExecutor
+that executes any ``thread_sensitive`` code on the outermost synchronous
+thread - either the main thread, or a single spawned subthread.
 
 This means you now have two basic states:
 
diff --git a/asgiref/py.typed b/asgiref/py.typed
new file mode 100644
index 0000000..e69de29
diff --git a/asgiref/sync.py b/asgiref/sync.py
index d62ee86..13a8422 100644
--- a/asgiref/sync.py
+++ b/asgiref/sync.py
@@ -27,6 +27,59 @@ def _restore_context(context):
             cvar.set(context.get(cvar))
 
 
+class ThreadSensitiveContext:
+    """Async context manager to manage context for thread sensitive mode
+
+    This context manager controls which thread pool executor is used when in
+    thread sensitive mode. By default, a single thread pool executor is shared
+    within a process.
+
+    In Python 3.7+, the ThreadSensitiveContext() context manager may be used to
+    specify a thread pool per context.
+
+    In Python 3.6, usage of this context manager has no effect.
+
+    This context manager is re-entrant, so only the outer-most call to
+    ThreadSensitiveContext will set the context.
+
+    Usage:
+
+    >>> import time
+    >>> async with ThreadSensitiveContext():
+    ...     await sync_to_async(time.sleep, 1)()
+    """
+
+    def __init__(self):
+        self.token = None
+
+    if contextvars:
+
+        async def __aenter__(self):
+            try:
+                SyncToAsync.thread_sensitive_context.get()
+            except LookupError:
+                self.token = SyncToAsync.thread_sensitive_context.set(self)
+
+            return self
+
+        async def __aexit__(self, exc, value, tb):
+            if not self.token:
+                return
+
+            executor = SyncToAsync.context_to_thread_executor.pop(self, None)
+            if executor:
+                executor.shutdown()
+            SyncToAsync.thread_sensitive_context.reset(self.token)
+
+    else:
+
+        async def __aenter__(self):
+            return self
+
+        async def __aexit__(self, exc, value, tb):
+            pass
+
+
 class AsyncToSync:
     """
     Utility class which turns an awaitable that only works on the thread with
@@ -232,12 +285,6 @@ class SyncToAsync:
     outer code. This is needed for underlying Python code that is not
     threadsafe (for example, code which handles SQLite database connections).
 
-    If current_context_func is passed, the code will run 1 thread per context.
-    As an example, this may be used to create a per-request synchronous thread
-    by specifying the request object as the context. Thread scheduling will
-    occur by request in this scenario - each request will execute synchronous
-    work within the same thread.
-
     If the outermost program is async (i.e. SyncToAsync is outermost), then
     this will be a dedicated single sub-thread that all sync code runs in,
     one after the other. If the outermost program is sync (i.e. AsyncToSync is
@@ -262,16 +309,21 @@ class SyncToAsync:
     # Single-thread executor for thread-sensitive code
     single_thread_executor = ThreadPoolExecutor(max_workers=1)
 
+    # Maintain a contextvar for the current execution context. Optionally used
+    # for thread sensitive mode.
+    thread_sensitive_context = (
+        contextvars.ContextVar("thread_sensitive_context") if contextvars else None
+    )
+
     # Maintaining a weak reference to the context ensures that thread pools are
     # erased once the context goes out of scope. This terminates the thread pool.
     context_to_thread_executor = weakref.WeakKeyDictionary()
 
-    def __init__(self, func, thread_sensitive=True, current_context_func=None):
+    def __init__(self, func, thread_sensitive=True):
         self.func = func
         functools.update_wrapper(self, func)
         self._thread_sensitive = thread_sensitive
         self._is_coroutine = asyncio.coroutines._is_coroutine
-        self._current_context_func = current_context_func
         try:
             self.__self__ = func.__self__
         except AttributeError:
@@ -285,18 +337,20 @@ class SyncToAsync:
             if hasattr(AsyncToSync.executors, "current"):
                 # If we have a parent sync thread above somewhere, use that
                 executor = AsyncToSync.executors.current
-            elif self._current_context_func:
+            elif self.thread_sensitive_context and self.thread_sensitive_context.get(
+                None
+            ):
                 # If we have a way of retrieving the current context, attempt
                 # to use a per-context thread pool executor
-                current_context = self._current_context_func()
+                thread_sensitive_context = self.thread_sensitive_context.get()
 
-                if current_context in self.context_to_thread_executor:
+                if thread_sensitive_context in self.context_to_thread_executor:
                     # Re-use thread executor in current context
-                    executor = self.context_to_thread_executor[current_context]
+                    executor = self.context_to_thread_executor[thread_sensitive_context]
                 else:
                     # Create new thread executor in current context
                     executor = ThreadPoolExecutor(max_workers=1)
-                    self.context_to_thread_executor[current_context] = executor
+                    self.context_to_thread_executor[thread_sensitive_context] = executor
             else:
                 # Otherwise, we run it in a fixed single thread
                 executor = self.single_thread_executor
@@ -393,15 +447,13 @@ class SyncToAsync:
 async_to_sync = AsyncToSync
 
 
-def sync_to_async(func=None, thread_sensitive=True, current_context_func=None):
+def sync_to_async(func=None, thread_sensitive=True):
     if func is None:
         return lambda f: SyncToAsync(
             f,
             thread_sensitive=thread_sensitive,
-            current_context_func=current_context_func,
         )
     return SyncToAsync(
         func,
         thread_sensitive=thread_sensitive,
-        current_context_func=current_context_func,
     )
diff --git a/asgiref/typing.py b/asgiref/typing.py
new file mode 100644
index 0000000..af0ac3f
--- /dev/null
+++ b/asgiref/typing.py
@@ -0,0 +1,205 @@
+import sys
+from typing import Awaitable, Callable, Dict, Iterable, Optional, Tuple, Type, Union
+
+if sys.version_info >= (3, 8):
+    from typing import Literal, Protocol, TypedDict
+else:
+    from typing_extensions import Literal, Protocol, TypedDict
+
+
+class ASGIVersions(TypedDict):
+    spec_version: str
+    version: Union[Literal["2.0"], Literal["3.0"]]
+
+
+class HTTPScope(TypedDict):
+    type: Literal["http"]
+    asgi: ASGIVersions
+    http_version: str
+    method: str
+    scheme: str
+    path: str
+    raw_path: bytes
+    query_string: bytes
+    root_path: str
+    headers: Iterable[Tuple[bytes, bytes]]
+    client: Optional[Tuple[str, int]]
+    server: Optional[Tuple[str, Optional[int]]]
+    extensions: Dict[str, Dict[object, object]]
+
+
+class WebsocketScope(TypedDict):
+    type: Literal["websocket"]
+    asgi: ASGIVersions
+    http_version: str
+    scheme: str
+    path: str
+    raw_path: bytes
+    query_string: bytes
+    root_path: str
+    headers: Iterable[Tuple[bytes, bytes]]
+    client: Optional[Tuple[str, int]]
+    server: Optional[Tuple[str, Optional[int]]]
+    subprotocols: Iterable[str]
+    extensions: Dict[str, Dict[object, object]]
+
+
+class LifespanScope(TypedDict):
+    type: Literal["lifespan"]
+    asgi: ASGIVersions
+
+
+WWWScope = Union[HTTPScope, WebsocketScope]
+Scope = Union[HTTPScope, WebsocketScope, LifespanScope]
+
+
+class HTTPRequestEvent(TypedDict):
+    type: Literal["http.request"]
+    body: bytes
+    more_body: bool
+
+
+class HTTPResponseStartEvent(TypedDict):
+    type: Literal["http.response.start"]
+    status: int
+    headers: Iterable[Tuple[bytes, bytes]]
+
+
+class HTTPResponseBodyEvent(TypedDict):
+    type: Literal["http.response.body"]
+    body: bytes
+    more_body: bool
+
+
+class HTTPServerPushEvent(TypedDict):
+    type: Literal["http.response.push"]
+    path: str
+    headers: Iterable[Tuple[bytes, bytes]]
+
+
+class HTTPDisconnectEvent(TypedDict):
+    type: Literal["http.disconnect"]
+
+
+class WebsocketConnectEvent(TypedDict):
+    type: Literal["websocket.connect"]
+
+
+class WebsocketAcceptEvent(TypedDict):
+    type: Literal["websocket.accept"]
+    subprotocol: Optional[str]
+    headers: Iterable[Tuple[bytes, bytes]]
+
+
+class WebsocketReceiveEvent(TypedDict):
+    type: Literal["websocket.receive"]
+    bytes: Optional[bytes]
+    text: Optional[str]
+
+
+class WebsocketSendEvent(TypedDict):
+    type: Literal["websocket.send"]
+    bytes: Optional[bytes]
+    text: Optional[str]
+
+
+class WebsocketResponseStartEvent(TypedDict):
+    type: Literal["websocket.http.response.start"]
+    status: int
+    headers: Iterable[Tuple[bytes, bytes]]
+
+
+class WebsocketResponseBodyEvent(TypedDict):
+    type: Literal["websocket.http.response.body"]
+    body: bytes
+    more_body: bool
+
+
+class WebsocketDisconnectEvent(TypedDict):
+    type: Literal["websocket.disconnect"]
+    code: int
+
+
+class WebsocketCloseEvent(TypedDict):
+    type: Literal["websocket.close"]
+    code: int
+
+
+class LifespanStartupEvent(TypedDict):
+    type: Literal["lifespan.startup"]
+
+
+class LifespanShutdownEvent(TypedDict):
+    type: Literal["lifespan.shutdown"]
+
+
+class LifespanStartupCompleteEvent(TypedDict):
+    type: Literal["lifespan.startup.complete"]
+
+
+class LifespanStartupFailedEvent(TypedDict):
+    type: Literal["lifespan.startup.failed"]
+    message: str
+
+
+class LifespanShutdownCompleteEvent(TypedDict):
+    type: Literal["lifespan.shutdown.complete"]
+
+
+class LifespanShutdownFailedEvent(TypedDict):
+    type: Literal["lifespan.shutdown.failed"]
+    message: str
+
+
+ASGIReceiveEvent = Union[
+    HTTPRequestEvent,
+    HTTPDisconnectEvent,
+    WebsocketConnectEvent,
+    WebsocketReceiveEvent,
+    WebsocketDisconnectEvent,
+    LifespanStartupEvent,
+    LifespanShutdownEvent,
+]
+
+
+ASGISendEvent = Union[
+    HTTPResponseStartEvent,
+    HTTPResponseBodyEvent,
+    HTTPServerPushEvent,
+    HTTPDisconnectEvent,
+    WebsocketAcceptEvent,
+    WebsocketSendEvent,
+    WebsocketResponseStartEvent,
+    WebsocketResponseBodyEvent,
+    WebsocketCloseEvent,
+    LifespanStartupCompleteEvent,
+    LifespanStartupFailedEvent,
+    LifespanShutdownCompleteEvent,
+    LifespanShutdownFailedEvent,
+]
+
+
+ASGIReceiveCallable = Callable[[], Awaitable[ASGIReceiveEvent]]
+ASGISendCallable = Callable[[ASGISendEvent], Awaitable[None]]
+
+
+class ASGI2Protocol(Protocol):
+    def __init__(self, scope: Scope) -> None:
+        ...
+
+    async def __call__(
+        self, receive: ASGIReceiveCallable, send: ASGISendCallable
+    ) -> None:
+        ...
+
+
+ASGI2Application = Type[ASGI2Protocol]
+ASGI3Application = Callable[
+    [
+        Scope,
+        ASGIReceiveCallable,
+        ASGISendCallable,
+    ],
+    Awaitable[None],
+]
+ASGIApplication = Union[ASGI2Application, ASGI3Application]
diff --git a/setup.cfg b/setup.cfg
index b335e29..0bd7dc9 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -31,6 +31,8 @@ project_urls =
 python_requires = >=3.5
 packages = find:
 include_package_data = true
+install_requires =
+    typing_extensions; python_version < "3.8"
 zip_safe = false
 
 [options.extras_require]
@@ -45,3 +47,6 @@ testpaths = tests
 exclude = venv/*,tox/*,specs/*
 ignore = E123,E128,E266,E402,W503,E731,W601
 max-line-length = 119
+
+[isort]
+line_length = 119


Test patch:
diff --git a/tests/test_sync.py b/tests/test_sync.py
index 168cb39..5e9c171 100644
--- a/tests/test_sync.py
+++ b/tests/test_sync.py
@@ -8,7 +8,7 @@ from unittest import TestCase
 
 import pytest
 
-from asgiref.sync import async_to_sync, sync_to_async
+from asgiref.sync import ThreadSensitiveContext, async_to_sync, sync_to_async
 
 
 @pytest.mark.asyncio
@@ -323,54 +323,54 @@ async def test_thread_sensitive_outside_async():
 
 
 @pytest.mark.asyncio
-async def test_thread_sensitive_with_context_different():
+async def test_thread_sensitive_with_context_matches():
     result_1 = {}
     result_2 = {}
 
-    class Context:
-        pass
-
-    context_1 = Context()
-    context_2 = Context()
-
     def store_thread(result):
         result["thread"] = threading.current_thread()
 
-    fn_1 = sync_to_async(store_thread, current_context_func=lambda: context_1)
-    fn_2 = sync_to_async(store_thread, current_context_func=lambda: context_2)
+    store_thread_async = sync_to_async(store_thread)
+
+    async def fn():
+        async with ThreadSensitiveContext():
+            # Run it (in supposed parallel!)
+            await asyncio.wait(
+                [store_thread_async(result_1), store_thread_async(result_2)]
+            )
 
-    # Run it (in true parallel!)
-    await asyncio.wait([fn_1(result_1), fn_2(result_2)])
+    await fn()
 
-    # They should not have run in the main thread, and on different threads
+    # 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_2["thread"]
 
 
 @pytest.mark.asyncio
-async def test_thread_sensitive_with_context_matches():
+async def test_thread_sensitive_nested_context():
     result_1 = {}
     result_2 = {}
 
-    class Context:
-        pass
-
-    context = Context()
-
+    @sync_to_async
     def store_thread(result):
         result["thread"] = threading.current_thread()
 
-    fn_1 = sync_to_async(store_thread, current_context_func=lambda: context)
-    fn_2 = sync_to_async(store_thread, current_context_func=lambda: context)
-
-    # Run it (in supposed parallel!)
-    await asyncio.wait([fn_1(result_1), fn_2(result_2)])
+    async with ThreadSensitiveContext():
+        await store_thread(result_1)
+        async with ThreadSensitiveContext():
+            await store_thread(result_2)
 
-    # They should not have run in the main thread, and on different threads
+    # 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"]
 
 
+@pytest.mark.asyncio
+async def test_thread_sensitive_context_without_sync_work():
+    async with ThreadSensitiveContext():
+        pass
+
+
 def test_thread_sensitive_double_nested_sync():
     """
     Tests that thread_sensitive SyncToAsync nests inside itself where the
diff --git a/tests/test_sync_contextvars.py b/tests/test_sync_contextvars.py
index af3e3be..b1027aa 100644
--- a/tests/test_sync_contextvars.py
+++ b/tests/test_sync_contextvars.py
@@ -1,15 +1,37 @@
 import asyncio
+import threading
 import time
 
 import pytest
 
-from asgiref.sync import async_to_sync, sync_to_async
+from asgiref.sync import ThreadSensitiveContext, async_to_sync, sync_to_async
 
 contextvars = pytest.importorskip("contextvars")
 
 foo = contextvars.ContextVar("foo")
 
 
+@pytest.mark.asyncio
+async def test_thread_sensitive_with_context_different():
+    result_1 = {}
+    result_2 = {}
+
+    @sync_to_async
+    def store_thread(result):
+        result["thread"] = threading.current_thread()
+
+    async def fn(result):
+        async with ThreadSensitiveContext():
+            await store_thread(result)
+
+    # Run it (in true parallel!)
+    await asyncio.wait([fn(result_1), fn(result_2)])
+
+    # They should not have run in the main thread, and on different threads
+    assert result_1["thread"] != threading.current_thread()
+    assert result_1["thread"] != result_2["thread"]
+
+
 @pytest.mark.asyncio
 async def test_sync_to_async_contextvars():
     """