diff --git a/CONTEXT.md b/CONTEXT.md index 498833a..6759955 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -18,8 +18,9 @@ package's own. **Root container**: The `Container` that `setup_di` attaches to `broker.state`: the one a caller constructs and hands -in, and the one `fetch_di_container` returns. By default the worker's startup and shutdown events -open and close it; with `setup_di(..., manage_lifespan=False)` the caller does. +in, and the one `fetch_di_container` returns. By default the worker's startup event opens it, and +shutdown closes it after all `WORKER_SHUTDOWN` handlers and middleware `shutdown()` calls; with +`setup_di(..., manage_lifespan=False)` the caller does. _Avoid_: APP-scope container — that names a scope where what matters is the position at the top of the tree, and `Scope.APP` is `modern-di`'s word for a band in the hierarchy, not for this object. diff --git a/README.md b/README.md index b01241b..930b75e 100644 --- a/README.md +++ b/README.md @@ -31,7 +31,7 @@ uv add modern-di-taskiq # or: pip install modern-di-taskiq ## Usage -`setup_di` stores the container on `broker.state` and registers `WORKER_STARTUP`/`WORKER_SHUTDOWN` handlers that open/close it. `FromDI` resolves a provider (or type) into a task parameter through a taskiq dependency, which builds a `Scope.REQUEST` child container for the task and closes it when the task finishes. No per-task decorator is needed. +`setup_di` stores the container on `broker.state` and opens it on `WORKER_STARTUP`. At shutdown it closes the container after every `WORKER_SHUTDOWN` handler and every middleware `shutdown()`, so your shutdown code can still resolve `Scope.APP` dependencies and their finalizers run last. `FromDI` resolves a provider (or type) into a task parameter through a taskiq dependency, which builds a `Scope.REQUEST` child container for the task and closes it when the task finishes. No per-task decorator is needed. ```python import typing diff --git a/modern_di_taskiq/main.py b/modern_di_taskiq/main.py index 2312a3e..83bc4b1 100644 --- a/modern_di_taskiq/main.py +++ b/modern_di_taskiq/main.py @@ -3,7 +3,7 @@ import taskiq from modern_di import Container, Scope, integrations, providers -from taskiq import AsyncBroker, Context, TaskiqDepends, TaskiqEvents, TaskiqState +from taskiq import AsyncBroker, Context, TaskiqDepends, TaskiqEvents, TaskiqMiddleware, TaskiqState taskiq_message_provider = providers.ContextProvider(taskiq.TaskiqMessage, scope=Scope.REQUEST) @@ -16,17 +16,46 @@ _CONNECTION_PROVIDERS = (taskiq_message_provider,) +class _RootCloser(TaskiqMiddleware): + def __init__(self, container: Container) -> None: + super().__init__() + self.container = container + self.broker_runs_middlewares = False + + def startup(self) -> None: + self.broker_runs_middlewares = True + + async def shutdown(self) -> None: + if self.broker_runs_middlewares: + await self.container.close_async() + + +def _move_to_end(items: list[typing.Any], item: object) -> None: + if item in items: + items.remove(item) + items.append(item) + + def setup_di(broker: AsyncBroker, container: Container, *, manage_lifespan: bool = True) -> Container: setattr(broker.state, _ROOT_CONTAINER_ATTR, container) container.add_providers(*_CONNECTION_PROVIDERS) if not manage_lifespan: return container + closer = _RootCloser(container) + def _open_root(_state: TaskiqState) -> None: container.open() + closer.broker_runs_middlewares = False + if closer not in broker.middlewares: + broker.add_middlewares(closer) + _move_to_end(broker.event_handlers[TaskiqEvents.WORKER_SHUTDOWN], _close_root) async def _close_root(_state: TaskiqState) -> None: - await container.close_async() + if closer.broker_runs_middlewares and closer in broker.middlewares: + _move_to_end(broker.middlewares, closer) + else: + await container.close_async() broker.add_event_handler(TaskiqEvents.WORKER_STARTUP, _open_root) broker.add_event_handler(TaskiqEvents.WORKER_SHUTDOWN, _close_root) diff --git a/pyproject.toml b/pyproject.toml index b95f998..7c44fc3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,7 +17,7 @@ classifiers = [ "Typing :: Typed", "Topic :: Software Development :: Libraries", ] -dependencies = ["taskiq>=0.11,<0.13", "modern-di>=4,<5"] +dependencies = ["taskiq>=0.11.8,<0.13", "modern-di>=4,<5"] version = "0" [project.urls] diff --git a/tests/test_shutdown_order.py b/tests/test_shutdown_order.py new file mode 100644 index 0000000..31870d8 --- /dev/null +++ b/tests/test_shutdown_order.py @@ -0,0 +1,148 @@ +import collections.abc + +import pytest +from modern_di import Container, Group, Scope, providers +from taskiq import AsyncBroker, BrokerMessage, InMemoryBroker, TaskiqEvents, TaskiqMiddleware, TaskiqState + +from modern_di_taskiq import setup_di + + +class Pool: + pass + + +class StubBroker(AsyncBroker): + def __init__(self) -> None: + super().__init__() + self.is_worker_process = True + + async def kick(self, message: BrokerMessage) -> None: # pragma: no cover + pass + + async def listen(self) -> collections.abc.AsyncGenerator[bytes]: # pragma: no cover + yield b"" + + +def _container(events: list[str]) -> Container: + class Resources(Group): + pool = providers.Factory( + Pool, scope=Scope.APP, cache=providers.CacheSettings(finalizer=lambda _: events.append("finalized")) + ) + + return Container(groups=[Resources]) + + +def _use(container: Container, events: list[str], label: str) -> None: + container.resolve(Pool) + events.append(label) + + +class UserMiddleware(TaskiqMiddleware): + def __init__(self, container: Container, events: list[str], label: str) -> None: + super().__init__() + self.container = container + self.events = events + self.label = label + + async def shutdown(self) -> None: + _use(self.container, self.events, self.label) + + +def _register_user_code(broker: AsyncBroker, container: Container, events: list[str], label: str) -> None: + async def hook(_state: TaskiqState) -> None: + _use(container, events, f"{label} hook") + + broker.add_event_handler(TaskiqEvents.WORKER_SHUTDOWN, hook) + broker.add_middlewares(UserMiddleware(container, events, f"{label} middleware")) + + +async def test_root_closes_after_all_shutdown_handlers_and_middlewares() -> None: + events: list[str] = [] + broker = StubBroker() + container = _container(events) + _register_user_code(broker, container, events, "before") + setup_di(broker, container) + _register_user_code(broker, container, events, "after") + + await broker.startup() + container.resolve(Pool) + await broker.shutdown() + + assert events == ["before hook", "after hook", "before middleware", "after middleware", "finalized"] + assert container.closed is True + + +async def test_root_closes_after_a_middleware_added_after_startup() -> None: + events: list[str] = [] + broker = StubBroker() + container = setup_di(broker, _container(events)) + + await broker.startup() + broker.add_middlewares(UserMiddleware(container, events, "late middleware")) + container.resolve(Pool) + await broker.shutdown() + + assert events == ["late middleware", "finalized"] + + +async def test_root_closes_once_per_cycle_without_duplicate_middlewares() -> None: + events: list[str] = [] + broker = StubBroker() + container = setup_di(broker, _container(events)) + _register_user_code(broker, container, events, "user") + + for _ in range(2): + await broker.startup() + middleware_count = len(broker.middlewares) + container.resolve(Pool) + await broker.shutdown() + assert len(broker.middlewares) == middleware_count + assert container.closed is True + + assert events == ["user hook", "user middleware", "finalized"] * 2 + + +async def test_root_closes_once_when_the_broker_skips_middleware_shutdown() -> None: + events: list[str] = [] + broker = InMemoryBroker() + container = setup_di(broker, _container(events)) + + async def hook(_state: TaskiqState) -> None: + _use(container, events, "hook") + + broker.add_event_handler(TaskiqEvents.WORKER_SHUTDOWN, hook) + + for _ in range(2): + await broker.startup() + container.resolve(Pool) + await broker.shutdown() + assert container.closed is True + + assert events == ["hook", "finalized"] * 2 + + +async def test_root_closes_when_shutdown_runs_without_startup() -> None: + events: list[str] = [] + broker = StubBroker() + container = setup_di(broker, _container(events)) + container.resolve(Pool) + + await broker.shutdown() + + assert container.closed is True + assert events == ["finalized"] + + +@pytest.mark.parametrize("broker_type", [StubBroker, InMemoryBroker]) +async def test_manage_lifespan_false_adds_no_middleware_and_keeps_the_root_open(broker_type: type[AsyncBroker]) -> None: + events: list[str] = [] + broker = broker_type() + container = setup_di(broker, _container(events), manage_lifespan=False) + container.resolve(Pool) + + await broker.startup() + assert broker.middlewares == [] + await broker.shutdown() + + assert container.closed is False + assert events == [] diff --git a/tests/test_taskiq_di.py b/tests/test_taskiq_di.py index 77eb927..5143bc6 100644 --- a/tests/test_taskiq_di.py +++ b/tests/test_taskiq_di.py @@ -10,7 +10,7 @@ async def test_resolves_app_request_and_context(broker: InMemoryBroker) -> None: # Explicit startup/shutdown (not `async with broker`) so the suite runs on - # the whole taskiq>=0.11 floor — InMemoryBroker gained async-context-manager + # the whole taskiq>=0.11.8 floor — InMemoryBroker gained async-context-manager # support only in 0.12.3, and the integration itself needs neither. @broker.task(task_name="my_task") async def my_task(