0

Fault-tolerant work queue with leases and retries

by Frontier Interview· 1 day ago· ⑂ 0 remixes

Question: Fault-tolerant work queue with leases and retries

Companies: OpenAI

Difficulty: Medium

Problem statement

Implement an in-memory work queue that keeps working when workers crash. Workers take a task, process it, and report success or failure. A worker that dies silently must not lose the task: every reservation is a lease that runs out.

class WorkQueue:
    def __init__(self, lease_seconds: int, max_attempts: int): ...
    def add(self, task_id: str, now: int) -> bool: ...
    def reserve(self, now: int) -> str | None: ...
    def complete(self, task_id: str, now: int) -> bool: ...
    def fail(self, task_id: str, now: int) -> str | None: ...
    def status(self, task_id: str, now: int) -> str | None: ...

Every call passes the current time now. Times never go down between calls.

Rules

  • A task is in one of four states: AVAILABLE, RESERVED, COMPLETED, DEAD.
  • add puts a new task at the back of the available line and returns True. Adding an id that already exists returns False and changes nothing.
  • reserve takes the task at the front of the available line, marks it RESERVED, counts one attempt, and returns its id. The lease ends at now + lease_seconds. Return None if no task is available.
  • complete succeeds (returns True) only if the task is currently RESERVED; it becomes COMPLETED. Otherwise return False.
  • fail works only on a RESERVED task (otherwise return None). If the task has used fewer than max_attempts attempts, it goes to the back of the available line and fail returns "AVAILABLE". If it has used all of them, it becomes DEAD and fail returns "DEAD".
  • Lease expiry: before handling any call at time now, every lease whose end time is <= now expires. An expired lease is handled exactly like fail (retry or dead-letter). Leases that end at the same time expire in the order they were reserved.
  • status returns the task's state, or None for an unknown id.

Aim for O(log n) or better per call. Talk through how you find expired leases without scanning every task.

Constraints: up to 2 * 10^5 calls; times up to 10^9; 1 <= max_attempts <= 10^5; 0 <= lease_seconds <= 10^9.

Example 1

q = WorkQueue(lease_seconds=30, max_attempts=3)
q.add("resize-7", 0)    -> True
q.add("resize-8", 0)    -> True
q.reserve(5)            -> "resize-7"
q.fail("resize-7", 9)   -> "AVAILABLE"   # retried, now behind resize-8
q.reserve(10)           -> "resize-8"
q.reserve(11)           -> "resize-7"
q.complete("resize-8", 20) -> True

Example 2

q = WorkQueue(lease_seconds=4, max_attempts=2)
q.add("emb-1", 0)       -> True
q.reserve(1)            -> "emb-1"       # lease ends at 5
q.status("emb-1", 5)    -> "AVAILABLE"   # the worker went silent
q.reserve(6)            -> "emb-1"       # second attempt, ends at 10
q.complete("emb-1", 11) -> False         # too late: the lease expired at 10
q.status("emb-1", 11)   -> "DEAD"

Approach

The queue is a small state machine per task, driven by a clock that only moves forward:

AVAILABLE --reserve--> RESERVED --complete--> COMPLETED
    ^                     |
    |   fail / lease      |
    +---- expires --------+  (attempts < max)
                          |
                          +--> DEAD             (attempts == max)

A worker that crashes never calls fail, so the lease deadline is what makes the queue fault tolerant: after lease_seconds the task is treated exactly as if the worker had failed it.

Brute force. Keep a list of tasks. Before each call, scan every reserved task and expire the ones past their deadline. reserve scans for the oldest available task. That is O(n) per call, which is too slow for 10^5 operations, but it pins down the rules and is a fine first version to say out loud.

Optimal. Three structures, each answering one question fast:

  • tasks: dict[id, Task] holds state, attempt count and the current lease id. It answers status, complete and fail in O(1).
  • ready: deque holds AVAILABLE task ids in FIFO order. A new task and a retried task both go to the tail. A task sits in ready at most once, and it only leaves through reserve, so the deque never holds stale ids.
  • leases: min-heap of (deadline, lease_id, task_id) finds the next lease to expire. Before every operation we pop every entry with deadline <= now.

Lazy deletion. When a task is completed or failed we do not dig its entry out of the heap. Each reservation gets a fresh, increasing lease_id, and an entry is acted on only if the task is still RESERVED under that same lease id. Otherwise it is stale and is dropped. This is also what stops an old deadline from expiring a later reservation of the same task.

Ordering. lease_id grows with every reservation, so it is the heap tie-breaker. Leases with the same deadline expire in the order they were reserved, and their tasks re-enter ready in that order.

Invariants.

  • Every AVAILABLE task appears exactly once in ready; no other task appears there.
  • Every RESERVED task has exactly one live heap entry (matching lease_id).
  • attempts counts reservations. A task is DEAD exactly when an attempt ended without success and attempts == max_attempts.
  • After _expire(now), no live lease has a deadline <= now, so every operation sees a consistent picture of time now.

Code

from __future__ import annotations

import heapq
from collections import deque
from dataclasses import dataclass

AVAILABLE, RESERVED, COMPLETED, DEAD = "AVAILABLE", "RESERVED", "COMPLETED", "DEAD"


@dataclass
class Task:
    task_id: str
    state: str = AVAILABLE
    attempts: int = 0
    lease_id: int = 0  # id of the current reservation; 0 = none


class WorkQueue:
    def __init__(self, lease_seconds: int, max_attempts: int) -> None:
        self.lease_seconds = lease_seconds
        self.max_attempts = max_attempts
        self.tasks: dict[str, Task] = {}
        self.ready: deque[str] = deque()  # AVAILABLE tasks, FIFO
        # (deadline, lease_id, task_id); stale entries are skipped lazily.
        self.leases: list[tuple[int, int, str]] = []
        self.next_lease = 1

    # -- internal ---------------------------------------------------------
    def _expire(self, now: int) -> None:
        """Retire every lease whose deadline is <= now, oldest lease first."""
        while self.leases and self.leases[0][0] <= now:
            _, lease_id, task_id = heapq.heappop(self.leases)
            task = self.tasks[task_id]
            if task.state == RESERVED and task.lease_id == lease_id:
                self._release(task)

    def _release(self, task: Task) -> str:
        """End the current attempt without success: retry or dead-letter."""
        task.lease_id = 0
        if task.attempts >= self.max_attempts:
            task.state = DEAD
        else:
            task.state = AVAILABLE
            self.ready.append(task.task_id)
        return task.state

    def _held(self, task_id: str) -> Task | None:
        task = self.tasks.get(task_id)
        return task if task is not None and task.state == RESERVED else None

    # -- public API -------------------------------------------------------
    def add(self, task_id: str, now: int) -> bool:
        self._expire(now)
        if task_id in self.tasks:
            return False
        self.tasks[task_id] = Task(task_id)
        self.ready.append(task_id)
        return True

    def reserve(self, now: int) -> str | None:
        self._expire(now)
        if not self.ready:
            return None
        task = self.tasks[self.ready.popleft()]
        task.state = RESERVED
        task.attempts += 1
        task.lease_id = self.next_lease
        self.next_lease += 1
        heapq.heappush(self.leases, (now + self.lease_seconds, task.lease_id, task.task_id))
        return task.task_id

    def complete(self, task_id: str, now: int) -> bool:
        self._expire(now)
        task = self._held(task_id)
        if task is None:
            return False
        task.state, task.lease_id = COMPLETED, 0  # heap entry becomes stale
        return True

    def fail(self, task_id: str, now: int) -> str | None:
        """Returns AVAILABLE (requeued), DEAD, or None if not reserved."""
        self._expire(now)
        task = self._held(task_id)
        return None if task is None else self._release(task)

    def status(self, task_id: str, now: int) -> str | None:
        self._expire(now)
        task = self.tasks.get(task_id)
        return None if task is None else task.state

Complexity

Let n be the number of tasks and r the number of reservations so far.

  • add, complete, fail, status: O(1) plus the expiry work.
  • reserve: O(log r) for the heap push, plus the expiry work.
  • Expiry: each heap entry is pushed once and popped once, so the total over the whole run is O(r log r). That is O(log r) amortized per operation.
  • Space: O(n + r). Stale heap entries stay until their deadline passes. If that matters, rebuild the heap when stale entries outnumber live ones.

Test cases

With WorkQueue(lease_seconds=10, max_attempts=2):

  • add("a",0), add("b",0), add("c",1) -> all True; add("a",1) -> False (duplicate id).
  • reserve(2) -> "a"; fail("a",3) -> "AVAILABLE"; then reserve(4) -> "b", reserve(4) -> "c", reserve(5) -> "a" (a retried task goes to the tail).
  • reserve(5) with nothing available -> None.
  • complete("b",6) -> True; complete("b",7) -> False; status("b",7) -> "COMPLETED".
  • c was reserved at 4, so its deadline is 14: status("c",13) -> "RESERVED", status("c",14) -> "AVAILABLE" (deadline is inclusive). complete("c",14) -> False, because the worker's lease is gone.
  • a's second attempt was reserved at 5: status("a",15) -> "DEAD"; fail("a",15) -> None.
  • status("zzz",15) -> None (never added).
  • Same deadline: with lease 5, reserve x, y, z at time 1; at time 6 all three expire and reserve(6) returns x then y.
  • max_attempts=1: one fail sends the task straight to DEAD.
  • lease_seconds=0: a task reserved at time 0 is AVAILABLE again for any later call at time 0.
  • Stale heap entry: reserve at 0 (deadline 10), fail at 1, reserve again at 2 (deadline 12). status(t,10) -> "RESERVED", status(t,12) -> "AVAILABLE".

Follow-ups

  • Stale workers (fencing tokens). Right now any caller can complete a task by id, including a worker whose lease already expired and was handed to someone else, if that someone has not finished yet. Return the lease_id from reserve and require it in complete, fail and heartbeat. A mismatched token is rejected. This is the same idea as fencing tokens in distributed locks.
  • Heartbeats / lease extension. Add extend(task_id, lease_id, now) that pushes the deadline out. Push a new heap entry with the new deadline and leave the old one to be skipped lazily (store the current deadline on the task and compare it on pop, or issue a new lease id).
  • Retry backoff and delayed tasks. Instead of requeueing at the tail right away, put the task in a second heap keyed by now + base * 2**attempts (with jitter). Before each operation, move due tasks from that heap into ready. Scheduled add(..., run_at) uses the same heap.
  • Durability / crash recovery. Append every state change (add, reserve with deadline, complete, fail) to a write-ahead log and fsync before replying. On restart, replay the log to rebuild the dict, deque and heap; snapshot periodically so replay stays short. Leases that were live at the crash expire naturally on replay. Discuss that this gives at-least-once delivery, so task handlers must be idempotent.
  • Concurrency and scale. One lock around every public method is correct and simple; reserve is the hot path. To scale out, shard tasks by id hash across queues and let workers poll several shards, or move to a store with atomic compare-and-set on (state, lease_id). A dead-letter listing with redrive(task_id) (reset attempts, requeue) is a common last part.

Interviewer expectations

Mid-level (L4)

  • Gets the state machine right: add, reserve, complete, fail, status with correct return values.
  • Expires leases before each operation and treats expiry like a failure, including dead-lettering.
  • Reaches O(log n) operations with a heap, even if stale entries are first handled awkwardly.

Senior (L5)

  • Uses lease ids for lazy heap deletion and explains why an old deadline must not expire a newer reservation.
  • States the edge cases unprompted: inclusive deadline, tie order on equal deadlines, complete after expiry, duplicate add, max_attempts=1.
  • Writes focused tests and adds fencing tokens or heartbeats cleanly when asked.

Staff (L6+)

  • Frames the design as at-least-once delivery with visibility timeouts (as in SQS) and explains why consumers must be idempotent.
  • Covers durability (WAL, snapshot, replay), clock assumptions (monotonic time, skew between workers and the queue), and backoff with jitter.
  • Discusses sharding, lock contention on reserve, and how to operate a dead-letter queue (alerts, redrive).

Sign in to comment, vote or remix.