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