perf: resolve FromDI without a worker thread (#85)

Author: lesnik512

modern_di_faststream/main.py

@@ -1,3 +1,4 @@
+import contextvars
 import dataclasses
 import typing
 from collections.abc import Awaitable, Callable
@@ -13,11 +14,10 @@
 
 faststream_message_provider = providers.ContextProvider(scope=Scope.REQUEST, context_type=faststream.StreamMessage)
 
-# Keys under which the containers live in FastStream's ``ContextRepo``. Each is
-# written in one place and read in another; naming them keeps writer and reader
-# in provable agreement instead of relying on two matching string literals.
 _ROOT_CONTAINER_KEY = "modern_di_root_container"
-_REQUEST_CONTAINER_KEY = "modern_di_request_container"
+_request_container: contextvars.ContextVar[Container | None] = contextvars.ContextVar(
+    "modern_di_request_container", default=None
+)
 
 _MISSING_REQUEST_CONTAINER = (
     "No request container for this message, so the DI middleware did not run for it. "
@@ -56,11 +56,14 @@ async def consume_scope(
         async with self.di_container.build_child_container(
             scope=match.scope, context=match.context
         ) as request_container:
-            with self.context.scope(_REQUEST_CONTAINER_KEY, request_container):
+            token = _request_container.set(request_container)
+            try:
                 return typing.cast(
                     typing.AsyncIterator[DecodedMessage],
                     await call_next(msg),
                 )
+            finally:
+                _request_container.reset(token)
 
 
 def fetch_di_container(app_: faststream.FastStream | AsgiFastStream) -> Container:
@@ -79,7 +82,7 @@ def setup_di(
 
     # Installed on startup rather than here so a broker added after ``setup_di`` is covered,
     # and skipped where present so a restart adds no second copy: docs/adr/0002-install-middleware-on-startup.md
-    def install_middleware() -> None:
+    async def install_middleware() -> None:
         if not app.brokers:
             raise RuntimeError(_NO_BROKER_AT_STARTUP)
         for broker in app.brokers:
@@ -91,7 +94,7 @@ def install_middleware() -> None:
         close_registered = False
 
         # Appended on startup so the close runs after every after_shutdown hook registered before it.
-        def open_root() -> None:
+        async def open_root() -> None:
             nonlocal close_registered
             container.open()
             if not close_registered:
@@ -106,8 +109,8 @@ def open_root() -> None:
 class Dependency(typing.Generic[T_co]):
     marker: integrations.Marker[T_co]
 
-    async def __call__(self, context: faststream.ContextRepo) -> T_co:
-        request_container: Container | None = context.get_local(_REQUEST_CONTAINER_KEY)
+    async def __call__(self) -> T_co:
+        request_container = _request_container.get()
         if request_container is None:
             raise RuntimeError(_MISSING_REQUEST_CONTAINER)
         return self.marker.resolve(request_container)

tests/test_faststream_di.py

@@ -1,5 +1,7 @@
 import typing
+from unittest import mock
 
+import anyio.to_thread
 import faststream
 import pytest
 from faststream import TestApp
@@ -33,6 +35,26 @@ async def index_subscriber(
         await br.publish("test", TEST_SUBJECT)
 
 
+async def test_from_di_resolves_without_a_worker_thread(
+    app: faststream.FastStream, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    broker = typing.cast(NatsBroker, app.broker)
+    resolved: list[SimpleCreator] = []
+
+    @broker.subscriber(TEST_SUBJECT)
+    async def index_subscriber(instance: typing.Annotated[SimpleCreator, FromDI(SimpleCreator)]) -> None:
+        resolved.append(instance)
+
+    run_sync = mock.Mock(wraps=anyio.to_thread.run_sync)
+    async with TestNatsBroker(broker) as br, TestApp(app):
+        monkeypatch.setattr(anyio.to_thread, "run_sync", run_sync)
+        await br.publish(None, TEST_SUBJECT)
+        monkeypatch.undo()
+
+    assert [type(instance) for instance in resolved] == [SimpleCreator]
+    assert run_sync.call_count == 0
+
+
 async def test_context_adapter(app: faststream.FastStream) -> None:
     broker = typing.cast(NatsBroker, app.broker)
 

tests/test_lifespan.py

@@ -1,6 +1,8 @@
 import typing
 from collections.abc import Callable
+from unittest import mock
 
+import anyio.to_thread
 import faststream
 import pytest
 from faststream import TestApp
@@ -146,3 +148,16 @@ def user_hook() -> None:
         assert [(instance.finalized, at_hook) for instance, at_hook in seen] == [(1, 0)] * cycle
         assert container.closed
         assert container.close_calls == cycle
+
+
+async def test_startup_and_shutdown_run_no_worker_thread(
+    app: faststream.FastStream, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    run_sync = mock.Mock(wraps=anyio.to_thread.run_sync)
+    monkeypatch.setattr(anyio.to_thread, "run_sync", run_sync)
+    broker = typing.cast(NatsBroker, app.broker)
+
+    async with TestNatsBroker(broker), TestApp(app):
+        pass
+
+    assert run_sync.call_count == 0