From db7cce15f48dca3d9f0f5a5646aa4601cdb7fd05 Mon Sep 17 00:00:00 2001 From: Adhik Joshi Date: Sat, 15 Aug 2026 15:38:17 +0530 Subject: [PATCH 1/2] 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 c8ca61d..87466f7 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", @@ -795,6 +806,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() @@ -813,6 +830,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 @@ -829,91 +850,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( @@ -939,6 +973,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. @@ -976,6 +1069,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) # prune_old_task_results() is deliberately NOT called here. Every # task_result key is written with a TTL, so Redis expires it on 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" From e71c9f324b895341446e5a0c2fe78d355550c9cc Mon Sep 17 00:00:00 2001 From: Tanmay patil Date: Thu, 20 Aug 2026 13:03:47 +0530 Subject: [PATCH 2/2] fix: make in-flight recovery immediately claimable --- modelq/app/base.py | 64 +++++++++++++++++++++++++++++++--- tests/test_inflight_custody.py | 30 ++++++++++++++++ 2 files changed, 89 insertions(+), 5 deletions(-) diff --git a/modelq/app/base.py b/modelq/app/base.py index 87466f7..e04f8e5 100644 --- a/modelq/app/base.py +++ b/modelq/app/base.py @@ -981,19 +981,73 @@ 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. + they are older than anything already queued and go back to the head. A + recovered task may already be present in ``processing_tasks`` if its + worker died after pickup. The queue handoff, processing-marker cleanup, + task-state reset, and queue-index restore therefore happen in one Redis + transaction. A new worker can never observe the queued copy while the + stale processing marker is still present and reject the only copy as a + duplicate. """ # 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") + while True: + pipe = self.redis_client.pipeline() + try: + # Multiple healthy servers can discover the same abandoned + # list on the same pruning pass. WATCH makes the tail read + # below and the subsequent RPOP refer to the same item. + pipe.watch(inflight_key) + item = pipe.lindex(inflight_key, -1) + if item is None: + pipe.unwatch() + break + + recovered_item = item + task_id = None + queued_at = time.time() + try: + task_dict = json.loads(item) + task_id = task_dict.get("task_id") + if task_id: + # Match the established stuck-task recovery contract: + # recovered work is queued again with a fresh queue + # timestamp and a bounded task record. + task_dict["status"] = "queued" + task_dict["queued_at"] = queued_at + recovered_item = json.dumps(task_dict) + except Exception as exc: + logger.warning( + f"Could not decode in-flight entry from '{inflight_key}' " + f"during recovery: {exc}. Returning it unchanged." + ) + + pipe.multi() + pipe.rpop(inflight_key) + pipe.lpush("ml_tasks", recovered_item) + if task_id: + pipe.set( + f"task:{task_id}", + recovered_item, + ex=self.task_ttl, + ) + pipe.zadd("queued_requests", {task_id: queued_at}) + pipe.srem("processing_tasks", task_id) + results = pipe.execute() + if results[0] is not None: + moved += 1 + break + except redis.WatchError: + # Another recovery worker won the race. Re-read the new tail + # rather than applying cleanup for the item it moved. + continue + finally: + pipe.reset() 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}'.") diff --git a/tests/test_inflight_custody.py b/tests/test_inflight_custody.py index 6f3acb6..22765cc 100644 --- a/tests/test_inflight_custody.py +++ b/tests/test_inflight_custody.py @@ -60,6 +60,36 @@ def test_drain_returns_held_tasks_to_the_queue(mq): assert mq.redis_client.smembers(mq.INFLIGHT_REGISTRY) == set() +def test_drain_clears_stale_processing_marker_before_task_is_visible(mq): + """A mid-task crash must not trip the next worker's duplicate guard.""" + inflight = mq._inflight_key(0) + mq.redis_client.rpush(inflight, _task("crashed")) + mq.redis_client.sadd(mq.INFLIGHT_REGISTRY, inflight) + mq.redis_client.sadd("processing_tasks", "crashed") + + assert mq.drain_inflight(inflight) == 1 + + recovered = json.loads(mq.redis_client.lindex("ml_tasks", 0)) + assert recovered["task_id"] == "crashed" + assert recovered["status"] == "queued" + assert not mq.redis_client.sismember("processing_tasks", "crashed") + assert mq.redis_client.zscore("queued_requests", "crashed") is not None + assert mq.redis_client.ttl("task:crashed") > 0 + + # This is the exact guard that previously discarded the recovered copy. + assert mq.redis_client.sadd("processing_tasks", "crashed") == 1 + + +def test_drain_keeps_malformed_entries_recoverable(mq): + """Bad payloads still move atomically instead of blocking the whole list.""" + inflight = mq._inflight_key(0) + mq.redis_client.rpush(inflight, "not-json") + mq.redis_client.sadd(mq.INFLIGHT_REGISTRY, inflight) + + assert mq.drain_inflight(inflight) == 1 + assert mq.redis_client.lpop("ml_tasks") == b"not-json" + + 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