diff --git a/README.md b/README.md index abb93c6..5f06c06 100644 --- a/README.md +++ b/README.md @@ -29,7 +29,7 @@ uv add modern-di-arq # or: pip install modern-di-arq ## Usage -`setup_di` seeds the root container into arq's `ctx` dict and wires four of arq's lifecycle hooks: `on_startup`/`on_shutdown` open and close the root container, and `on_job_start`/`on_job_end` build and close a `Scope.REQUEST` child container around each job. Decorate a task with `@inject` to resolve its `FromDI`-marked parameters from that per-job child. +`setup_di` seeds the root container into arq's `ctx` dict and wires four of arq's lifecycle hooks: `on_startup`/`on_shutdown` open and close the root container, and `on_job_start` builds a `Scope.REQUEST` child container per job. Decorate a task with `@inject` to resolve its `FromDI`-marked parameters from that per-job child — the child is open only while `@inject`-decorated task body/bodies are running, reference-counted so nested and concurrent (`asyncio.gather`) `@inject` calls over the same job share one open child and close it exactly once, guaranteeing teardown even if arq skips `on_job_end`. ```python import typing diff --git a/architecture/dependency-injection.md b/architecture/dependency-injection.md index 6770b0a..f7ceea3 100644 --- a/architecture/dependency-injection.md +++ b/architecture/dependency-injection.md @@ -13,16 +13,22 @@ version that keeps that shape. arq's `ctx` dict is the framework's state store — the worker builds `self.ctx`, then merges `{**self.ctx, **job_ctx}` into a fresh dict per job and passes that same dict through `on_job_start`, the task coroutine, and -`on_job_end`. Two named constants hand the container through it, so the -writer and every reader stay in provable agreement instead of relying on bare -string literals: +`on_job_end`. Three named constants hand the container (and its per-job +open/close bookkeeping) through it, so the writer and every reader stay in +provable agreement instead of relying on bare string literals: - `_ROOT_CONTAINER_KEY = "modern_di_container"` — the root container, seeded into the worker `ctx` by `setup_di` and read back by `fetch_di_container` and by `on_job_start` (to build the per-job child). - `_CHILD_CONTAINER_KEY = "modern_di_request_container"` — the per-job - `Scope.REQUEST` child, stashed on the per-job `ctx` by `on_job_start`, read - by `inject`'s wrapper, and closed by `on_job_end`. + `Scope.REQUEST` child, built (unopened) on the per-job `ctx` by + `on_job_start`, and opened/closed by `inject`'s wrapper(s) around the task + body/bodies they decorate (see "Per-job scope" below). `on_job_end` only + closes it as a safety net. +- `_CHILD_DEPTH_KEY = "modern_di_request_container_depth"` — the per-job + open depth counter `inject`'s wrapper reads and writes, so nested and + concurrent `@inject` calls sharing one `ctx` share one open child and close + it exactly once (see "Guaranteed teardown" below). ## `setup_di` and hook composition @@ -48,17 +54,19 @@ hooks, so a user hook always observes a live container: - `on_startup`: `container.open()` runs, *then* the existing hook (the user hook starts with an open root). -- `on_job_start`: the `Scope.REQUEST` child is built, *then* the existing - hook (the user hook can resolve from the child). -- `on_job_end`: the existing hook runs *first*, *then* the child is closed - (the user hook still has a live child). +- `on_job_start`: the `Scope.REQUEST` child is built (unopened), *then* the + existing hook. The child is **closed** at this point — a user + `on_job_start` hook cannot resolve from it; only an `@inject`-decorated + task can (see "Per-job scope"). +- `on_job_end`: the existing hook runs *first*, *then* the safety-net close. + By this point the owning `@inject` wrapper has normally already closed the + child, so both the existing hook and the safety net usually see it closed. - `on_shutdown`: the existing hook runs *first*, *then* `container.close_async()` (the user hook still has an open root). -`on_job_end` pops the child with `ctx.pop(_CHILD_CONTAINER_KEY, None)` and is -a no-op if absent — defensive only; arq always pairs `on_job_start`/`on_job_end` -around a job, including the error path (arq catches a task's exception itself -and still fires `on_job_end`). +`on_job_end` pops the child with `ctx.pop(_CHILD_CONTAINER_KEY, None)` and +closes it only `if child is not None and not child.closed` — a safety net, +not the primary teardown path (see "Per-job scope"). ## Root lifecycle @@ -76,19 +84,87 @@ and still fires `on_job_end`). ## Per-job scope `on_job_start` builds one `Scope.REQUEST` child per job — -`root.build_child_container(scope=Scope.REQUEST)` — opens it with a bare -`child.open()` (modern-di 3.x's mandatory-open lifecycle: resolving from an -unopened container raises `ContainerClosedError`), and stashes it under -`_CHILD_CONTAINER_KEY` on that job's `ctx`. The build happens in -`on_job_start` and the close happens in `on_job_end` — two different -hooks — so the child cannot be opened via a `with` block; it is opened with -a bare call instead. `on_job_end` closes it with -`close_async()`, unconditionally: because the compose ordering runs the -user's `on_job_end` first and the close second, the child is torn down -whether the task raised or returned, and it happens even for a task that was -never decorated with `@inject` (a child is built for every job regardless of -whether anything reads it back — see the non-goals note in the design; not -optimized). +`root.build_child_container(scope=Scope.REQUEST)` — and stashes it, +**unopened**, under `_CHILD_CONTAINER_KEY` on that job's `ctx`. It is not +opened here: a `Scope.REQUEST` child is only ever open while an `@inject` +wrapper's call is on the stack (see below), so building it unopened means +nothing has been resolved yet if a user `on_job_start` hook then raises — +there is nothing to leak. + +### Guaranteed teardown: the wrapper reference-counts open/close + +`inject`'s wrapper is the task body and runs inside arq's own caught +`try` (`run_job` catches a task's exception itself), so a `finally` there is +guaranteed to run — including when arq's `on_job_end` is later skipped, e.g. +it fails to serialize an unpicklable job result/exception after the task +already returned. The wrapper tracks a per-job **depth counter** under +`_CHILD_DEPTH_KEY` on `ctx`, opening on the 0→1 transition and closing on the +1→0 transition: + +```python +depth = ctx.get(_CHILD_DEPTH_KEY, 0) +if depth == 0: + child.open() +ctx[_CHILD_DEPTH_KEY] = depth + 1 +try: + ... # resolve markers, call func +finally: + ctx[_CHILD_DEPTH_KEY] -= 1 + if ctx[_CHILD_DEPTH_KEY] == 0: + await child.close_async() +``` + +The check/open/increment (and the decrement/close) run with no `await` in +between, so a single `@inject` call's entry and exit are each atomic from +asyncio's point of view — the event loop cannot interleave another task's +entry into the middle of one. This is what makes both **nested and +concurrent `@inject`** safe: + +- **Nested**: an outer `@inject` task that awaits an inner + `@inject`-decorated function passing the same `ctx` finds the child already + open on the inner call (`depth=1`), increments to `depth=2`, and its own + `finally` only decrements back to `1` — it does not close. Only the outer + call's `finally`, decrementing `1→0`, closes the child. +- **Concurrent fan-out**: a job coroutine that is *not* itself `@inject` can + `asyncio.gather` two or more `@inject`-decorated helpers over the same + `ctx`. Whichever sibling's task step runs first opens the child (`0→1`); + every other sibling that enters before the child is closed again just + increments the depth and reuses it. Each sibling's own `finally` + decrements; the child is closed exactly once, when the *last* sibling to + exit brings the depth back to `0` — regardless of which sibling happened to + open it or which finishes first. Earlier boolean ownership (`owns = + child.closed`) got this wrong under fan-out: whichever sibling finished + first believed it alone owned the child and closed it in its `finally`, + even while a slower sibling was still mid-body — a use-after-finalize on + the REQUEST-scoped resource the slower sibling was still holding. + +A task with no `FromDI` parameter, or one that isn't decorated with +`@inject`, never opens the child at all — it stays closed for that job's +entire span, and `on_job_end`'s safety net finds nothing to close. + +### `on_job_end` is a safety net, not the primary teardown + +`on_job_end` pops the child and the depth counter +(`ctx.pop(_CHILD_DEPTH_KEY, None)`, for cleanliness — the depth is job-scoped +and the `ctx` dict does not outlive the job) and closes the child **only if +still open** (`if child is not None and not child.closed`). In the ordinary +case this is a no-op: the last `@inject` wrapper to exit already closed the +child before the task returned, or no `@inject` task ever opened it. It only +does real work if something left the child open across the task boundary +(e.g. code that opens and resolves from the child directly instead of via +`@inject`). `Container.open()`/`close_async()` are both idempotent, so this +composes safely with the wrapper regardless of ordering. + +### Contract change from `on_job_start` unconditionally opening the child + +Earlier versions opened the child in `on_job_start` and closed it +unconditionally in `on_job_end`, so it was live across the whole +`on_job_start`→`on_job_end` span — a user-supplied `on_job_start` or +`on_job_end` hook could resolve straight from +`ctx[_CHILD_CONTAINER_KEY]`. That window is gone: the child is now open only +inside an `@inject` wrapper's call. `_CHILD_CONTAINER_KEY` was always a +private key, and the supported way to resolve a per-job dependency is +`@inject`, not reading the hook's `ctx` directly. ## Resolution diff --git a/modern_di_arq/main.py b/modern_di_arq/main.py index 092edc8..b42e448 100644 --- a/modern_di_arq/main.py +++ b/modern_di_arq/main.py @@ -13,6 +13,7 @@ _ROOT_CONTAINER_KEY = "modern_di_container" _CHILD_CONTAINER_KEY = "modern_di_request_container" +_CHILD_DEPTH_KEY = "modern_di_request_container_depth" _WRAPPED_MARKER = "__modern_di_wrapped__" @@ -54,8 +55,9 @@ async def on_shutdown(ctx: dict[str, typing.Any]) -> None: def _wrap_job_start(existing: _Hook | None) -> _Hook: async def on_job_start(ctx: dict[str, typing.Any]) -> None: root = typing.cast(Container, ctx[_ROOT_CONTAINER_KEY]) + # Built unopened: `@inject`'s wrapper(s) refcount open/close (see `inject`), + # so the child stays closed here even if a user hook resolves-then-raises below. child = root.build_child_container(scope=Scope.REQUEST) - child.open() ctx[_CHILD_CONTAINER_KEY] = child if existing is not None: await existing(ctx) @@ -68,8 +70,13 @@ async def on_job_end(ctx: dict[str, typing.Any]) -> None: if existing is not None: await existing(ctx) child = ctx.pop(_CHILD_CONTAINER_KEY, None) - if child is not None: - await child.close_async() # never leak the per-job child, even on the error path + ctx.pop(_CHILD_DEPTH_KEY, None) + # Safety net only: normally the owning `@inject` wrapper(s) already closed the + # child in their own `finally`, or it was never opened (no `@inject` task) — + # both are no-ops here. Only closes if still open (e.g. a non-`@inject` task + # left it open, or arq itself raised between the task and this hook). + if child is not None and not child.closed: + await child.close_async() return on_job_end @@ -79,10 +86,14 @@ def setup_di(worker_settings: typing.Any, container: Container) -> Container: # Seeds the root container into the worker ``ctx`` (arq's state store) and wraps the worker's lifecycle hooks: ``on_startup``/``on_shutdown`` open and - close the root; ``on_job_start``/``on_job_end`` build and close a - ``Scope.REQUEST`` child per job. Any hook the user already set still runs. - Accepts a class/object ``worker_settings`` (attribute access) or a ``dict`` - (item access). Returns *container*. + close the root; ``on_job_start`` builds an unopened ``Scope.REQUEST`` child + per job, opened and closed by ``@inject``-decorated task(s) around their + own bodies — reference-counted, so nested and concurrent (``asyncio.gather``) + ``@inject`` calls over the same job share one open child, closed exactly + once by the last to exit (``on_job_end`` only closes it as a safety net). + Any hook the user already set still runs. Accepts a class/object + ``worker_settings`` (attribute access) or a ``dict`` (item access). + Returns *container*. Raises: TypeError: *worker_settings* was already wired by a prior ``setup_di`` call. @@ -154,9 +165,24 @@ def inject(func: typing.Callable[..., typing.Awaitable[T]]) -> typing.Callable[. async def wrapper(*args: typing.Any, **kwargs: typing.Any) -> T: # noqa: ANN401 ctx = typing.cast("dict[str, typing.Any]", args[0]) child = typing.cast(Container, ctx[_CHILD_CONTAINER_KEY]) - resolved = integrations.resolve_markers(child, di_params) - bound = visible_signature.bind(*args, **kwargs) - bound.apply_defaults() - return await func(**bound.arguments, **resolved) + # Reference-count opens so nested AND concurrent (@inject fan-out via gather) + # share one open child and close it exactly once, when the LAST @inject body + # exits. The check/open/increment run without an await, so asyncio cannot + # interleave them — the count is consistent under concurrency. + depth = ctx.get(_CHILD_DEPTH_KEY, 0) + if depth == 0: + child.open() + ctx[_CHILD_DEPTH_KEY] = depth + 1 + try: + resolved = integrations.resolve_markers(child, di_params) + bound = visible_signature.bind(*args, **kwargs) + bound.apply_defaults() + return await func(**bound.arguments, **resolved) + finally: + # Guarantees teardown even if arq skips on_job_end (e.g. it fails to + # serialize an unpicklable job result/exception after the task runs). + ctx[_CHILD_DEPTH_KEY] -= 1 + if ctx[_CHILD_DEPTH_KEY] == 0: + await child.close_async() return wrapper diff --git a/planning/changes/2026-07-21.01-guaranteed-child-teardown.md b/planning/changes/2026-07-21.01-guaranteed-child-teardown.md new file mode 100644 index 0000000..3a8c5ba --- /dev/null +++ b/planning/changes/2026-07-21.01-guaranteed-child-teardown.md @@ -0,0 +1,96 @@ +--- +summary: inject's wrapper reference-counts opening and closing the per-job Scope.REQUEST child around the task body/bodies it decorates, so nested AND concurrent (asyncio.gather) @inject calls over the same ctx share one open child closed exactly once, guaranteeing teardown even when arq skips on_job_end; on_job_start builds the child unopened, on_job_end becomes a safety net. +--- + +# Change: Guaranteed per-job child teardown + +**Lane:** lightweight — 1 file changed in `modern_di_arq/`, no new file, no +public-API change, tests added/updated in `tests/test_jobs.py` and +`tests/test_lifespan.py`. + +## Goal + +`on_job_start` built and opened the per-job `Scope.REQUEST` child; `on_job_end` +closed it unconditionally. arq's `run_job` does not call `on_job_end` inside a +`finally`, so two narrow paths skipped it and leaked the opened child's +REQUEST-scope finalizers: (1) arq fails to serialize an unpicklable job +result/exception after the task runs but before `on_job_end`; (2) a user +`on_job_start` hook resolves-then-raises. Both are pre-existing, narrow, and +leak a per-job resource rather than crash anything — but avoidable. + +## Approach + +Move open/close into `inject`'s wrapper, which *is* the task body and runs +inside arq's own caught `try`, so a `finally` there is guaranteed to run even +when `on_job_end` is later skipped: + +- `on_job_start` builds the child **unopened** — nothing resolved yet if a + user hook then raises, so nothing to leak. +- `inject`'s wrapper reference-counts the open depth on a `ctx`-scoped + counter (`_CHILD_DEPTH_KEY`): it opens the child on the 0→1 transition and + closes it in a `finally` on the 1→0 transition. This is what makes both + **nested `@inject`** (an inner call sees `depth>=1`, increments, and its + own `finally` doesn't bring it back to 0) and **concurrent `@inject` + fan-out** (a non-`@inject` job coroutine `asyncio.gather`s two or more + `@inject`-decorated helpers over the same `ctx`) safe — the child is closed + exactly once, by whichever call is *last* to exit, regardless of which one + opened it or finished first. + An earlier boolean-ownership design (`owns = child.closed`) handled nesting + but not concurrent fan-out: an adversarial review found that under + `asyncio.gather`, the first sibling to finish believed it alone owned the + child and closed it in its `finally` — even while a slower sibling fanned + out alongside it was still resolving/using a REQUEST-scoped resource, a + use-after-finalize. The refcount fixes this. +- `on_job_end` becomes a safety net: it closes the child only if it's still + open, which is normally a no-op (the last wrapper to exit already closed + it, or no `@inject` task ever opened it); it also pops the depth counter + for cleanliness. + +Contract change: the per-job child is now open only within an `@inject` +wrapper's span, not across the whole `on_job_start`→`on_job_end` span, so a +user-supplied `on_job_start`/`on_job_end` hook can no longer resolve from it +directly — `_CHILD_CONTAINER_KEY` was always private; `@inject` is the +supported path. See `architecture/dependency-injection.md`'s "Per-job scope" +section, rewritten in this change. + +## Files + +- `modern_di_arq/main.py` — `_wrap_job_start` drops `child.open()`; + `inject`'s `wrapper` gains the refcounted open/close around a `try`/ + `finally`, keyed by the new `_CHILD_DEPTH_KEY`; `_wrap_job_end` pops the + depth key and only closes `if child is not None and not child.closed`. +- `tests/test_jobs.py` — added: guaranteed-teardown-without-`on_job_end` + (return and raise paths), nested-`@inject` (inner doesn't close a child it + doesn't own), the `on_job_end` safety-net path (still-open child gets + closed), and `test_concurrent_inject_fanout_shares_one_child` (two + `@inject` siblings fanned out via `asyncio.gather` over the same `ctx`; + the slow sibling must not observe its REQUEST resource finalized while the + fast sibling exits first — RED under the boolean-ownership version, GREEN + under the refcount). +- `tests/test_lifespan.py` — `resolve_job` now goes through `@inject` instead + of resolving directly from the child (the child is closed outside an + `@inject` span under the new contract); `test_setup_di_composes_with_user_hooks` + updated for the child being closed at `on_job_start`/`on_job_end` time. +- `architecture/dependency-injection.md`, `README.md` — promoted the + refcounted lifecycle (nested + concurrent) and the contract change. +- `planning/releases/3.1.0.md` — release note updated: both nested and + concurrent (`asyncio.gather`) `@inject` are refcount-safe, not just nested. + +## Verification + +- [x] Failing test first — new tests in `tests/test_jobs.py` fail against the + old implementation (child not closed without `on_job_end`; nested `@inject` + closes early); `test_setup_di_composes_with_user_hooks` fails on the + updated `closed is True` assertions. +- [x] Apply the change — open/close in `inject`'s wrapper; drop + `on_job_start`'s open; guard `on_job_end`'s close. +- [x] Tests pass — `just test`. +- [x] `just test-ci` — full suite green, 100% line coverage. +- [x] `just lint-ci` — clean. +- [x] Adversarial review found the boolean-ownership version unsafe under + concurrent `@inject` fan-out (`asyncio.gather` over the same `ctx`) — a + faster sibling closes the child while a slower sibling is still using a + REQUEST-scoped resource. Reproduced with a new failing test + (`test_concurrent_inject_fanout_shares_one_child`, confirmed RED), fixed by + replacing the boolean with a per-job reference count (open on 0→1, close on + 1→0); confirmed GREEN, `just test-ci` 100% coverage, `just lint-ci` clean. diff --git a/planning/releases/3.1.0.md b/planning/releases/3.1.0.md new file mode 100644 index 0000000..feaa235 --- /dev/null +++ b/planning/releases/3.1.0.md @@ -0,0 +1,50 @@ +# modern-di-arq 3.1.0 — guaranteed per-job child teardown + +Closes two narrow leak windows around the per-job `Scope.REQUEST` child +container. **No public API change** — `FromDI`, `fetch_di_container`, +`inject`, and `setup_di` keep their signatures. One contract change: see +below. + +## Fix + +- **The per-job child is now guaranteed to close, even if arq skips + `on_job_end`.** Previously `on_job_start` opened the child and `on_job_end` + closed it — but arq's `run_job` doesn't call `on_job_end` inside a + `finally`, so two narrow paths leaked the child's REQUEST-scope + finalizers: (1) arq fails to serialize an unpicklable job result/exception + after the task ran but before `on_job_end`; (2) a user `on_job_start` hook + resolves-then-raises. `@inject`'s wrapper now opens and closes the child + around the task body it decorates, inside arq's own caught `try`, so its + `finally` runs even when `on_job_end` is later skipped. +- **Both nested and concurrent `@inject` calls sharing a `ctx` are safe.** + The wrapper reference-counts the child's open depth on `ctx` instead of a + boolean "did I open it" flag: an outer `@inject` task awaiting an inner + `@inject`-decorated function over the same `ctx` shares one open child, and + so does a job coroutine that fans two or more `@inject`-decorated helpers + out concurrently via `asyncio.gather` over the same `ctx` — whichever + sibling's task step runs first opens it, every other sibling that joins in + before it's closed again just reuses it, and it's closed exactly once, by + whichever sibling is *last* to exit. An earlier boolean-ownership version + of this fix got the concurrent case wrong: the first sibling to finish + believed it alone owned the child and closed it, even while a slower + sibling fanned out alongside it was still using a REQUEST-scoped resource — + a use-after-finalize. The reference count closes this window. + +## Contract change + +The per-job child is now open only for the duration of an `@inject`-decorated +task's body, not across the whole `on_job_start`→`on_job_end` span. A +user-supplied `on_job_start` or `on_job_end` hook can no longer resolve from +the per-job child — it was reached via the private `ctx["modern_di_request_container"]` +key, never a supported path. `@inject` remains the supported way to resolve a +per-job dependency; `on_job_end` is now a safety net only, closing the child +if (and only if) it is somehow still open when it runs. + +## Downstream + +No action needed unless a hook reads `ctx["modern_di_request_container"]` +directly — switch that resolution to an `@inject`-decorated task instead. + +## Internals + +- 100% line coverage; `ruff`, `ty` clean. diff --git a/tests/test_jobs.py b/tests/test_jobs.py index cfd127b..d079b10 100644 --- a/tests/test_jobs.py +++ b/tests/test_jobs.py @@ -1,9 +1,11 @@ +import asyncio import typing import pytest from modern_di import Container, Group, Scope, providers from modern_di_arq import FromDI, inject, setup_di +from modern_di_arq.main import _CHILD_CONTAINER_KEY, _ROOT_CONTAINER_KEY, _wrap_job_end, _wrap_job_start from tests.conftest import run_burst_worker from tests.dependencies import ( AppResource, @@ -124,6 +126,152 @@ async def test_inject_closes_child_on_task_error(arq_redis) -> None: # noqa: AN assert boom_teardowns == ["closed"] +async def test_wrapper_guarantees_close_without_on_job_end() -> None: + """The wrapper's own `finally` closes the child, so a skipped on_job_end leaks nothing.""" + results.clear() + request_teardowns.clear() + container = Container(groups=[Dependencies], validate=True) + container.open() + ctx: dict[str, typing.Any] = {_ROOT_CONTAINER_KEY: container} + await _wrap_job_start(None)(ctx) + child = ctx[_CHILD_CONTAINER_KEY] + assert child.closed is True # built unopened by on_job_start + + await resolves_app_and_request(ctx, 7) # invoke the @inject wrapper directly; on_job_end never runs + + assert results == {"app_ok": True, "request_ok": True, "x": 7, "linked": True} + assert child.closed is True # wrapper opened it, then closed it in its own finally + assert request_teardowns == ["request-closed"] + + +async def test_wrapper_closes_child_when_task_raises_without_on_job_end() -> None: + """The wrapper's `finally` closes the child even when the task raises, without on_job_end.""" + boom_teardowns.clear() + container = Container(groups=[Boom], validate=True) + container.open() + ctx: dict[str, typing.Any] = {_ROOT_CONTAINER_KEY: container} + await _wrap_job_start(None)(ctx) + child = ctx[_CHILD_CONTAINER_KEY] + + with pytest.raises(ValueError, match="boom"): + await raiser(ctx) + + assert child.closed is True + assert boom_teardowns == ["closed"] + + +nested_calls: dict[str, typing.Any] = {} + + +@inject +async def inner_task( + ctx: dict[str, typing.Any], + app_instance: typing.Annotated[AppResource, FromDI(AppResource)], +) -> AppResource: + nested_calls["inner_child_closed_during_call"] = ctx[_CHILD_CONTAINER_KEY].closed + return app_instance + + +@inject +async def outer_task( + ctx: dict[str, typing.Any], + request_instance: typing.Annotated[RequestResource, FromDI(Dependencies.request_factory)], +) -> None: + nested_calls["request_ok"] = isinstance(request_instance, RequestResource) + inner_result = await inner_task(ctx) + nested_calls["inner_ok"] = isinstance(inner_result, AppResource) + nested_calls["child_closed_after_inner_returns"] = ctx[_CHILD_CONTAINER_KEY].closed + + +async def test_nested_inject_inner_does_not_close_shared_child() -> None: + """Only the outer (owning) wrapper closes the shared child; the inner call must not.""" + nested_calls.clear() + request_teardowns.clear() + container = Container(groups=[Dependencies], validate=True) + container.open() + ctx: dict[str, typing.Any] = {_ROOT_CONTAINER_KEY: container} + await _wrap_job_start(None)(ctx) + child = ctx[_CHILD_CONTAINER_KEY] + + await outer_task(ctx) + + assert nested_calls["request_ok"] is True + assert nested_calls["inner_ok"] is True + assert nested_calls["inner_child_closed_during_call"] is False + assert nested_calls["child_closed_after_inner_returns"] is False # inner did not close it + assert child.closed is True # outer, the owner, closed it once after returning + assert request_teardowns == ["request-closed"] # finalizer ran exactly once + + +concurrent_calls: dict[str, typing.Any] = {} + + +@inject +async def fast_sibling( + ctx: dict[str, typing.Any], # noqa: ARG001 + request_instance: typing.Annotated[RequestResource, FromDI(Dependencies.request_factory)], +) -> None: + concurrent_calls["fast_request_instance"] = request_instance + await asyncio.sleep(0) # yield once so the slow sibling can open concurrently + concurrent_calls["fast_alive_after_yield"] = not request_teardowns + + +@inject +async def slow_sibling( + ctx: dict[str, typing.Any], # noqa: ARG001 + request_instance: typing.Annotated[RequestResource, FromDI(Dependencies.request_factory)], +) -> None: + concurrent_calls["slow_request_instance"] = request_instance + await asyncio.sleep(0.05) # outlives fast_sibling + concurrent_calls["slow_alive_after_yield"] = not request_teardowns + + +async def fanout_job(ctx: dict[str, typing.Any]) -> None: + """Not @inject itself: fans out two @inject siblings over the same ctx.""" + await asyncio.gather(fast_sibling(ctx), slow_sibling(ctx)) + + +async def test_concurrent_inject_fanout_shares_one_child() -> None: + """Two `@inject` siblings fanned out via `asyncio.gather` over the same ctx. + + They must share one open child, refcounted, and close it exactly once — only + when the last sibling exits. + """ + concurrent_calls.clear() + request_teardowns.clear() + container = Container(groups=[Dependencies], validate=True) + container.open() + ctx: dict[str, typing.Any] = {_ROOT_CONTAINER_KEY: container} + await _wrap_job_start(None)(ctx) + child = ctx[_CHILD_CONTAINER_KEY] + + await fanout_job(ctx) + + assert concurrent_calls["fast_request_instance"] is concurrent_calls["slow_request_instance"] + assert concurrent_calls["fast_alive_after_yield"] is True + assert concurrent_calls["slow_alive_after_yield"] is True # not finalized while slow was still running + assert child.closed is True # closed after the last (slow) sibling exits + assert request_teardowns == ["request-closed"] # finalizer fired exactly once + + +async def test_on_job_end_safety_net_closes_a_still_open_child() -> None: + """on_job_end is a safety net: it closes the child only when some non-@inject path left it open.""" + request_teardowns.clear() + container = Container(groups=[Dependencies], validate=True) + container.open() + ctx: dict[str, typing.Any] = {_ROOT_CONTAINER_KEY: container} + await _wrap_job_start(None)(ctx) + child = ctx[_CHILD_CONTAINER_KEY] + child.open() # simulate a non-@inject task resolving directly, leaving the child open + child.resolve_dependency(Dependencies.request_factory) + + await _wrap_job_end(None)(ctx) + + assert child.closed is True + assert _CHILD_CONTAINER_KEY not in ctx + assert request_teardowns == ["request-closed"] + + def test_inject_rejects_var_positional_with_fromdi() -> None: async def bad_task( ctx: dict[str, typing.Any], # noqa: ARG001 diff --git a/tests/test_lifespan.py b/tests/test_lifespan.py index ed89277..b092c45 100644 --- a/tests/test_lifespan.py +++ b/tests/test_lifespan.py @@ -7,14 +7,17 @@ from modern_di import Container, exceptions import modern_di_arq -from modern_di_arq import fetch_di_container, setup_di +from modern_di_arq import FromDI, fetch_di_container, inject, setup_di from tests.conftest import REDIS_URL, run_burst_worker -from tests.dependencies import Dependencies, app_teardowns, request_teardowns +from tests.dependencies import Dependencies, RequestResource, app_teardowns, request_teardowns -async def resolve_job(ctx: dict[str, typing.Any]) -> None: - child = ctx["modern_di_request_container"] - child.resolve_dependency(Dependencies.request_factory) # REQUEST-scoped (child) +@inject +async def resolve_job( + ctx: dict[str, typing.Any], + _dep: typing.Annotated[RequestResource, FromDI(Dependencies.request_factory)], +) -> None: + pass # the @inject wrapper resolves the REQUEST-scoped dependency; nothing else to do def make_settings(container: Container) -> type: @@ -111,7 +114,9 @@ async def user_startup(ctx: dict[str, typing.Any]) -> None: async def user_job_start(ctx: dict[str, typing.Any]) -> None: calls.append("user_job_start") - assert ctx["modern_di_request_container"].closed is False # our on_job_start already built the child + # on_job_start builds the child but no longer opens it — it opens only within an + # @inject wrapper's ownership span, so the hook sees it closed + assert ctx["modern_di_request_container"].closed is True async def user_shutdown(ctx: dict[str, typing.Any]) -> None: calls.append("user_shutdown") @@ -119,8 +124,9 @@ async def user_shutdown(ctx: dict[str, typing.Any]) -> None: async def user_job_end(ctx: dict[str, typing.Any]) -> None: calls.append("user_job_end") - # our on_job_end closes the child after the user's hook, so it's still live here - assert ctx["modern_di_request_container"].closed is False + # the @inject wrapper already opened and closed the child around the task body, + # so by the time on_job_end's safety net runs, it's already closed + assert ctx["modern_di_request_container"].closed is True class WorkerSettings: functions: typing.ClassVar[list] = [resolve_job]