LLD: Task Scheduler¶

Difficulty: Advanced | Time: 45–60 minutes

1. Problem Statement¶

Design an in-process task scheduler. Callers submit tasks with a priority and/or a scheduled execution time — run now, run at a specific timestamp, or run after a delay. A fixed-size pool of worker threads executes submitted tasks concurrently. Among tasks that are currently eligible to run, higher-priority tasks execute before lower-priority ones. As a stretch goal, support recurring tasks that re-fire every N seconds until cancelled.

This is the single-process version of the problem. Distributed Job Scheduler covers the same idea at the system-design layer — leader election, leases, exactly-once execution across machines. Here, the constraint is purely about data-structure and concurrency correctness inside one process.

2. Requirements¶

Functional (in scope):

• submit(task, priority, run_at=None, delay=None) — schedule a task for immediate or future execution at a given priority
• A fixed-size worker pool (N threads) pulls and executes tasks
• Among tasks whose scheduled time has arrived, the highest-priority task runs next
• cancel(task_id) — prevent a not-yet-started task from running
• Stretch: recurring tasks (interval_seconds) that re-enqueue themselves after each successful run, until cancelled

Explicitly out of scope for v1: distributed/cross-machine scheduling, leader election, durability across process restarts, and exactly-once execution guarantees when a worker crashes mid-task — all covered by Distributed Job Scheduler, which picks up exactly where this one stops: single process → many machines.

3. Entities¶

The nouns in the problem statement: Task, TaskQueue (the priority/delay-aware ready queue), Scheduler (the promoter that moves due tasks into the ready queue), WorkerPool / Worker, SchedulingStrategy.

4. Class Design¶

classDiagram
    class Task {
        +task_id: str
        +priority: int
        +run_at: datetime
        +fn: Callable
        +interval_seconds: float
        +cancelled: bool
        +submitted_seq: int
    }

    class SchedulingStrategy {
        <<interface>>
        +sort_key(task)* tuple
    }
    class PriorityFirstStrategy
    class EarliestDeadlineFirstStrategy
    class FifoStrategy

    class TaskQueue {
        -heap: List~tuple~
        -lock: Lock
        -not_empty: Condition
        -strategy: SchedulingStrategy
        +push(task) void
        +pop_ready(timeout) Task
        +cancel(task_id) void
    }

    class Scheduler {
        -queue: TaskQueue
        -pending: List~Task~
        -promoter_thread: Thread
        +submit(task) str
        +cancel(task_id) void
        -_promote_due_tasks() void
    }

    class WorkerPool {
        -workers: List~Worker~
        -queue: TaskQueue
        +start() void
        +shutdown() void
    }

    class Worker {
        +run() void
    }

    Scheduler --> TaskQueue : promotes due tasks into
    Scheduler --> SchedulingStrategy : configured with
    TaskQueue --> SchedulingStrategy : uses (injected)
    WorkerPool *-- Worker : composition
    WorkerPool --> TaskQueue : consumes from
    TaskQueue o-- Task : holds references while ready/pending

stateDiagram-v2
    [*] --> Pending : submit() [run_at in future]
    [*] --> ReadyQueue : submit() [already due]
    Pending --> ReadyQueue : promoter thread\npromotes at run_at
    Pending --> Cancelled : cancel() [lazy flag]
    ReadyQueue --> Cancelled : cancel() [lazy flag]
    ReadyQueue --> Running : worker pop_ready()
    Cancelled --> [*] : skipped on pop (lazy deletion)
    Running --> Done : fn() returns
    Running --> Done : fn() raises\n(isolated, logged)
    Done --> Pending : reschedule_recurring()\n[interval_seconds set]
    Done --> [*] : [not recurring]

Why a min-heap keyed on (-priority, run_at, submitted_seq) in the ready queue, and not a plain FIFO queue or a plain priority queue alone:

• A plain FIFO queue loses priority entirely — a priority=10 task submitted after a priority=1 task would still run second. That directly violates "higher-priority tasks execute before lower-priority ones."
• A plain priority queue keyed on priority alone, with future tasks mixed in, loses time — it would happily pop a run_at two hours from now ahead of a lower-priority task that's due immediately. Priority only makes sense to compare among tasks that are already due — comparing priorities across "due now" and "due in two hours" is comparing the wrong axis.
• The fix is two structures, not a cleverer single key: the promoter holds not-yet-due tasks in a time-ordered pending heap and only pushes a task into the ready queue once run_at has arrived. The ready heap then orders by (-priority, run_at, submitted_seq) so among due tasks, higher priority wins — including a just-due high-priority task beating an overdue low-priority one. That matches the functional requirement. EarliestDeadlineFirstStrategy is the opt-in if you want time-first among due tasks instead.
• That still leaves a second problem a single heap doesn't solve on its own: a task due two hours from now sitting in the same heap as tasks due right now must not be poppable early just because it's the smallest item once its neighbors finish. The heap's job is ordering candidates for the moment they're actually eligible, not gating eligibility. This design keeps that gate explicit — see the single-threaded promoter in Core Code — rather than baking "am I due yet" into every worker's pop logic.

Why WorkerPool *-- Worker is composition: workers have no identity or lifecycle outside the pool that spawned them — shut down the pool, the worker threads terminate. Why TaskQueue o-- Task is aggregation: a submitted Task is a caller-owned unit of work the queue tracks and hands off, not something the queue's lifecycle defines — see Class Relationships.

5. Patterns Applied¶

• Strategy for SchedulingStrategy — "priority order" vs. "pure FIFO" vs. "earliest-deadline-first" are named, real variation points (the problem statement itself names two of them), so TaskQueue depends on the sort_key() interface rather than a hardcoded comparison, and a new policy is a new class with zero edits to TaskQueue. See Design Patterns.
• Producer-consumer, the load-bearing concurrency pattern here (not a GoF pattern — a concurrency idiom): submitting threads are producers pushing Tasks onto the shared queue, worker threads are consumers popping them off. The entire Concurrency section below is this pattern done correctly — bounded coordination through one shared, lock-protected structure with condition-variable signaling, instead of workers polling or submitters blocking on worker availability.
• Explicitly not using Observer for "notify on task completion" even though it's tempting — nothing in the problem statement asks for multiple interested listeners per task; a plain callback or Future-style result object covers it without the indirection. Don't add a pattern the requirements don't ask for.

6. Core Code¶

Two structural decisions worth stating before the code: (1) the ready queue is heapq protected by a Lock + Condition rather than queue.PriorityQueue — shown explicitly here so the wait/notify mechanics are visible; queue.PriorityQueue is a thin wrapper around exactly this (its own Lock, two Conditions for not-full/not-empty, a heapq list) and is the right call in production code once you don't need to explain what's inside it. (2) Delayed tasks are held in a separate min-heap of not-yet-due tasks, promoted into the ready queue by a single dedicated promoter thread that sleeps until the next due time — the sched-module approach — rather than encoding due-time directly into the ready heap and having every worker re-check it. A single promoter means the "is it due yet" logic exists in exactly one place, sleeping precisely until it next has work, instead of N workers each re-deriving that logic and busy-checking.

from __future__ import annotations

import heapq
import itertools
import threading
import time
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from typing import Callable, Protocol


# ---------- Task ----------

@dataclass
class Task:
    fn: Callable[[], None]
    priority: int = 0                 # higher runs first among due tasks
    run_at: datetime = field(default_factory=datetime.now)
    interval_seconds: float | None = None   # stretch: recurring tasks
    task_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    submitted_seq: int = 0            # tiebreak — assigned at submission, see Edge Cases
    cancelled: bool = False           # lazy-deletion flag, see Edge Cases


# ---------- Scheduling strategy ----------

class SchedulingStrategy(Protocol):
    def sort_key(self, task: Task) -> tuple: ...


class PriorityFirstStrategy:
    """Among due tasks, highest priority first; then earlier run_at; then submission order."""
    def sort_key(self, task: Task) -> tuple:
        return (-task.priority, task.run_at, task.submitted_seq)


class FifoStrategy:
    """Ignore priority entirely — pure submission order among due tasks."""
    def sort_key(self, task: Task) -> tuple:
        return (task.submitted_seq,)


class EarliestDeadlineFirstStrategy:
    """Due time is the only axis that matters; priority is a pure tiebreak."""
    def sort_key(self, task: Task) -> tuple:
        return (task.run_at, task.submitted_seq, -task.priority)


# ---------- Ready queue: thread-safe, condition-variable-backed ----------

class TaskQueue:
    """Holds tasks that are currently eligible to run, ordered by SchedulingStrategy."""

    def __init__(self, strategy: SchedulingStrategy):
        self._strategy = strategy
        self._heap: list[tuple[tuple, Task]] = []
        self._lock = threading.Lock()
        self._not_empty = threading.Condition(self._lock)

    def push(self, task: Task) -> None:
        with self._not_empty:                      # acquire, then notify under the same lock
            heapq.heappush(self._heap, (self._strategy.sort_key(task), task))
            self._not_empty.notify()                # wake exactly one waiting worker

    def pop_ready(self, timeout: float | None = None) -> Task | None:
        """Block until a non-cancelled task is available, or timeout elapses."""
        with self._not_empty:
            deadline = None if timeout is None else time.monotonic() + timeout
            while True:
                while not self._heap:
                    remaining = None if deadline is None else deadline - time.monotonic()
                    if remaining is not None and remaining <= 0:
                        return None
                    self._not_empty.wait(remaining)  # releases lock while blocked; no CPU spin
                _, task = heapq.heappop(self._heap)
                if task.cancelled:                   # lazy deletion — skip and keep looping
                    continue
                return task

    def cancel(self, task_id: str) -> bool:
        """O(1) lazy cancellation: mark and let pop_ready skip it — see Edge Cases."""
        with self._lock:
            for _, task in self._heap:
                if task.task_id == task_id:
                    task.cancelled = True
                    return True
            return False


# ---------- Scheduler: owns the not-yet-due tasks and promotes them ----------

class Scheduler:
    """Holds tasks scheduled for the future and promotes them into the ready
    TaskQueue exactly when they become due, via one dedicated promoter thread."""

    def __init__(self, ready_queue: TaskQueue):
        self._ready_queue = ready_queue
        self._pending: list[tuple[datetime, int, Task]] = []  # sequence prevents Task comparison on equal run_at
        self._lock = threading.Lock()
        self._wakeup = threading.Condition(self._lock)
        self._seq_counter = itertools.count()
        self._all_tasks: dict[str, Task] = {}
        self._running = True
        self._promoter = threading.Thread(target=self._promote_due_tasks, daemon=True)
        self._promoter.start()

    def submit(self, fn: Callable[[], None], priority: int = 0,
               run_at: datetime | None = None, delay: float | None = None,
               interval_seconds: float | None = None) -> str:
        if run_at is None:
            run_at = datetime.now() + timedelta(seconds=delay or 0)
        task = Task(fn=fn, priority=priority, run_at=run_at,
                    interval_seconds=interval_seconds,
                    submitted_seq=next(self._seq_counter))
        with self._wakeup:
            self._all_tasks[task.task_id] = task
            if run_at <= datetime.now():
                self._ready_queue.push(task)          # already due — skip the pending heap
            else:
                heapq.heappush(self._pending, (run_at, task.submitted_seq, task))
                self._wakeup.notify()                  # promoter may need to wake earlier now
        return task.task_id

    def cancel(self, task_id: str) -> bool:
        task = self._all_tasks.get(task_id)
        if task is None:
            return False
        task.cancelled = True                          # lazy — cheap regardless of which heap it's in
        return True

    def reschedule_recurring(self, task: Task) -> None:
        """Called by a worker after a recurring task completes successfully."""
        if task.cancelled or task.interval_seconds is None:
            return
        next_run = datetime.now() + timedelta(seconds=task.interval_seconds)
        self.submit(task.fn, task.priority, run_at=next_run,
                    interval_seconds=task.interval_seconds)

    def shutdown(self) -> None:
        with self._wakeup:
            self._running = False
            self._wakeup.notify()
        self._promoter.join(timeout=1.0)

    def _promote_due_tasks(self) -> None:
        """Single promoter thread — sleeps exactly until the next task is due,
        instead of every worker polling 'am I due yet' independently."""
        with self._wakeup:
            while self._running:
                if not self._pending:
                    self._wakeup.wait()                 # nothing scheduled — sleep until submit() wakes us
                    continue
                run_at, _, _ = self._pending[0]
                remaining = (run_at - datetime.now()).total_seconds()
                if remaining > 0:
                    self._wakeup.wait(remaining)         # sleep exactly until the next due time
                    continue
                _, _, task = heapq.heappop(self._pending)
                if not task.cancelled:
                    self._ready_queue.push(task)


# ---------- Worker pool: producer-consumer, N concurrent consumers ----------

class WorkerPool:
    def __init__(self, ready_queue: TaskQueue, scheduler: Scheduler, num_workers: int):
        self._queue = ready_queue
        self._scheduler = scheduler
        self._threads = [threading.Thread(target=self._worker_loop, daemon=True)
                          for _ in range(num_workers)]
        self._running = True

    def start(self) -> None:
        for t in self._threads:
            t.start()

    def shutdown(self) -> None:
        self._running = False
        for t in self._threads:
            self._queue.push(Task(fn=lambda: None, run_at=datetime.min))  # unblock waiters
        for t in self._threads:
            t.join(timeout=1.0)

    def _worker_loop(self) -> None:
        while self._running:
            task = self._queue.pop_ready(timeout=0.5)
            if task is None or task.cancelled:
                continue
            try:
                task.fn()                                # isolate failures — see Edge Cases
            except Exception:
                pass                                       # a task's exception must not kill the worker
            else:
                if task.interval_seconds is not None:
                    self._scheduler.reschedule_recurring(task)

7. Edge Cases¶

8. Concurrency¶

This is the exercise's crux: multiple submitter threads calling submit() concurrently, one promoter thread moving due tasks from the pending heap into the ready queue, and N worker threads concurrently popping from that same ready queue — three populations of threads touching two shared heaps.

Why a Condition, not polling. A naive worker loop —

while True:
    if not queue.empty():
        task = queue.pop()
        ...
    else:
        time.sleep(0.1)

— has two real costs. First, wasted CPU: every idle worker wakes up 10 times a second forever, even when the queue is empty for hours, multiplied by every worker in the pool. Second, latency: a task that becomes ready right after a worker just checked and found nothing waits up to a full poll interval before any worker notices it — and that interval is a tension between "responsive" and "not burning CPU," with no value that's good at both. TaskQueue.pop_ready() avoids both: Condition.wait() blocks the thread with zero CPU usage while releasing the lock, and push() calls notify() at the moment a task becomes available, waking exactly one blocked worker immediately rather than after some polling delay. This is the same wait/notify discipline covered in Concurrency Basics — Locks.

Why the pending-heap promoter also uses a Condition and not a fixed-interval sleep. _promote_due_tasks() computes exactly how long until the next pending task is due and calls wait(remaining) — it sleeps precisely that long, not "poll every second and see." Crucially, submit() calls notify() after pushing a new task into _pending, which matters when the new task is due sooner than whatever the promoter was already sleeping toward: without that notify, a task submitted with run_at five seconds from now could sit unpromoted for up to however long the promoter's previous sleep had left to run.

Why notify(), not notify_all(), in push(). Exactly one task became available, so exactly one blocked worker needs to wake and claim it; waking all of them would have every worker but one immediately re-check the heap, find it empty (or find the one item already claimed), and go back to waiting — correct, but pure wasted wakeup for a pool of any size. This is the "efficiency," not just "correctness," half of the thread-safety story — see Concurrency Basics — Thread Safety Without Locks for the broader point that a correct-but-inefficient synchronization scheme is still a design smell worth naming out loud in an interview.

Why the ready-queue lock and the pending-heap lock are separate. TaskQueue and Scheduler each own their own Lock/Condition pair rather than sharing one. A submitter pushing a future task only ever touches the pending heap and its lock; a worker popping a ready task only ever touches the ready queue and its lock. Sharing one lock across both would serialize submitters against workers for no reason — they operate on different data. The only thread that legitimately needs both is the promoter, and it always acquires the pending lock (implicitly, as the owner of _wakeup) before calling ready_queue.push(), which acquires the ready lock — a single fixed acquisition order, so there's no two-lock deadlock risk of the kind discussed in Concurrency Basics — Deadlocks.

9. Extensibility¶

Interview Questions¶

Key Takeaways¶

Previous: Pub/Sub | Next: LLD Problem Roadmap