-
Notifications
You must be signed in to change notification settings - Fork 4.5k
fix(sandbox): single-flight cached dependency factories #3935
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,5 +1,6 @@ | ||
| from __future__ import annotations | ||
|
|
||
| import asyncio | ||
| import inspect | ||
| from collections.abc import Awaitable, Callable, Mapping | ||
| from dataclasses import dataclass | ||
|
|
@@ -73,7 +74,10 @@ class Dependencies: | |
| def __init__(self) -> None: | ||
| self._bindings: dict[DependencyKey, _Binding] = {} | ||
| self._cache: dict[DependencyKey, object] = {} | ||
| self._pending: dict[DependencyKey, asyncio.Task[object]] = {} | ||
| self._active_tasks: set[asyncio.Task[object]] = set() | ||
| self._owned_results: list[object] = [] | ||
| self._close_task: asyncio.Task[None] | None = None | ||
| self._closed = False | ||
|
|
||
| @classmethod | ||
|
|
@@ -144,6 +148,9 @@ def _bind( | |
| raise DependenciesBindingError(f"Dependency `{key}` is already bound") | ||
| self._bindings[key] = binding | ||
| self._cache.pop(key, None) | ||
| pending = self._pending.pop(key, None) | ||
| if pending is not None: | ||
| pending.cancel() | ||
|
|
||
| async def get(self, key: DependencyKey) -> object | None: | ||
| binding = self._bindings.get(key) | ||
|
|
@@ -173,24 +180,86 @@ async def _resolve(self, key: DependencyKey, binding: _Binding) -> object: | |
| return binding.value | ||
|
|
||
| assert isinstance(binding, _FactoryBinding) | ||
| if self._closed: | ||
| raise DependenciesError(f"Dependencies container is closed; cannot resolve `{key}`") | ||
| if binding.cache and key in self._cache: | ||
| return self._cache[key] | ||
|
|
||
| produced = binding.factory(self) | ||
| value = ( | ||
| await cast(Awaitable[object], produced) if inspect.isawaitable(produced) else produced | ||
| ) | ||
|
|
||
| if binding.cache: | ||
| self._cache[key] = value | ||
| if binding.owns_result: | ||
| self._owned_results.append(value) | ||
| return value | ||
| task = self._pending.get(key) | ||
| if task is None: | ||
| task = self._create_factory_task(key, binding) | ||
| self._pending[key] = task | ||
|
Comment on lines
+191
to
+192
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
On Python 3.12+ when the event loop uses Useful? React with 👍 / 👎. |
||
| return await asyncio.shield(task) | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
When an owned factory completes and another task calls AGENTS.md reference: AGENTS.md:L115-L115 Useful? React with 👍 / 👎. |
||
|
|
||
| task = self._create_factory_task(key, binding) | ||
| return await task | ||
|
|
||
| def _create_factory_task( | ||
| self, key: DependencyKey, binding: _FactoryBinding | ||
| ) -> asyncio.Task[object]: | ||
| task = asyncio.create_task(self._run_factory(key, binding)) | ||
| self._active_tasks.add(task) | ||
| task.add_done_callback(_consume_task_exception) | ||
| return task | ||
|
|
||
| async def _run_factory(self, key: DependencyKey, binding: _FactoryBinding) -> object: | ||
| try: | ||
| produced = binding.factory(self) | ||
| value = ( | ||
| await cast(Awaitable[object], produced) | ||
| if inspect.isawaitable(produced) | ||
| else produced | ||
| ) | ||
|
|
||
| if self._closed: | ||
| if binding.owns_result: | ||
| self._owned_results.append(value) | ||
| raise DependenciesError(f"Dependencies container closed while resolving `{key}`") | ||
|
|
||
| if self._bindings.get(key) is not binding: | ||
| if binding.owns_result: | ||
| self._owned_results.append(value) | ||
| raise DependenciesBindingError( | ||
| f"Dependency `{key}` was rebound while its factory was resolving" | ||
| ) | ||
|
|
||
| if binding.cache: | ||
| self._cache[key] = value | ||
| if binding.owns_result: | ||
| self._owned_results.append(value) | ||
| return value | ||
| except asyncio.CancelledError: | ||
| if self._closed: | ||
| raise DependenciesError( | ||
| f"Dependencies container closed while resolving `{key}`" | ||
| ) from None | ||
| if self._bindings.get(key) is not binding: | ||
| raise DependenciesBindingError( | ||
| f"Dependency `{key}` was rebound while its factory was resolving" | ||
| ) from None | ||
| raise | ||
| finally: | ||
| task = asyncio.current_task() | ||
| if task is not None: | ||
| self._active_tasks.discard(task) | ||
| if self._pending.get(key) is task: | ||
| self._pending.pop(key, None) | ||
|
|
||
| async def aclose(self) -> None: | ||
| if self._closed: | ||
| return | ||
| self._closed = True | ||
| task = self._close_task | ||
| if task is None: | ||
| self._closed = True | ||
| task = asyncio.create_task(self._close()) | ||
| self._close_task = task | ||
| await asyncio.shield(task) | ||
|
|
||
| async def _close(self) -> None: | ||
| active_tasks = tuple(self._active_tasks) | ||
| for task in active_tasks: | ||
| task.cancel() | ||
| if active_tasks: | ||
| await asyncio.gather(*active_tasks, return_exceptions=True) | ||
|
|
||
| seen_ids: set[int] = set() | ||
| for value in reversed(self._owned_results): | ||
|
|
@@ -199,3 +268,12 @@ async def aclose(self) -> None: | |
| continue | ||
| seen_ids.add(value_id) | ||
| await _close_best_effort(value) | ||
|
|
||
| self._pending.clear() | ||
| self._cache.clear() | ||
| self._owned_results.clear() | ||
|
|
||
|
|
||
| def _consume_task_exception(task: asyncio.Task[object]) -> None: | ||
| if not task.cancelled(): | ||
| task.exception() | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
When an overwrite races immediately after
require()creates the pending task but before_run_factoryreceives its first timeslice, this cancellation prevents the coroutine'stry/finallyfrom running at all. The stale waiter consequently receives a rawCancelledErrorinstead of the intendedDependenciesBindingError, and the completed task remains permanently in_active_tasks; uncached requests cancelled in the same pre-start window also leave retained tasks. Ensure active-set cleanup runs from a task completion callback and translate pre-start rebind cancellation consistently.AGENTS.md reference: AGENTS.md:L115-L115
Useful? React with 👍 / 👎.