feat: hold tasks in a per-worker in-flight list so none can be lost in transit

BLPOP removes the task from the queue and *then* writes it to the client. If
that write is lost, the task exists nowhere: not in ml_tasks, not in
processing_tasks (which is only populated after the reply arrives), and so
invisible to requeue_stuck_processing_tasks and every other recovery path.

This is not theoretical. Reproduced against production on 2026-08-15 by
blackholing a worker's connection at the Redis host and pushing a marked probe
task. Redis served it to the blocked client — the client's idle timer reset at
the moment of the push, so the pop and the write both happened — and the task
was then absent from the queue with no worker ever logging receipt. It very
likely explains the generation records seen stuck in "processing" with nothing
queued behind them.

BLMOVE makes taking the task and recording who took it a single atomic step:

- the worker moves the task into `inflight:{server_id}:{worker_id}` and holds
  custody until it is done, whatever the outcome, via a finally
- every in-flight list is registered in a set, so recovery enumerates them
  instead of SCANning for them
- at startup a server drains its own lists, which can only be debris from a
  previous run of the same id
- the pruning loop drains lists whose owning server is no longer registered,
  running after prune_inactive_servers so the dead are already deregistered

include_self is the one flag that must not be wrong. At startup our lists are
debris and must be drained; from the periodic sweep they belong to our own
running workers, and draining those would hand a live task to a second worker.
Both directions are tested.

drain_inflight is bounded by the list length read up front rather than looping
until empty, so it cannot spin if the no-appender assumption ever breaks.

Requires Redis 6.2+ for BLMOVE/LMOVE. The fleet runs 7.4.2.

The bounded-read test from the connection-liveness work now asserts the
timeout on BLMOVE rather than BLPOP; the property it pins is unchanged.

Mutation-tested: reverting to BLPOP, dropping the custody release, and letting
the sweep drain live workers' lists each turn a test red. The first two
initially survived against weaker assertions, which is why the custody test now
observes a genuinely in-flight task mid-run rather than inspecting list
lengths after the fact.

Author: adhikjoshi

modelq/app/base.py

@@ -138,6 +138,17 @@ class ModelQ:
     HEALTH_CHECK_INTERVAL = 30  # seconds: PING a pooled connection idle this long
     BACKGROUND_LOOP_BACKOFF = 5  # seconds: pause before retrying a crashed loop body
 
+    # --- in-flight task custody ------------------------------------------------
+    # BLPOP removes the task from the list and *then* writes it to the client. If
+    # that write is lost, the task exists nowhere: not in the queue, not in
+    # processing_tasks (which is only populated after the reply arrives), and so
+    # invisible to every recovery sweep. BLMOVE instead moves the task into a
+    # per-worker in-flight list as a single atomic step, so a task in transit to
+    # a worker that never receives it stays parked somewhere we can find it.
+    INFLIGHT_PREFIX = "inflight"
+    # Registry of every in-flight list, so recovery never has to SCAN for them.
+    INFLIGHT_REGISTRY = "inflight_lists"
+
     def __init__(
         self,
         host: str = "localhost",
@@ -733,6 +744,12 @@ def start_workers(self, no_of_workers: int = 1):
         else:
             self.check_middleware("before_worker_boot")
 
+        # Anything still in OUR in-flight lists is debris from a previous run of
+        # this server_id — a live worker of ours cannot exist yet. Drain before
+        # starting workers, never after, or we would yank a task out from under
+        # a worker that had just claimed it.
+        self.recover_abandoned_inflight_tasks(include_self=True)
+
         # 1) Delayed re-queue thread
         requeue_thread = threading.Thread(target=self.requeue_delayed_tasks, daemon=True)
         requeue_thread.start()
@@ -751,6 +768,10 @@ def start_workers(self, no_of_workers: int = 1):
         # 4) Worker threads
         def worker_loop(worker_id):
             self.check_middleware("after_worker_boot")
+            inflight_key = self._inflight_key(worker_id)
+            # Register before taking custody of anything, so a task can never be
+            # held in a list that recovery does not know to look in.
+            self.redis_client.sadd(self.INFLIGHT_REGISTRY, inflight_key)
             while True:
                 try:
                     # Check worker health before picking up tasks
@@ -767,91 +788,104 @@ def worker_loop(worker_id):
                     # forgotten, while the queue behind it grows unattended.
                     # Timing out and looping forces the read to complete, which
                     # is what lets keepalive/health-check reap the dead socket.
-                    task_data = self.redis_client.blpop("ml_tasks", timeout=self.BLPOP_TIMEOUT)
-                    if not task_data:
-                        continue
-
-                    self.update_server_status(f"worker_{worker_id}: busy")
-                    _, task_json = task_data
-                    task_dict = json.loads(task_json)
-                    task = Task.from_dict(task_dict)
-
-                    # Mark task as 'processing'
-                    added = self.redis_client.sadd("processing_tasks", task.task_id)
-                    if added == 0:
-                        logger.warning(
-                            f"Task {task.task_id} is already being processed. Skipping duplicate."
-                        )
+                    # Custody handoff, not a handoff-and-hope. BLPOP removes the
+                    # task and *then* writes it; if that write is lost the task is
+                    # gone from every structure that could recover it. BLMOVE makes
+                    # taking the task and recording who took it one atomic step, so
+                    # a task in transit to a worker that never receives it stays in
+                    # `inflight_key` until a sweep returns it to the queue.
+                    task_json = self.redis_client.blmove(
+                        "ml_tasks", inflight_key, self.BLPOP_TIMEOUT, "LEFT", "RIGHT"
+                    )
+                    if not task_json:
                         continue
-                    task.status = "processing"
-
-                    # The task has left the queue (claimed for processing). Keep the
-                    # `queued_requests` index in sync with `ml_tasks`; otherwise it
-                    # accumulates every completed/failed task forever and badly
-                    # inflates queue_num / queue_time.
-                    self.redis_client.zrem("queued_requests", task.task_id)
-
-                    # Set started_at
-                    task_dict["started_at"] = time.time()
-
-                    # Update in Redis
-                    self.redis_client.set(f"task:{task.task_id}", json.dumps(task_dict),ex=86400)
-
-                    if task.task_name in self.allowed_tasks:
-                        try:
-                            logger.info(f"Worker {worker_id} started processing: {task.task_name}")
-
-                            # Add Sentry breadcrumb for task processing
-                            if self.sentry_enabled:
-                                add_breadcrumb(
-                                    message=f"Processing task: {task.task_name}",
-                                    category="task",
-                                    level="info",
-                                    data={"task_id": task.task_id, "worker_id": worker_id},
-                                )
 
-                            start_time = time.time()
-                            self.process_task(task)
-                            end_time = time.time()
-                            logger.info(
-                                f"Worker {worker_id} finished {task.task_name} "
-                                f"in {end_time - start_time:.2f} seconds"
+                    # Release custody only when we are finished with it, whatever
+                    # the outcome. If this process dies mid-task the entry stays
+                    # put on purpose — that is what makes it recoverable.
+                    try:
+                        self.update_server_status(f"worker_{worker_id}: busy")
+                        task_dict = json.loads(task_json)
+                        task = Task.from_dict(task_dict)
+
+                        # Mark task as 'processing'
+                        added = self.redis_client.sadd("processing_tasks", task.task_id)
+                        if added == 0:
+                            logger.warning(
+                                f"Task {task.task_id} is already being processed. Skipping duplicate."
                             )
-
-                        except TaskProcessingError as e:
-                            if self._should_ignore_sentry_exception(e.__cause__):
-                                logger.warning(
-                                    "Worker %s encountered an ignored Sentry TaskProcessingError: %s",
-                                    worker_id,
-                                    e,
+                            continue
+                        task.status = "processing"
+
+                        # The task has left the queue (claimed for processing). Keep the
+                        # `queued_requests` index in sync with `ml_tasks`; otherwise it
+                        # accumulates every completed/failed task forever and badly
+                        # inflates queue_num / queue_time.
+                        self.redis_client.zrem("queued_requests", task.task_id)
+
+                        # Set started_at
+                        task_dict["started_at"] = time.time()
+
+                        # Update in Redis
+                        self.redis_client.set(f"task:{task.task_id}", json.dumps(task_dict),ex=86400)
+
+                        if task.task_name in self.allowed_tasks:
+                            try:
+                                logger.info(f"Worker {worker_id} started processing: {task.task_name}")
+
+                                # Add Sentry breadcrumb for task processing
+                                if self.sentry_enabled:
+                                    add_breadcrumb(
+                                        message=f"Processing task: {task.task_name}",
+                                        category="task",
+                                        level="info",
+                                        data={"task_id": task.task_id, "worker_id": worker_id},
+                                    )
+
+                                start_time = time.time()
+                                self.process_task(task)
+                                end_time = time.time()
+                                logger.info(
+                                    f"Worker {worker_id} finished {task.task_name} "
+                                    f"in {end_time - start_time:.2f} seconds"
                                 )
-                            else:
+
+                            except TaskProcessingError as e:
+                                if self._should_ignore_sentry_exception(e.__cause__):
+                                    logger.warning(
+                                        "Worker %s encountered an ignored Sentry TaskProcessingError: %s",
+                                        worker_id,
+                                        e,
+                                    )
+                                else:
+                                    logger.error(
+                                        f"Worker {worker_id} encountered a TaskProcessingError: {e}"
+                                    )
+                                if task.payload.get("retries", 0) > 0:
+                                    new_task_dict = task.to_dict()
+                                    new_task_dict["payload"] = task.original_payload
+                                    new_task_dict["payload"]["retries"] -= 1
+                                    self.enqueue_delayed_task(new_task_dict, delay_seconds=self.delay_seconds)
+
+                            except Exception as e:
                                 logger.error(
-                                    f"Worker {worker_id} encountered a TaskProcessingError: {e}"
+                                    f"Worker {worker_id} encountered an unexpected error: {e}"
                                 )
-                            if task.payload.get("retries", 0) > 0:
-                                new_task_dict = task.to_dict()
-                                new_task_dict["payload"] = task.original_payload
-                                new_task_dict["payload"]["retries"] -= 1
-                                self.enqueue_delayed_task(new_task_dict, delay_seconds=self.delay_seconds)
-
-                        except Exception as e:
-                            logger.error(
-                                f"Worker {worker_id} encountered an unexpected error: {e}"
+                                if task.payload.get("retries", 0) > 0:
+                                    new_task_dict = task.to_dict()
+                                    new_task_dict["payload"] = task.original_payload
+                                    new_task_dict["payload"]["retries"] -= 1
+                                    self.enqueue_delayed_task(new_task_dict, delay_seconds=self.delay_seconds)
+                        else:
+                            # If task is not allowed on this server, re-queue it
+                            logger.warning(
+                                f"Worker {worker_id} cannot process task {task.task_name}, re-queueing..."
                             )
-                            if task.payload.get("retries", 0) > 0:
-                                new_task_dict = task.to_dict()
-                                new_task_dict["payload"] = task.original_payload
-                                new_task_dict["payload"]["retries"] -= 1
-                                self.enqueue_delayed_task(new_task_dict, delay_seconds=self.delay_seconds)
-                    else:
-                        # If task is not allowed on this server, re-queue it
-                        logger.warning(
-                            f"Worker {worker_id} cannot process task {task.task_name}, re-queueing..."
-                        )
-                        self.redis_client.rpush("ml_tasks", task_json)
-                        self.redis_client.zadd("queued_requests", {task.task_id: task_dict.get("queued_at", time.time())})
-                        self.redis_client.srem("processing_tasks", task.task_id)
+                            self.redis_client.rpush("ml_tasks", task_json)
+                            self.redis_client.zadd("queued_requests", {task.task_id: task_dict.get("queued_at", time.time())})
+                            self.redis_client.srem("processing_tasks", task.task_id)
+                    finally:
+                        self.redis_client.lrem(inflight_key, 1, task_json)
 
                 except Exception as e:
                     logger.error(
@@ -877,6 +911,65 @@ def worker_loop(worker_id):
             f"Registered tasks: {task_names}"
         )
 
+    def _inflight_key(self, worker_id: int) -> str:
+        """Per-worker custody list. Scoped by server so recovery can attribute it."""
+        return f"{self.INFLIGHT_PREFIX}:{self.server_id}:{worker_id}"
+
+    def drain_inflight(self, inflight_key: str) -> int:
+        """Return every task held in `inflight_key` to the front of the queue.
+
+        Tasks land here only when a worker took custody but never finished, so
+        they are older than anything already queued and go back to the head.
+        Uses LMOVE so a crash mid-drain cannot lose a task: it is in one list or
+        the other at every instant, never in neither.
+        """
+        # Bounded by the length read up front. The owner is gone, so nothing is
+        # appending; an unbounded `while True` here would spin forever the day
+        # that assumption breaks.
+        moved = 0
+        for _ in range(self.redis_client.llen(inflight_key) or 0):
+            item = self.redis_client.lmove(inflight_key, "ml_tasks", "RIGHT", "LEFT")
+            if item is None:
+                break
+            moved += 1
+        self.redis_client.srem(self.INFLIGHT_REGISTRY, inflight_key)
+        if moved:
+            logger.warning(f"Recovered {moved} in-flight task(s) from '{inflight_key}'.")
+        return moved
+
+    def recover_abandoned_inflight_tasks(
+        self, active_server_ids=None, include_self: bool = False
+    ) -> int:
+        """Re-queue tasks stranded in the in-flight lists of dead workers.
+
+        `include_self` is the difference between the two callers, and getting it
+        wrong is the one way this can lose work. At startup our own lists are
+        debris from a previous run and must be drained. From the periodic sweep
+        they belong to our own live workers, which are mid-task — draining those
+        would hand the same task to somebody else while it is still running.
+        """
+        if active_server_ids is None:
+            active_server_ids = set(self.get_registered_server_ids() or [])
+        active_server_ids = {
+            s.decode() if isinstance(s, bytes) else s for s in active_server_ids
+        }
+
+        recovered = 0
+        for raw_key in self.redis_client.smembers(self.INFLIGHT_REGISTRY) or []:
+            key = raw_key.decode() if isinstance(raw_key, bytes) else raw_key
+            try:
+                _, owner, _ = key.split(":", 2)
+            except ValueError:
+                logger.warning(f"Ignoring malformed in-flight key '{key}'.")
+                continue
+            if owner == self.server_id:
+                if not include_self:
+                    continue
+            elif owner in active_server_ids:
+                continue  # a live worker elsewhere still owns it
+            recovered += self.drain_inflight(key)
+        return recovered
+
     @contextlib.contextmanager
     def _guarded_iteration(self, loop_name: str):
         """Swallow and log any exception raised by one background-loop iteration.
@@ -914,6 +1007,9 @@ def _pruning_loop(self):
         while True:
             with self._guarded_iteration("pruning"):
                 self.prune_inactive_servers(timeout_seconds=self.PRUNE_TIMEOUT)
+                # Runs after the prune so dead servers are already deregistered
+                # and their in-flight lists read as abandoned on this same pass.
+                self.recover_abandoned_inflight_tasks()
                 self.requeue_stuck_processing_tasks(threshold=180)
                 self.prune_old_task_results(older_than_seconds=self.TASK_RESULT_RETENTION)
             time.sleep(self.PRUNE_CHECK_INTERVAL)

tests/test_connection_liveness.py

@@ -32,33 +32,34 @@ def modelq_instance():
 # 1. the queue read must be bounded
 # ---------------------------------------------------------------------------
 
-def test_worker_blpop_passes_a_timeout(modelq_instance):
-    """The worker must never issue an unbounded BLPOP.
+def test_worker_queue_read_passes_a_timeout(modelq_instance):
+    """The worker must never issue an unbounded blocking read.
 
-    This is the actual outage. `blpop(key)` with no timeout blocks forever on a
-    socket Redis has already forgotten.
+    This is the actual outage. A blocking pop with no timeout waits forever on
+    a socket Redis has already forgotten. The command is now BLMOVE rather than
+    BLPOP (see the in-flight custody work), but the property under test is the
+    same one: the read is bounded.
     """
     seen = {}
     stop = threading.Event()
 
-    def fake_blpop(key, *args, **kwargs):
-        seen["args"] = args
-        seen["kwargs"] = kwargs
+    def fake_blmove(src, dst, timeout, *args, **kwargs):
+        seen["timeout"] = timeout
         stop.set()
-        # Behave like a real timed-out BLPOP so the worker loops rather than
-        # trying to decode a task.
+        # Behave like a real timed-out blocking move so the worker loops rather
+        # than trying to decode a task.
         time.sleep(0.01)
         return None
 
     modelq_instance.redis_client = MagicMock(wraps=modelq_instance.redis_client)
-    modelq_instance.redis_client.blpop.side_effect = fake_blpop
+    modelq_instance.redis_client.blmove.side_effect = fake_blmove
 
     modelq_instance.start_workers(no_of_workers=1)
-    assert stop.wait(timeout=5), "worker never called blpop"
+    assert stop.wait(timeout=5), "worker never issued a blocking queue read"
 
-    timeout = seen["kwargs"].get("timeout", seen["args"][0] if seen["args"] else None)
-    assert timeout is not None, "BLPOP was issued without a timeout"
-    assert timeout > 0, f"BLPOP timeout must be positive, got {timeout!r}"
+    timeout = seen["timeout"]
+    assert timeout is not None, "blocking read was issued without a timeout"
+    assert timeout > 0, f"timeout must be positive, got {timeout!r}"
 
 
 def test_blpop_timeout_stays_below_socket_timeout():