From 51b2954abc4fb50006fffa275ac5eef0000ff405 Mon Sep 17 00:00:00 2001 From: Adhik Joshi Date: Sat, 15 Aug 2026 15:38:17 +0530 Subject: [PATCH] feat: hold tasks in a per-worker in-flight list so none can be lost in transit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- modelq/app/base.py | 252 +++++++++++++++++++++--------- tests/test_connection_liveness.py | 29 ++-- tests/test_inflight_custody.py | 181 +++++++++++++++++++++ 3 files changed, 370 insertions(+), 92 deletions(-) create mode 100644 tests/test_inflight_custody.py diff --git a/modelq/app/base.py b/modelq/app/base.py index b786361..03ee2cc 100644 --- a/modelq/app/base.py +++ b/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) diff --git a/tests/test_connection_liveness.py b/tests/test_connection_liveness.py index fcfbdaa..e9f3482 100644 --- a/tests/test_connection_liveness.py +++ b/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(): diff --git a/tests/test_inflight_custody.py b/tests/test_inflight_custody.py new file mode 100644 index 0000000..6f3acb6 --- /dev/null +++ b/tests/test_inflight_custody.py @@ -0,0 +1,181 @@ +"""Tests for in-flight task custody. + +Proven against production on 2026-08-15: with a worker's connection blackholed, +a task pushed to `ml_tasks` was popped by Redis, written into the dead socket, +and lost. It was not in the queue, not in `processing_tasks`, and no worker ever +logged receiving it — invisible to every recovery path. + +BLMOVE closes that hole by making "take the task" and "record who took it" one +atomic step. +""" + +import json +import threading +import time + +import fakeredis +import pytest + +from modelq import ModelQ + + +@pytest.fixture +def mq(): + return ModelQ(redis_client=fakeredis.FakeStrictRedis(), server_id="srv-a") + + +def _task(task_id="t1", name="noop"): + return json.dumps( + {"task_id": task_id, "task_name": name, "payload": {}, "status": "queued"} + ) + + +# --------------------------------------------------------------------------- +# custody +# --------------------------------------------------------------------------- + +def test_a_task_in_transit_is_never_in_neither_list(mq): + """The whole point: after the move the task is still somewhere findable. + + Under BLPOP this window is where the task ceased to exist. + """ + inflight = mq._inflight_key(0) + mq.redis_client.rpush("ml_tasks", _task()) + + moved = mq.redis_client.blmove("ml_tasks", inflight, 1, "LEFT", "RIGHT") + + assert moved is not None + assert mq.redis_client.llen("ml_tasks") == 0, "left the queue" + assert mq.redis_client.llen(inflight) == 1, "but is held in custody, not lost" + + +def test_drain_returns_held_tasks_to_the_queue(mq): + inflight = mq._inflight_key(0) + mq.redis_client.rpush(inflight, _task("t1"), _task("t2")) + mq.redis_client.sadd(mq.INFLIGHT_REGISTRY, inflight) + + assert mq.drain_inflight(inflight) == 2 + assert mq.redis_client.llen("ml_tasks") == 2 + assert mq.redis_client.llen(inflight) == 0 + assert mq.redis_client.smembers(mq.INFLIGHT_REGISTRY) == set() + + +def test_drain_is_bounded_and_terminates_on_empty(mq): + """An unbounded drain loop spins forever the day the list refills.""" + assert mq.drain_inflight(mq._inflight_key(9)) == 0 + + +# --------------------------------------------------------------------------- +# recovery: whose lists get drained +# --------------------------------------------------------------------------- + +def test_dead_servers_inflight_tasks_are_recovered(mq): + dead = f"{mq.INFLIGHT_PREFIX}:srv-ghost:0" + mq.redis_client.rpush(dead, _task("orphan")) + mq.redis_client.sadd(mq.INFLIGHT_REGISTRY, dead) + + recovered = mq.recover_abandoned_inflight_tasks(active_server_ids={"srv-a"}) + + assert recovered == 1 + assert mq.redis_client.llen("ml_tasks") == 1 + + +def test_a_live_servers_inflight_tasks_are_left_alone(mq): + """Draining a running worker's list hands its task to somebody else.""" + live = f"{mq.INFLIGHT_PREFIX}:srv-b:0" + mq.redis_client.rpush(live, _task("in-progress")) + mq.redis_client.sadd(mq.INFLIGHT_REGISTRY, live) + + recovered = mq.recover_abandoned_inflight_tasks(active_server_ids={"srv-a", "srv-b"}) + + assert recovered == 0 + assert mq.redis_client.llen(live) == 1 + assert mq.redis_client.llen("ml_tasks") == 0 + + +def test_own_list_is_skipped_during_the_periodic_sweep(mq): + """Our own workers are mid-task; the sweep must not touch them.""" + mine = mq._inflight_key(0) + mq.redis_client.rpush(mine, _task("mine")) + mq.redis_client.sadd(mq.INFLIGHT_REGISTRY, mine) + + assert mq.recover_abandoned_inflight_tasks(active_server_ids={"srv-a"}) == 0 + assert mq.redis_client.llen(mine) == 1 + + +def test_own_list_is_drained_at_startup(mq): + """At startup our own lists are debris from a previous run of this id.""" + mine = mq._inflight_key(0) + mq.redis_client.rpush(mine, _task("leftover")) + mq.redis_client.sadd(mq.INFLIGHT_REGISTRY, mine) + + recovered = mq.recover_abandoned_inflight_tasks( + active_server_ids={"srv-a"}, include_self=True + ) + + assert recovered == 1 + assert mq.redis_client.llen("ml_tasks") == 1 + + +def test_malformed_registry_entry_does_not_abort_recovery(mq): + """One bad key must not strand every other server's tasks.""" + mq.redis_client.sadd(mq.INFLIGHT_REGISTRY, "garbage") + dead = f"{mq.INFLIGHT_PREFIX}:srv-ghost:0" + mq.redis_client.rpush(dead, _task("orphan")) + mq.redis_client.sadd(mq.INFLIGHT_REGISTRY, dead) + + assert mq.recover_abandoned_inflight_tasks(active_server_ids={"srv-a"}) == 1 + + +# --------------------------------------------------------------------------- +# end-to-end through the worker loop +# --------------------------------------------------------------------------- + +def test_worker_registers_its_inflight_list_before_taking_work(mq): + """A task must never be held in a list recovery does not know about.""" + mq.start_workers(no_of_workers=1) + deadline = time.time() + 5 + while time.time() < deadline: + if mq.redis_client.smembers(mq.INFLIGHT_REGISTRY): + break + time.sleep(0.05) + + members = { + m.decode() if isinstance(m, bytes) else m + for m in mq.redis_client.smembers(mq.INFLIGHT_REGISTRY) + } + assert mq._inflight_key(0) in members + + +def test_worker_holds_custody_while_running_then_releases_it(mq): + """The two halves of custody, observed on a real in-flight task. + + Holding it is what BLPOP cannot do — under BLPOP the task is in no list at + all while it runs. Releasing it is what stops the list growing forever. + """ + inflight = mq._inflight_key(0) + running = threading.Event() + may_finish = threading.Event() + + @mq.task() + def slow_task(): + running.set() + may_finish.wait(timeout=10) + return "done" + + mq.start_workers(no_of_workers=1) + slow_task() + + assert running.wait(timeout=10), "worker never started the task" + + # Mid-flight: the task must be recorded as held by this worker. + assert mq.redis_client.llen(inflight) == 1, ( + "task is in flight but held in no list — it would be unrecoverable" + ) + + may_finish.set() + + deadline = time.time() + 10 + while time.time() < deadline and mq.redis_client.llen(inflight) != 0: + time.sleep(0.05) + assert mq.redis_client.llen(inflight) == 0, "custody was never released"