diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 73c61b8..af12980 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -6,6 +6,12 @@ Changelog next ==== +Bugfixes +-------- + +* Reset per-consumer state to fix the hangs after a consumer restart. + (`#107 `_) + Maintenance ----------- diff --git a/celery_batches/__init__.py b/celery_batches/__init__.py index f70b6ec..521c011 100644 --- a/celery_batches/__init__.py +++ b/celery_batches/__init__.py @@ -195,6 +195,17 @@ def Strategy(self, task: "Batches", app: Celery, consumer: Consumer) -> Callable # the default strategy does. # # See Batches._do_flush for ETA handling. + + # Strategy() runs every time the consumer (re)starts, including after a + # lost broker connection is re-established. Reset the per-consumer state + # so that the next message properly re-arms the flush timer. + if self._tref is not None: + self._tref.cancel() + self._tref = None + self._buffer = Queue() + self._pending = Queue() + self._count = count(1) + self._pool = consumer.pool hostname = consumer.hostname diff --git a/t/unit/__init__.py b/t/unit/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/t/unit/test_strategy.py b/t/unit/test_strategy.py new file mode 100644 index 0000000..a615456 --- /dev/null +++ b/t/unit/test_strategy.py @@ -0,0 +1,66 @@ +from typing import Any, cast +from unittest.mock import MagicMock + +from celery_batches import Batches + +from celery import Celery + + +def _make_batch_task() -> Batches: + app = Celery("test_batches_strategy", set_as_current=False) + + # Celery's task decorator is untyped, so it returns Any. + @app.task(base=Batches, flush_every=2, flush_interval=10) + def dummy(requests: list) -> None: + return None + + return cast(Batches, dummy) + + +def _mock_consumer() -> MagicMock: + consumer = MagicMock() + consumer.connection_errors = () + return consumer + + +def test_strategy_rearms_flush_timer_after_reconnect() -> None: + """Strategy() must reset per-consumer state so a reconnect re-arms the timer. + + Regression test for the wedge where, after a broker reconnect, the stale + flush timer (``self._tref``) stayed set and was never re-armed on the new + event-loop hub, so the worker stopped flushing batches. + """ + task = _make_batch_task() + + # Simulate state left over from a previous (now dead) consumer. The mock is + # typed as Any so that assigning it does not narrow ``_tref`` away from + # Optional, which would make the ``is None`` assertion below unreachable. + stale_timer: Any = MagicMock() + task._tref = stale_timer + task._buffer.put(MagicMock()) + task._pending.put(MagicMock()) + next(task._count) + + handler = task.Strategy(task, task.app, _mock_consumer()) + + # The stale timer is cancelled and cleared so the next message re-arms it + # via the ``if self._tref is None`` guard in the message handler. + stale_timer.cancel.assert_called_once_with() + assert task._tref is None + # Buffers tied to the old connection are dropped, and the flush counter is + # reset so flush_every alignment starts fresh. + assert task._buffer.empty() + assert task._pending.empty() + assert next(task._count) == 1 + assert callable(handler) + + +def test_strategy_without_existing_timer_is_a_noop_reset() -> None: + """First consumer start (no prior timer) must not raise.""" + task = _make_batch_task() + assert task._tref is None + + handler = task.Strategy(task, task.app, _mock_consumer()) + + assert task._tref is None + assert callable(handler)