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"