Design a concurrent priority task scheduler

01The design, in the six-step framework

1. Requirements and use cases

A client submits a function with a priority and gets back a handle. Workers run the highest-priority task first and, among equal priorities, the oldest first. A client can cancel a task that has not started, and a task that waits longer than its timeout is cancelled by the scheduler itself. A client can block on the result. Shutdown stops intake and lets running tasks finish.

Out of scope: dependencies between tasks (see the dependency scheduler), persistence, and interrupting a task that is already running, which Python threads cannot do safely.

Constraints: many worker threads and many submitting threads touch the queue at once, so no task may run twice or be lost; a cancelled or expired task must never start; a worker must not spin while the queue is empty.

2. Core entities and relationships

• PriorityScheduler: owns the heap, the condition variable and the worker threads; the only class with a lock around shared state.
• Entry: the heap record: priority, a sequence number for FIFO ties, the function and its handle.
• Handle: the caller's view of one submission: a small state machine (RunState) plus the result slot and its own condition variable, so waiting on one result does not block the scheduler.
• Priority: an ordered enum, converted to a negative integer so the min-heap pops HIGH first.

The scheduler composes entries and workers. Entries only aggregate handles because the caller holds the handle too; that shared ownership is exactly why the claim step lives on the handle and not on the heap.

3. Interfaces and key methods

class Priority(IntEnum):
    LOW = 0; NORMAL = 1; HIGH = 2

class Handle:
    """What submit() returns: lets the caller cancel or wait."""
    def cancel(self) -> bool: ...        # True iff cancelled before it started
    def result(self, timeout: float | None = None) -> object: ...
    def done(self) -> bool: ...

class PriorityScheduler:
    def __init__(self, workers: int, clock: Callable[[], float] = time.monotonic) -> None: ...
    def submit(self, fn: Callable[[], object], priority: Priority, timeout: float | None = None) -> Handle:
        """Enqueue fn; a task not started within `timeout` seconds is cancelled."""
    def shutdown(self, wait: bool = True) -> None:
        """Stop accepting work; optionally wait for running tasks."""

Invariants: a handle moves QUEUED to RUNNING exactly once and only through _claim; CANCELLED is reachable only from QUEUED; the heap may contain entries whose handle is already CANCELLED, and the worker skips them, so cancel is O(1) and never searches the heap.

4. Design patterns and the choices behind them

Claim-before-run, not remove-on-cancel

The force is the race between a client cancelling and a worker starting the same task. Removing an entry from the middle of a heap is O(n) and must happen under the scheduler lock; instead the handle owns a tiny compare-and-set, _claim, and both sides go through it. Whoever moves the state first wins, so a task is never both cancelled and run. This is the single-process version of idempotent operations: a second cancel or a second claim is a harmless no-op returning False.

One condition variable for the queue, one per handle

Two different things are waited on: workers wait for work, callers wait for results. Sharing one lock would wake every waiting caller on every submit. The scheduler uses a producer-consumer condition on the heap; each handle has its own condition for result. The alternative, one global lock, was rejected in the reported interview as fully serialising the pool.

Priority plus sequence in the heap key

A heap keyed on priority alone makes equal-priority tasks pop in arbitrary order, which looks like starvation of early submitters. Adding a monotonically increasing sequence number gives FIFO within a priority at zero cost.

Timeout evaluated at dequeue

Expiry is checked when a worker pops an entry, with an injected clock. A timer thread per task was rejected: thousands of timers for thousands of tasks, and the same race with starting. The cost is that an expired task is noticed only when it reaches the top; a sweep every few seconds fixes that if the caller needs prompt cancellation events.

5. Core flow and code

1. submit(fn, HIGH, timeout=2) creates a Handle with a deadline, pushes an Entry under the lock and notifies one waiting worker.
2. A worker in _take pops the entry, checks the deadline, then calls handle._claim(). It returns True, so the worker leaves the lock holding the entry.
3. Meanwhile the client calls handle.cancel(); the state is RUNNING, so it returns False.
4. The worker runs fn outside every lock and calls _finish with the value or the exception, which notifies waiters on that handle only.
5. The client's result() wakes, sees DONE and returns the value, or re-raises the task's exception.
6. shutdown() flips _closed, wakes all workers; each drains what remains and exits when the heap is empty.

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

class Priority(IntEnum):
    LOW = 0; NORMAL = 1; HIGH = 2

class RunState(Enum):
    QUEUED = auto(); RUNNING = auto(); DONE = auto(); CANCELLED = auto(); FAILED = auto()

class Handle:
    def __init__(self, deadline: float | None) -> None:
        self.state = RunState.QUEUED
        self.deadline = deadline
        self._value: object = None
        self._error: BaseException | None = None
        self._cv = threading.Condition()

    def cancel(self) -> bool:
        with self._cv:
            if self.state is not RunState.QUEUED:
                return False
            self.state = RunState.CANCELLED
            self._cv.notify_all()
            return True

    def _claim(self) -> bool:
        """Worker calls this: QUEUED -> RUNNING exactly once."""
        with self._cv:
            if self.state is not RunState.QUEUED:
                return False
            self.state = RunState.RUNNING
            return True

    def _finish(self, value: object = None, error: BaseException | None = None) -> None:
        with self._cv:
            self._value, self._error = value, error
            self.state = RunState.FAILED if error else RunState.DONE
            self._cv.notify_all()

    def done(self) -> bool:
        return self.state in (RunState.DONE, RunState.CANCELLED, RunState.FAILED)

    def result(self, timeout: float | None = None) -> object:
        with self._cv:
            if not self._cv.wait_for(self.done, timeout):
                raise TimeoutError("still pending")
            if self.state is RunState.CANCELLED:
                raise RuntimeError("cancelled")
            if self._error:
                raise self._error
            return self._value

@dataclass(order=True)
class Entry:
    neg_priority: int
    seq: int                                  # FIFO among equal priorities
    fn: Callable[[], object] = field(compare=False)
    handle: Handle = field(compare=False)

class PriorityScheduler:
    def __init__(self, workers: int, clock: Callable[[], float] = time.monotonic) -> None:
        self._clock = clock
        self._heap: list[Entry] = []
        self._seq = itertools.count()
        self._cv = threading.Condition()
        self._closed = False
        self._threads = [threading.Thread(target=self._loop, daemon=True) for _ in range(workers)]
        for t in self._threads:
            t.start()

    def submit(self, fn: Callable[[], object], priority: Priority, timeout: float | None = None) -> Handle:
        deadline = None if timeout is None else self._clock() + timeout
        handle = Handle(deadline)
        with self._cv:
            if self._closed:
                raise RuntimeError("scheduler is shut down")
            heapq.heappush(self._heap, Entry(-int(priority), next(self._seq), fn, handle))
            self._cv.notify()
        return handle

    def _take(self) -> Entry | None:
        with self._cv:
            while True:
                if self._closed and not self._heap:
                    return None
                while self._heap:
                    entry = heapq.heappop(self._heap)
                    h = entry.handle
                    if h.deadline is not None and self._clock() > h.deadline:
                        h.cancel()                      # expired while queued
                        continue
                    if h._claim():                      # skips user-cancelled entries
                        return entry
                self._cv.wait(0.05)

    def _loop(self) -> None:
        while (entry := self._take()) is not None:
            try:
                entry.handle._finish(value=entry.fn())
            except BaseException as e:               # noqa: BLE001 - report, never kill the worker
                entry.handle._finish(error=e)

    def shutdown(self, wait: bool = True) -> None:
        with self._cv:
            self._closed = True
            self._cv.notify_all()
        if wait:
            for t in self._threads:
                t.join()

6. Extensions, concurrency and what interviewers probe

• Cooldown between runs of the same key (a reported variant): add a not_before field to Entry keyed on the task's group, and let _take re-push entries that are not yet due; the heap key becomes (not_before, neg_priority, seq).
• Exactly-once under worker crash: a worker that dies mid-task leaves the handle RUNNING forever. Wrapping fn() in try/finally that always calls _finish covers exceptions; a hard process crash needs persistence, which is out of scope and worth saying.
• Where the atomic step is: _claim. Everything else can tolerate a stale read. Interviewers in the reported round pushed on exactly this: show the two interleavings and why both end with one winner.
• Priority starvation: a steady stream of HIGH tasks starves LOW ones. The fix is ageing: boost priority by elapsed queue time at pop; mention it, do not build it unless asked.
• Test first: a fake clock, one worker, three tasks submitted LOW, HIGH, NORMAL; assert run order HIGH, NORMAL, LOW; then a task with timeout 1 submitted behind a 5-second task, assert it ends CANCELLED and never runs.
• Common mistake: holding the scheduler lock while running the task, which turns N workers into one.

02Where this question was reported

• FastPrep · Priority queue + ReentrantLock; getting the concurrency rules right under time pressure was the hard part
• algoroq.io · 'Design a task scheduler that respects cooldown constraints' listed as commonly asked
• LLD Mastery · Commercial tag: task scheduler with DAG dependencies and exponential backoff; task management system (state transitions, observer)