Design a task scheduler with dependencies and retries

01The design, in the six-step framework

1. Requirements and use cases

A user registers tasks and says which other tasks each one needs finished first. The scheduler runs a task only when every dependency has succeeded, runs higher-priority tasks before lower ones when several are ready, retries a failed task with exponential backoff up to a limit, and lets a caller subscribe to success, failure and retry events. A caller can block until the whole graph has settled.

Out of scope: persisting the graph across process restarts, distributing work across machines (that is the distributed job scheduler), and task cancellation.

Constraints that shape the design: several worker threads share one queue, so every state transition must be atomic; adding a task that would create a cycle must be rejected at add time, not discovered as a hang; and a backoff must not make a worker sleep while other tasks are ready.

2. Core entities and relationships

• Task: a unit of work with an id, a priority, an attempt counter and its current TaskState. It owns its own retry arithmetic.
• Scheduler: the orchestrator. It owns the dependency graph (two adjacency maps), the ready heap, the worker threads and the listener lists.
• ScheduledRun: a small value placed on the heap: when the task may run and with what priority. It exists so that a retry waiting on backoff and a fresh ready task live in one ordered structure.
• Worker: a thread loop that takes the next runnable entry and reports the outcome back.
• Listener: any callable that wants to hear about outcomes (an observer).

The scheduler composes the heap entries and the worker threads: they have no meaning outside it and die with it. It only aggregates tasks: the caller creates them and may keep references to inspect state afterwards. The dependency graph is stored twice, forward and reverse, because add needs the forward edges for the cycle check and completion needs the reverse edges to find who is unblocked (see graphs inside objects).

3. Interfaces and key methods

class TaskState(Enum):
    PENDING = auto(); READY = auto(); RUNNING = auto()
    SUCCEEDED = auto(); FAILED = auto(); RETRY_WAIT = auto()

@dataclass(order=True)
class ScheduledRun:
    """Heap entry: earliest not_before first, then higher priority."""
    not_before: float
    neg_priority: int
    task_id: str = field(compare=False)

class Task:
    def run(self) -> None: ...              # the work; raises on failure
    def can_retry(self) -> bool: ...        # attempts < max_attempts

class Scheduler:
    def add(self, task: Task, deps: list[str] = ()) -> None:
        """Register a task; raises CycleError if deps would form a cycle."""
    def start(self, workers: int = 4) -> None:
        """Begin dispatching READY tasks to a worker pool."""
    def on(self, event: str, cb: Callable[[Task], None]) -> None:
        """Subscribe to 'succeeded' / 'failed' / 'retry' notifications."""
    def wait(self) -> None:
        """Block until every task is SUCCEEDED or permanently FAILED."""

Invariants: a task is on the heap at most once; it is on the heap only when its unfinished-dependency set is empty; _open equals the number of tasks that are not SUCCEEDED or FAILED, so wait can return exactly when the graph settles.

4. Design patterns and the choices behind them

Readiness as a graph, ordering as a heap

What varies is which tasks are runnable and in what order. A topological sort computed once up front would be wrong as soon as a task fails and is retried later, so readiness is tracked incrementally: each task carries the set of dependencies still unfinished, and a completion removes itself from its dependents. Ordering among ready tasks is a priority queue keyed by (not_before, -priority), which gives backoff and priority in one structure. The rejected alternative was a separate timer thread for retries; it doubles the synchronisation surface for no gain.

An explicit state machine on the task

Six states and a handful of legal transitions make the concurrency rules checkable: only READY or RETRY_WAIT entries are on the heap, only RUNNING tasks can succeed or fail. Booleans like done and retrying were rejected because their combinations include impossible ones (see state machines).

Observer for outcomes

The caller wants to react to outcomes without the scheduler knowing why. Listeners are called outside the lock so a slow listener cannot stall workers, which is the usual trap with the observer pattern.

Injected clock

clock is a constructor argument so tests can advance time deterministically instead of sleeping, following designing for unit tests.

5. Core flow and code

1. The caller calls add(B, deps=["A"]). The scheduler checks that A cannot reach B, stores both adjacency directions and, because B has a dependency, leaves it PENDING.
2. add(A) has no deps, so A is pushed to the heap as READY and the condition variable is notified.
3. A worker in _next wakes, sees the heap top is due, pops it, marks A RUNNING and runs it outside the lock.
4. A raises. _on_failure re-pushes A with not_before = now + 0.5 s in state RETRY_WAIT and emits retry. Workers that wake see the top is not yet due and wait with a timeout equal to the remaining delay.
5. A succeeds on the second attempt. _on_success decrements _open, removes A from B's pending set, finds it empty and pushes B as READY.
6. B runs and succeeds; _open reaches zero and wait returns.

import heapq, threading, time
from dataclasses import dataclass, field
from enum import Enum, auto
from typing import Callable

class CycleError(Exception): pass

class TaskState(Enum):
    PENDING = auto(); READY = auto(); RUNNING = auto()
    SUCCEEDED = auto(); FAILED = auto(); RETRY_WAIT = auto()

@dataclass
class Task:
    task_id: str
    fn: Callable[[], None]
    priority: int = 0
    max_attempts: int = 3
    base_backoff: float = 0.5
    attempts: int = 0
    state: TaskState = TaskState.PENDING

    def run(self) -> None:
        self.attempts += 1
        self.fn()

    def can_retry(self) -> bool:
        return self.attempts < self.max_attempts

    def next_backoff(self) -> float:
        return self.base_backoff * (2 ** (self.attempts - 1))

@dataclass(order=True)
class ScheduledRun:
    not_before: float
    neg_priority: int
    task_id: str = field(compare=False)

class Scheduler:
    def __init__(self, clock: Callable[[], float] = time.monotonic) -> None:
        self._clock = clock
        self._tasks: dict[str, Task] = {}
        self._deps: dict[str, set[str]] = {}       # task -> unfinished deps
        self._dependents: dict[str, set[str]] = {} # dep -> tasks waiting on it
        self._heap: list[ScheduledRun] = []
        self._lock = threading.Lock()
        self._wake = threading.Condition(self._lock)
        self._listeners: dict[str, list[Callable[[Task], None]]] = {}
        self._open = 0                              # tasks not yet terminal

    def add(self, task: Task, deps: list[str] = ()) -> None:
        with self._lock:
            if self._reaches(deps, task.task_id):
                raise CycleError(task.task_id)
            self._tasks[task.task_id] = task
            self._deps[task.task_id] = set(deps)
            for d in deps:
                self._dependents.setdefault(d, set()).add(task.task_id)
            self._open += 1
            if not deps:
                self._make_ready(task)

    def _reaches(self, start: list[str], target: str) -> bool:
        stack = list(start); seen = set()
        while stack:
            cur = stack.pop()
            if cur == target: return True
            if cur in seen: continue
            seen.add(cur); stack.extend(self._deps.get(cur, ()))
        return False

    def _make_ready(self, task: Task, delay: float = 0.0) -> None:
        task.state = TaskState.READY if delay == 0 else TaskState.RETRY_WAIT
        heapq.heappush(self._heap, ScheduledRun(self._clock() + delay, -task.priority, task.task_id))
        self._wake.notify_all()

    def on(self, event: str, cb: Callable[[Task], None]) -> None:
        self._listeners.setdefault(event, []).append(cb)

    def _emit(self, event: str, task: Task) -> None:
        for cb in self._listeners.get(event, []):
            cb(task)

    def start(self, workers: int = 4) -> None:
        for _ in range(workers):
            threading.Thread(target=self._worker, daemon=True).start()

    def _next(self) -> Task:
        with self._wake:
            while True:
                if self._heap and self._heap[0].not_before <= self._clock():
                    entry = heapq.heappop(self._heap)
                    task = self._tasks[entry.task_id]
                    task.state = TaskState.RUNNING
                    return task
                timeout = (self._heap[0].not_before - self._clock()) if self._heap else None
                self._wake.wait(timeout)

    def _worker(self) -> None:
        while True:
            task = self._next()
            try:
                task.run()
            except Exception:
                self._on_failure(task)
            else:
                self._on_success(task)

    def _on_success(self, task: Task) -> None:
        with self._lock:
            task.state = TaskState.SUCCEEDED
            self._open -= 1
            for child in self._dependents.get(task.task_id, ()):
                self._deps[child].discard(task.task_id)
                if not self._deps[child]:
                    self._make_ready(self._tasks[child])
            self._wake.notify_all()
        self._emit("succeeded", task)

    def _on_failure(self, task: Task) -> None:
        with self._lock:
            if task.can_retry():
                self._make_ready(task, task.next_backoff())
                event = "retry"
            else:
                task.state = TaskState.FAILED
                self._open -= 1
                event = "failed"
            self._wake.notify_all()
        self._emit(event, task)

    def wait(self) -> None:
        with self._wake:
            while self._open > 0:
                self._wake.wait(0.1)

6. Extensions, concurrency and what interviewers probe

• Exponential backoff with jitter: add a random fraction in next_backoff; the heap ordering already handles arbitrary delays, nothing else changes.
• A failed dependency: today dependents stay PENDING forever and wait never returns. The reported follow-up is to mark them SKIPPED and decrement _open for each; that is one more state and a reverse-graph walk in _on_failure.
• Where the lock matters: pop-and-mark-RUNNING is one critical section, otherwise two workers can pop the same entry; completion and unblocking are another, otherwise a dependent can be pushed twice. Running the task itself is outside the lock, or the pool serialises.
• Cycle detection at add time is a depth-first reachability check from the new deps to the new task; interviewers ask what happens with a self-dependency (caught, since the start list contains the target).
• Test first: a diamond graph (A before B and C, both before D) with a fake clock, asserting D runs once after both; then a task that fails twice and succeeds, asserting the two retry events and their timestamps.
• Common mistake: computing a topological order once and iterating it, which cannot express retries, priorities among ready tasks, or tasks added while the scheduler runs.

02Where this question was reported

• LLD Mastery · Google-tagged (Hard); DAG dependencies, exponential backoff
• LLD Mastery · Google-tagged with Amazon/Microsoft/Meta/Uber; topological sort, retries
• PracHub · PracHub Google-tagged
• Codemia · Codemia Advanced 'OS Task Scheduler'
• thita.ai · thita.ai LLD sheet, advanced