diff --git a/README.md b/README.md index 862d5c5..9b4d196 100644 --- a/README.md +++ b/README.md @@ -31,7 +31,7 @@ uv add modern-di-celery # or: pip install modern-di-celery ## Usage -`setup_di` registers the container for the app and connects handlers that open and close it on `worker_process_init`/`worker_process_shutdown` (once per forked process under the prefork and solo pools) and on `worker_init`/`worker_shutdown` (once in the main worker process, the only signals the threads, gevent, and eventlet pools send). `@inject` builds a `Scope.REQUEST` child container per task call and resolves the `FromDI`-marked parameters from it; alternatively set `task_cls=DITask` to inject every task without a per-task decorator. +`setup_di` registers the container for the app and connects handlers that open and close it on `worker_process_init`/`worker_process_shutdown` (once per forked process under the prefork and solo pools) and on `worker_init`/`worker_shutdown` (once in the main worker process, the only signals the threads, gevent, and eventlet pools send). The close runs after every other receiver of the shutdown signal, including ones connected after `setup_di`, so your shutdown code can still resolve `Scope.APP` dependencies and their finalizers run last. `@inject` builds a `Scope.REQUEST` child container per task call and resolves the `FromDI`-marked parameters from it; alternatively set `task_cls=DITask` to inject every task without a per-task decorator. ```python import typing diff --git a/modern_di_celery/main.py b/modern_di_celery/main.py index 3f73422..72827bc 100644 --- a/modern_di_celery/main.py +++ b/modern_di_celery/main.py @@ -4,6 +4,7 @@ from celery import Celery, Task, current_app, signals from celery.local import Proxy +from celery.utils.dispatch import Signal from celery.utils.functional import head_from_fun from modern_di import Container, Scope, integrations @@ -33,6 +34,11 @@ def _bind_current_app(main: str | None, container: Container) -> Celery: return process_app +def _move_to_end(signal: Signal, receiver: typing.Callable[..., None]) -> None: + signal.disconnect(receiver) + signal.connect(receiver, weak=False) + + def setup_di(app: Celery, container: Container, *, manage_lifespan: bool = True) -> Container: _ROOT_CONTAINERS[app] = container main = app.main @@ -50,12 +56,14 @@ def _close_for(owner: Celery) -> None: def _open_for_worker(sender: object = None, **_: typing.Any) -> None: # noqa: ANN401 _open_for(_worker_app(sender)) + _move_to_end(signals.worker_shutdown, _close_for_worker) def _close_for_worker(sender: object = None, **_: typing.Any) -> None: # noqa: ANN401 _close_for(_worker_app(sender)) def _open_for_process(**_: typing.Any) -> None: # noqa: ANN401 _open_for(_bind_current_app(main, container)) + _move_to_end(signals.worker_process_shutdown, _close_for_process) def _close_for_process(**_: typing.Any) -> None: # noqa: ANN401 _close_for(_current_app()) diff --git a/tests/test_shutdown_order.py b/tests/test_shutdown_order.py new file mode 100644 index 0000000..701c614 --- /dev/null +++ b/tests/test_shutdown_order.py @@ -0,0 +1,106 @@ +import typing + +import pytest +from celery import Celery, signals +from celery.utils.dispatch import Signal +from modern_di import Container, Group, Scope, providers + +from modern_di_celery import fetch_di_container, setup_di + + +class Pool: + pass + + +class _FakeWorker: + def __init__(self, app: Celery) -> None: + self.app = app + + +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 _make_app(main: str) -> Celery: + app = Celery(main, broker="memory://", backend="cache+memory://") + app.set_current() + return app + + +_WORKER_PAIR = (signals.worker_init, signals.worker_shutdown, True) +_PROCESS_PAIR = (signals.worker_process_init, signals.worker_process_shutdown, False) +_PAIRS = pytest.mark.parametrize( + ("init", "shutdown", "sent_by_worker"), [_WORKER_PAIR, _PROCESS_PAIR], ids=["worker", "process"] +) + + +def _cycle(app: Celery, init: Signal, shutdown: Signal, *, sent_by_worker: bool) -> None: + sender = _FakeWorker(app) if sent_by_worker else None + init.send(sender=sender) + fetch_di_container(app).resolve(Pool) + shutdown.send(sender=sender) + + +@_PAIRS +def test_root_closes_after_shutdown_receivers_connected_after_setup_di( + init: Signal, shutdown: Signal, sent_by_worker: bool +) -> None: + events: list[str] = [] + app = _make_app(f"late-receiver-{shutdown.name}") + container = setup_di(app, _container(events)) + + def hook(**_: typing.Any) -> None: # noqa: ANN401 + container.resolve(Pool) + events.append("user") + + shutdown.connect(hook, weak=False) + try: + _cycle(app, init, shutdown, sent_by_worker=sent_by_worker) + finally: + shutdown.disconnect(hook) + + assert events == ["user", "finalized"] + assert container.closed is True + + +def _receiver_counts(*watched: Signal) -> list[int]: + return [len(signal.receivers or ()) for signal in watched] + + +@_PAIRS +def test_restarts_do_not_add_shutdown_receivers(init: Signal, shutdown: Signal, sent_by_worker: bool) -> None: + events: list[str] = [] + app = _make_app(f"restart-{shutdown.name}") + setup_di(app, _container(events)) + receivers_after_setup = _receiver_counts(init, shutdown) + + _cycle(app, init, shutdown, sent_by_worker=sent_by_worker) + _cycle(app, init, shutdown, sent_by_worker=sent_by_worker) + + assert _receiver_counts(init, shutdown) == receivers_after_setup + assert events == ["finalized", "finalized"] + + +@_PAIRS +def test_manage_lifespan_false_still_neither_opens_nor_closes_the_root( + init: Signal, shutdown: Signal, sent_by_worker: bool +) -> None: + events: list[str] = [] + app = _make_app(f"unmanaged-{shutdown.name}") + container = setup_di(app, _container(events), manage_lifespan=False) + container.close_sync() + + sender = _FakeWorker(app) if sent_by_worker else None + init.send(sender=sender) + assert container.closed is True + + container.open() + container.resolve(Pool) + shutdown.send(sender=sender) + assert container.closed is False + assert events == []