feat: guaranteed per-job child teardown (3.1.0)

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

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
 

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

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.