From 6c9add6ec581e883c3452d60f602a15380a3fe52 Mon Sep 17 00:00:00 2001 From: Yasha Date: Thu, 18 Jun 2026 15:11:04 +0200 Subject: [PATCH] fix(celery): drop consumer connection before re-enqueue to stop self-grab The reschedule-on-SIGTERM e2e test (test_reschedule_e2e) was flaky (~1 in 5), ending 'suspended' instead of 'completed'. Root cause is a real message-loss race in the reschedule path, not a test artifact. On the Redis transport, cancelling the task consumer does not abort a BRPOP the dying worker has already issued: the server still serves the *next* message pushed to those queues to that pending BRPOP. So the copy we re-enqueue on shutdown can be grabbed straight back by the dying worker, which then exits before recording it in `unacked`. The message is lost entirely -- not in the ready list, and not restorable via the broker visibility_timeout -- so the parent workflow stays SUSPENDED (until visibility_timeout, 3600s). Confirmed by inspecting Redis on a failing run: after the reschedule there is no ready-list entry at all and `unacked` holds only the original task. Passing runs have the re-enqueued message sitting in the ready list for a fresh worker. Fix: after cancelling the consumer, drop its broker connection (`connection.collect()`) to abort any in-flight BRPOP *before* re-publishing. The re-enqueue uses the producer pool (a separate connection) and the lock release uses the singleton backend, so both are unaffected. 15/15 stress runs and 6/6 e2e runs pass after the change (was ~1-in-5 failing). Co-Authored-By: Claude Opus 4.8 (1M context) --- pyworkflow/celery/reschedule.py | 23 +++++++++++++++++++++++ tests/unit/test_reschedule.py | 22 ++++++++++++++++++++++ 2 files changed, 45 insertions(+) diff --git a/pyworkflow/celery/reschedule.py b/pyworkflow/celery/reschedule.py index 15c8edc..9c5a1e7 100644 --- a/pyworkflow/celery/reschedule.py +++ b/pyworkflow/celery/reschedule.py @@ -290,9 +290,32 @@ def on_consumer_stop(consumer: Any) -> None: except Exception as exc: logger.warning(f"reschedule: failed to cancel task consumer: {exc}") + # Cancelling is not enough on the Redis transport: a BRPOP this worker already + # issued is accepted by the server and will still pop the *next* message pushed + # to those queues -- including the copy we are about to re-enqueue. The dying + # worker then exits before recording it in `unacked`, so the message is lost + # entirely (neither in the ready list nor restorable via visibility_timeout). + # Drop the consumer's broker connection so any in-flight BRPOP is aborted + # server-side *before* we re-publish; the re-enqueue itself uses the producer + # pool (a separate connection), so it is unaffected. + _abort_consumer_fetching(consumer) + reschedule_inflight_on_shutdown() +def _abort_consumer_fetching(consumer: Any) -> None: + """Close the consumer's broker connection to abort any in-flight BRPOP.""" + conn = getattr(consumer, "connection", None) + if conn is None: + return + try: + # collect() tears down channels + the underlying socket, aborting the + # pending BRPOP. Never raises into shutdown. + conn.collect() + except Exception as exc: + logger.warning(f"reschedule: failed to drop consumer connection: {exc}") + + class RescheduleConsumerStep(bootsteps.StartStopStep): """ Consumer bootstep that re-enqueues this worker's in-flight tasks the moment diff --git a/tests/unit/test_reschedule.py b/tests/unit/test_reschedule.py index 355d7ce..ed0a92f 100644 --- a/tests/unit/test_reschedule.py +++ b/tests/unit/test_reschedule.py @@ -261,6 +261,28 @@ def test_prerun_then_consumer_stop_roundtrip(self, fake_backend, _stopping): consumer.task_consumer.cancel.assert_called_once() # stop fetching first task.apply_async.assert_called_once() # then re-enqueue + def test_consumer_connection_dropped_before_reenqueue(self, fake_backend, _stopping): + # Cancelling the consumer is not enough on Redis: a BRPOP this worker + # already issued can still grab the re-enqueued copy and lose it on exit. + # The connection must be dropped (aborting the in-flight BRPOP) *before* + # the message is re-published. + task = _make_singleton_task() + fake_backend.registry[task.name] = task + _track(fake_backend, task, "task-1", {"run_id": "r1", "step_id": "s1"}) + + order = [] + consumer = MagicMock() + consumer.connection.collect.side_effect = lambda *_, **__: order.append("collect") + task.apply_async.side_effect = lambda *_, **__: order.append("apply_async") + + reschedule.on_consumer_stop(consumer) + + consumer.connection.collect.assert_called_once() + assert order == [ + "collect", + "apply_async", + ], "connection must be dropped before the task is re-enqueued" + def test_consumer_stop_is_noop_when_not_shutting_down(self, fake_backend, monkeypatch): # Transient consumer restart (e.g. broker reconnect): flags are clear. from celery.worker import state