Don't nest contexts in modules (#137)

Author: davidbrochart

src/fps/_context.py

@@ -101,6 +101,7 @@ def __init__(
         self._exit_stack: ExitStack | None = None
         self._async_exit_stack: AsyncExitStack | None = None
         self._opened = False
+        self._closing = False
 
     def _drop(self, borrower: Value) -> None:
         if borrower in self._borrowers:
@@ -191,6 +192,11 @@ async def aclose(
         Raises:
             TimeoutError: If the shared value could not be closed in time.
         """
+        if self._closing:
+            return
+
+        self._closing = True
+
         if timeout is None:
             timeout = self._close_timeout
         if timeout is None:
@@ -228,7 +234,7 @@ class Context:
     def __init__(self) -> None:
         self._context: dict[int, SharedValue] = {}
         self._value_added = Event()
-        self._closed = Event()
+        self._closed = False
         self._teardown_callbacks: list[
             Callable[..., Any] | Callable[..., Awaitable[Any]]
         ] = []
@@ -246,21 +252,12 @@ async def __aenter__(self) -> Context:
         return self
 
     async def __aexit__(self, exc_type, exc_value, exc_tb):
-        for context in set(self._children):
-            await context._closed.wait()
-        async with create_task_group() as tg:
-            for shared_value in self._context.values():
-                tg.start_soon(
-                    partial(
-                        shared_value.aclose,
-                        _exc_type=exc_type,
-                        _exc_value=exc_value,
-                        _exc_tb=exc_tb,
-                    )
-                )
-            for callback in self._teardown_callbacks[::-1]:
-                await call(callback, exc_value)
-        self._closed.set()
+        await self.aclose(
+            timeout=None,
+            _exc_type=exc_type,
+            _exc_value=exc_value,
+            _exc_tb=exc_tb,
+        )
         _current_context.reset(self._token)
         if self._parent is not None:
             self._parent._children.remove(self)
@@ -277,7 +274,7 @@ def _get_value_types(
         return types
 
     def _check_closed(self):
-        if self._closed.is_set():
+        if self._closed:
             raise RuntimeError("Context is closed")
 
     def add_teardown_callback(
@@ -302,6 +299,7 @@ def put(
         teardown_callback: Callable[..., Any]
         | Callable[..., Awaitable[Any]]
         | None = None,
+        shared_value: SharedValue[T] | None = None,
     ) -> SharedValue[T]:
         """
         Put a value in the context so that it can be shared.
@@ -319,12 +317,16 @@ def put(
             The shared value.
         """
         self._check_closed()
-        _shared_value = SharedValue(
-            value,
-            max_borrowers=max_borrowers,
-            manage=manage,
-            teardown_callback=teardown_callback,
-        )
+        if shared_value is not None:
+            _shared_value = shared_value
+            value = _shared_value._value
+        else:
+            _shared_value = SharedValue(
+                value,
+                max_borrowers=max_borrowers,
+                manage=manage,
+                teardown_callback=teardown_callback,
+            )
         _types = self._get_value_types(value, types)
         for value_type in _types:
             value_type_id = id(value_type)
@@ -379,6 +381,41 @@ async def _get(self, value_type: type[T]) -> Value[T]:
                 return await shared_value.get()
             await self._value_added.wait()
 
+    async def aclose(
+        self,
+        *,
+        timeout: float | None = None,
+        _exc_type=None,
+        _exc_value: BaseException | None = None,
+        _exc_tb=None,
+    ) -> None:
+        """
+        Close the context, after all shared values that were borrowed have been dropped.
+        The shared values will be torn down, if applicable.
+
+        Args:
+            timeout: The time to wait for all shared values to be freed.
+
+        Raises:
+            TimeoutError: If the context could not be closed in time.
+        """
+        if timeout is None:
+            timeout = float("inf")
+        with fail_after(timeout):
+            async with create_task_group() as tg:
+                for shared_value in self._context.values():
+                    tg.start_soon(
+                        partial(
+                            shared_value.aclose,
+                            _exc_type=_exc_type,
+                            _exc_value=_exc_value,
+                            _exc_tb=_exc_tb,
+                        )
+                    )
+                for callback in self._teardown_callbacks[::-1]:
+                    await call(callback, _exc_value)
+        self._closed = True
+
 
 @lru_cache(maxsize=1024)
 def count_parameters(func: Callable) -> int:

src/fps/_module.py

@@ -6,11 +6,11 @@
 from collections.abc import Callable, Awaitable
 from contextlib import AsyncExitStack
 from inspect import isawaitable, signature, _empty
-from typing import TypeVar, Any, Iterable
+from typing import TypeVar, Any, Iterable, cast
 
 import anyio
 import structlog
-from anyio import Event, create_task_group, move_on_after
+from anyio import Event, create_task_group, fail_after, move_on_after
 from anyioutils import create_task, wait, FIRST_COMPLETED
 
 from ._context import Context, SharedValue, Value
@@ -69,6 +69,7 @@ def __init__(
         self._start_timeout = start_timeout
         self._stop_timeout = stop_timeout
         self._parent: Module | None = None
+        self._context = Context()
         self._prepared = Event()
         self._started = Event()
         self._stopped = Event()
@@ -256,24 +257,24 @@ def put(
             teardown_callback: A callback to call when the value is torn down.
             manage: Whether to use the (async) context manager of the value for its setup/teardown.
         """
-        if self.parent is None:
-            shared_value = self._context.put(
-                value,
-                types,
-                max_borrowers=max_borrowers,
-                manage=manage,
-                teardown_callback=teardown_callback,
-            )
-        else:
-            shared_value = self.parent._context.put(
+        value_id = id(value)
+        shared_value = self._context.put(
+            value,
+            types,
+            max_borrowers=max_borrowers,
+            manage=manage,
+            teardown_callback=teardown_callback,
+        )
+        self._published_values[value_id] = shared_value
+        if self.parent is not None:
+            self.parent._context.put(
                 value,
                 types,
                 max_borrowers=max_borrowers,
                 manage=manage,
                 teardown_callback=teardown_callback,
+                shared_value=shared_value,
             )
-        value_id = id(value)
-        self._published_values[value_id] = shared_value
         log.debug("Module added value", path=self.path, types=types)
 
     async def get(
@@ -290,10 +291,24 @@ async def get(
             The borrowed value.
         """
         log.debug("Module getting value", path=self.path, value_type=value_type)
+        tasks = [create_task(self._context.get(value_type), self._task_group)]
+        if self.parent is not None:
+            tasks.append(
+                create_task(self.parent._context.get(value_type), self._task_group)
+            )
         value_acquired = False
         try:
-            value = await self._context.get(value_type, timeout=timeout)
-            value_acquired = True
+            with fail_after(timeout):
+                done, pending = await wait(
+                    tasks, self._task_group, return_when=FIRST_COMPLETED
+                )
+                for task in pending:
+                    task.cancel()
+                for task in done:
+                    break
+                value = await task.wait()
+                value = cast(Value, value)
+                value_acquired = True
         finally:
             if not value_acquired:
                 log.critical(
@@ -441,27 +456,26 @@ async def _finish(self):
 
     async def _drop_and_wait_values(self):
         self.drop_all()
+        await self._context.aclose()
         self._stopped.set()
         log.debug("Module stopped", path=self.path)
 
     async def _prepare(self) -> None:
         log.debug("Preparing module", path=self.path)
-        async with Context() as self._context:
-            try:
-                async with create_task_group() as tg:
-                    for module in self._modules.values():
-                        module._task_group = tg
-                        module._phase = self._phase
-                        module._exceptions = self._exceptions
-                        tg.start_soon(module._prepare, name=f"{module.path} _prepare")
-                    tg.start_soon(
-                        self._prepare_and_done, name=f"{self.path} _prepare_and_done"
-                    )
-            except ExceptionGroup as exc:
-                self._exceptions.append(*exc.exceptions)
-                self._exit.set()
-                log.critical("Module failed while preparing", path=self.path)
-            await self._stopped.wait()
+        try:
+            async with create_task_group() as tg:
+                for module in self._modules.values():
+                    module._task_group = tg
+                    module._phase = self._phase
+                    module._exceptions = self._exceptions
+                    tg.start_soon(module._prepare, name=f"{module.path} _prepare")
+                tg.start_soon(
+                    self._prepare_and_done, name=f"{self.path} _prepare_and_done"
+                )
+        except ExceptionGroup as exc:
+            self._exceptions.append(*exc.exceptions)
+            self._exit.set()
+            log.critical("Module failed while preparing", path=self.path)
 
     async def _prepare_and_done(self) -> None:
         await self.prepare()

tests/test_context.py

@@ -54,6 +54,10 @@ async def test_nested_contexts():
             acquired_value_0 = await get(str)
             assert published_value_0 is acquired_value_0.unwrap()
             acquired_value_0.drop()
+            async with Context():
+                acquired_value_1 = await get(str)
+                assert published_value_0 is acquired_value_1.unwrap()
+                acquired_value_1.drop()
 
 
 async def test_context_cm():

tests/test_value.py

@@ -92,7 +92,10 @@ class Module2(Module):
         async def start(self):
             self.value2 = Value2()
             self.put(self.value2)
-            self.value0 = await self.get(Value0, timeout=0.1)
+            try:
+                self.value0 = await self.get(Value0, timeout=0.1)
+            except TimeoutError:
+                self.value0 = None
             self.value1 = await self.get(Value1, timeout=0.1)
 
     async with Module0("module0") as module0:
@@ -104,7 +107,7 @@ async def start(self):
     assert module0.value2 is None
     assert module1.value0 == module0.value0
     assert module1.value2 == module2.value2
-    assert module2.value0 == module0.value0
+    assert module2.value0 is None
     assert module2.value1 == module1.value1
 
 
@@ -251,8 +254,12 @@ async def stop(self):
     async with Module0("module0", stop_timeout=0.1) as module0:
         pass
 
-    assert len(module0.exceptions) == 1
-    assert str(module0.exceptions[0]) == "Module timed out while stopping: module0"
+    assert len(module0.exceptions) == 2
+    assert (
+        str(module0.exceptions[0])
+        == "Module timed out while stopping: module0.submodule0"
+    )
+    assert str(module0.exceptions[1]) == "Module timed out while stopping: module0"
 
 
 async def test_all_freed():