From 4c4aeeb341a1f37f6d1d47e17ef21cd253ce1c56 Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Sat, 17 Oct 2026 15:17:45 +0300 Subject: [PATCH 1/2] test: cover worker signals with two apps in one process --- tests/conftest.py | 11 ++++++++++- tests/test_lifespan.py | 39 +++++++++++++++++++++++++++++++++++++++ tests/test_spawn.py | 8 -------- 3 files changed, 49 insertions(+), 9 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 8382f26..a53bfa8 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,11 +1,20 @@ +import typing + import pytest -from celery import Celery +from celery import Celery, current_app from modern_di import Container from modern_di_celery import setup_di from tests.dependencies import Dependencies +@pytest.fixture(autouse=True) +def _restore_current_app() -> typing.Iterator[None]: + previous = current_app._get_current_object() # noqa: SLF001 + yield + previous.set_current() + + @pytest.fixture def app() -> Celery: app_ = Celery("test", broker="memory://", backend="cache+memory://") diff --git a/tests/test_lifespan.py b/tests/test_lifespan.py index d39b01e..2e75d50 100644 --- a/tests/test_lifespan.py +++ b/tests/test_lifespan.py @@ -64,3 +64,42 @@ def sample( signals.worker_shutdown.send(sender=None) assert container.closed is True + + +class _FakeWorker: + def __init__(self, app: Celery) -> None: + self.app = app + + +def _make_app(main: str) -> Celery: + return Celery(main, broker="memory://", backend="cache+memory://", set_as_current=False) + + +def test_worker_signals_act_only_on_the_senders_app() -> None: + app_a = _make_app("worker-a") + app_b = _make_app("worker-b") + container_a = modern_di_celery.setup_di(app_a, Container(groups=[Dependencies])) + container_b = modern_di_celery.setup_di(app_b, Container(groups=[Dependencies])) + worker_a = _FakeWorker(app_a) + + signals.worker_shutdown.send(sender=worker_a) + assert (container_a.closed, container_b.closed) == (True, False) + + container_b.close_sync() + signals.worker_init.send(sender=worker_a) + assert (container_a.closed, container_b.closed) == (False, True) + + +def test_process_signals_act_only_on_the_current_app() -> None: + app_a = _make_app("process-a") + app_b = _make_app("process-b") + container_a = modern_di_celery.setup_di(app_a, Container(groups=[Dependencies])) + container_b = modern_di_celery.setup_di(app_b, Container(groups=[Dependencies])) + app_a.set_current() + + signals.worker_process_shutdown.send(sender=None) + assert (container_a.closed, container_b.closed) == (True, False) + + container_b.close_sync() + signals.worker_process_init.send(sender=None) + assert (container_a.closed, container_b.closed) == (False, True) diff --git a/tests/test_spawn.py b/tests/test_spawn.py index ece0dfa..6b0dfd0 100644 --- a/tests/test_spawn.py +++ b/tests/test_spawn.py @@ -1,7 +1,6 @@ import pickle import typing -import pytest from celery import Celery, current_app, signals from modern_di import Container @@ -13,13 +12,6 @@ def _make_app(main: str, set_as_current: bool = True) -> Celery: return Celery(main, broker="memory://", backend="cache+memory://", set_as_current=set_as_current) -@pytest.fixture(autouse=True) -def _restore_current_app() -> typing.Iterator[None]: - previous = current_app._get_current_object() # noqa: SLF001 - yield - previous.set_current() - - def test_a_configured_app_survives_a_pickle_round_trip() -> None: app = _make_app("pickled") setup_di(app, Container(groups=[Dependencies])) From b53c4b51bef2d32293c9317857c444ec3289a9d3 Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Sat, 17 Oct 2026 15:18:34 +0300 Subject: [PATCH 2/2] fix: act on worker signals only for the app that owns the container --- modern_di_celery/main.py | 69 +++++++++++++++++++++++++--------------- 1 file changed, 44 insertions(+), 25 deletions(-) diff --git a/modern_di_celery/main.py b/modern_di_celery/main.py index 0dca1b0..b3f998f 100644 --- a/modern_di_celery/main.py +++ b/modern_di_celery/main.py @@ -17,44 +17,63 @@ def _real_app(app: Celery) -> Celery: return app +def _current_app() -> Celery: + return _real_app(typing.cast(Celery, current_app)) + + +def _worker_app(sender: object) -> Celery: + app = getattr(sender, "app", None) + return _current_app() if app is None else _real_app(app) + + +def _bind_current_app(main: str | None, container: Container) -> Celery: + process_app = _current_app() + if process_app.main == main: + _ROOT_CONTAINERS.setdefault(process_app, container) + return process_app + + def setup_di(app: Celery, container: Container) -> Container: _ROOT_CONTAINERS[app] = container main = app.main - def _open_container(**_: typing.Any) -> None: # noqa: ANN401 - container.open() + def _owns(owner: Celery) -> bool: + return _ROOT_CONTAINERS.get(owner) is container + + def _open_for_worker(sender: object = None, **_: typing.Any) -> None: # noqa: ANN401 + if _owns(_worker_app(sender)): + container.open() + + def _close_for_worker(sender: object = None, **_: typing.Any) -> None: # noqa: ANN401 + if _owns(_worker_app(sender)): + container.close_sync() - def _open_and_bind_container(**_: typing.Any) -> None: # noqa: ANN401 - container.open() - process_app = _real_app(typing.cast(Celery, current_app)) - if process_app.main == main: - _ROOT_CONTAINERS.setdefault(process_app, container) + def _open_for_process(**_: typing.Any) -> None: # noqa: ANN401 + if _owns(_bind_current_app(main, container)): + container.open() - def _close_container(**_: typing.Any) -> None: # noqa: ANN401 - container.close_sync() + def _close_for_process(**_: typing.Any) -> None: # noqa: ANN401 + if _owns(_current_app()): + container.close_sync() # weak=False is required: Celery signals default to weak refs, which would # garbage-collect these handlers before a worker process ever fires them. # # worker_process_init/worker_process_shutdown fire only under the prefork and - # solo pools (once per forked child for prefork), giving each forked process - # its own open/close so cached resources and finalizers are fork-safe. - # worker_init/worker_shutdown fire once in the main worker process for ALL - # pools, which is the only signal the gevent/eventlet/threads pools send — - # those pools never fork and never emit worker_process_init, so without this - # the root container would stay closed and every @inject task would raise - # ContainerClosedError. The overlap between the two pairs is harmless: - # Container.open() is a no-op when already open, and close_sync() is a no-op + # solo pools; worker_init/worker_shutdown fire once in the main worker process + # for all pools, and are the only pair gevent/eventlet/threads send. The + # overlap is harmless: open() is a no-op on an open container, close_sync() # when nothing was cached. # - # A spawn child unpickles the app it is sent and makes that copy current, then - # re-imports the task modules, so setup_di runs again on a new app object. The - # process-init handler binds its container to that current copy by main name, - # unless the copy already has one. - signals.worker_process_init.connect(_open_and_bind_container, weak=False) - signals.worker_process_shutdown.connect(_close_container, weak=False) - signals.worker_init.connect(_open_container, weak=False) - signals.worker_shutdown.connect(_close_container, weak=False) + # Each handler acts only when its app owns this container: the worker's app + # for the worker pair, the process's current app for the process pair. A + # spawn child makes an unpickled copy of the app current and re-runs + # setup_di on a new app object, so the process-init handler first binds its + # container to that copy by main name, unless the copy already has one. + signals.worker_process_init.connect(_open_for_process, weak=False) + signals.worker_process_shutdown.connect(_close_for_process, weak=False) + signals.worker_init.connect(_open_for_worker, weak=False) + signals.worker_shutdown.connect(_close_for_worker, weak=False) return container