Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
adhikjoshi wants to merge 1 commit into
base: fix/redis-blocking-read-hang
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
252 changes: 174 additions & 78 deletions modelq/app/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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()
Expand All @@ -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
Expand All @@ -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(
Expand All @@ -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.
Expand Down Expand Up @@ -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)
Expand Down
29 changes: 15 additions & 14 deletions tests/test_connection_liveness.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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():
Expand Down
Loading
Loading