Task scheduler

Design a scheduler that runs one-off and recurring tasks at their due time on a pool of workers, with retries and cancel. The key idea is a priority queue ordered by due time and workers that sleep until the next task is due.

Design an in-process task scheduler. Callers schedule a task to run once at a given time, or to repeat every fixed period. A pool of N worker threads runs the tasks when they are due. A failed task is retried, and a task that has not started can be cancelled.

Questions to ask first

• How big is it? Assume one process, N worker threads (4 to 16) and up to 100,000 waiting tasks.
• What kinds of tasks? Assume one-off tasks and recurring tasks with a fixed period such as every 10 minutes. Cron expressions are out of scope.
• What if every worker is busy when a task is due? Assume it waits in the queue and starts when a worker is free. It runs late, but it is never dropped.
• What counts as a failure, and what do we do? Assume the job returns false or throws. A one-off task is retried up to a limit with a fixed delay, then marked FAILED.
• Can we cancel a task? Assume yes until it starts. A task that is already running is not interrupted.
• Where does the time come from? Assume a clock that the caller can replace, so a test can move time forward without sleeping.

Requirements

• Schedule a one-off task for a due time, and a recurring task from a first time and then every period.
• Run due tasks on a pool of N workers, the earliest due task first. Tasks with the same due time run in the order they were scheduled.
• Retry a failed one-off task after a delay, up to a limit. After the limit its state is FAILED.
• Cancel a task that has not started. A task that is running or finished cannot be cancelled.
• Do not poll. An idle worker sleeps until the earliest task is due, or until a new task with an earlier due time arrives.
• Read time from a clock that can be replaced, so tests are fast and do not sleep.

Out of scope

• Cron expressions and calendar rules.
• Saving tasks to disk and recovering after a restart.
• Running one task on several machines.
• Stopping a job that is already running.

Entities

• Scheduler: schedule(task), cancel(id); owns the workers
• Task: id, dueTime, job, recurrence, retry policy, state; tryStart(), tryCancel()
• Job: the work; run(now) returns true when it succeeded
• Recurrence: nextDue(dueTime); OneTime and FixedRate
• RetryPolicy: delayAfterFailure(failures); FixedDelay
• TaskQueue: a priority queue by due time; add(task), take() blocks until a task is due
• TaskState: WAITING, RUNNING, DONE, FAILED, CANCELLED

Classes in the model design

• Scheduler (class)
• TaskQueue (class)
• Task (class)
• TaskState (enum)
• Job (interface)
• Recurrence (interface)
• OneTime (class)
• FixedRate (class)
• RetryPolicy (interface)
• FixedDelay (class)

The main use case, step by step

1. A caller schedules the task `report` for 9:00 with `scheduler.schedule(task)`. The scheduler stores it by id, which is why an id can be used only once, and adds it to the `TaskQueue`.
2. The queue is a `PriorityQueue` ordered by due time, then by a sequence number that keeps the order of equal times. Four workers are asleep inside `take()`, waiting with `awaitNanos` until the earliest task is due.
3. The new task is the new earliest, so `add` signals the condition. The workers wake up, read the head again, and go back to sleep until 9:00. Nobody polls.
4. At 9:00 one worker finds the head due and removes it. It calls `tryStart()`, which moves the state from WAITING to RUNNING only if nobody cancelled it. Then it calls `job.run(now)`.
5. The job returns true. `finished` asks the task's `Recurrence` for the next due time. `FixedRate` answers 9:00 plus one period, so the task gets that due time, goes back to WAITING and into the queue. No thread waits for it.
6. If the job throws, a one-off task asks its `RetryPolicy`. `FixedDelay` says to try again in 5 minutes, so the task is queued again. After the last retry the policy returns -1 and the state becomes FAILED.
7. A cancel at 8:59 calls `tryCancel()`, which changes WAITING to CANCELLED. A worker that later removes the task from the queue sees that `tryStart()` returns false and skips it. A cancel at 9:01 returns false because the task is already RUNNING.

Patterns used

• Producer-consumer with a blocking queue: Callers of schedule are the producers. The worker threads are the consumers of TaskQueue.. The two sides never call each other. The queue holds the work and tells the workers when it is time, so a slow job cannot block a caller.
• Strategy: Recurrence (OneTime, FixedRate) and RetryPolicy (FixedDelay).. Repeating and retrying are rules that change: every day at 9, with backoff, with a limit. Each rule is a small class, and Scheduler and Task do not change when a new one arrives.
• Command: Job is the work as an object that the scheduler can store, delay and run later.. The scheduler does not know what a job does. It only knows when to run it and what to do with the answer.
• State (as an enum with guarded moves): TaskState, changed only through tryStart, tryCancel and the scheduler.. Cancel and start race against each other. Both are synchronized moves out of WAITING, so exactly one of them wins.

Key code

Job, Recurrence and RetryPolicy: the rules that change

public enum TaskState { WAITING, RUNNING, DONE, FAILED, CANCELLED }

/** The work. Return true when it succeeded. False or an exception means it failed. */
public interface Job {
    boolean run(long now);
}

/** What happens after a run: when is the next one due? */
public interface Recurrence {
    long nextDue(long dueTime);        // -1 means there is no next run
}

public final class OneTime implements Recurrence {
    @Override public long nextDue(long dueTime) { return -1; }
}

public final class FixedRate implements Recurrence {
    private final long period;
    public FixedRate(long period) { this.period = period; }
    @Override public long nextDue(long dueTime) { return dueTime + period; }   // no drift: counted from the due time
}

public interface RetryPolicy {
    long delayAfterFailure(int failures);   // -1 means give up
}

public final class FixedDelay implements RetryPolicy {
    private final int maxRetries;
    private final long delay;
    public FixedDelay(int maxRetries, long delay) { this.maxRetries = maxRetries; this.delay = delay; }
    @Override public long delayAfterFailure(int failures) { return failures <= maxRetries ? delay : -1; }
}

Task: one state with two guarded moves

public final class Task implements Comparable<Task> {
    private final String id;
    private final Job job;
    private final Recurrence recurrence;
    private final RetryPolicy retry;
    private long dueTime;
    private long sequence;                  // equal due times run in the order they were added
    private int failures;
    private TaskState state = TaskState.WAITING;

    public Task(String id, long dueTime, Job job, Recurrence recurrence, RetryPolicy retry) {
        this.id = id; this.dueTime = dueTime; this.job = job; this.recurrence = recurrence; this.retry = retry;
    }

    @Override public int compareTo(Task o) {
        return dueTime != o.dueTime ? Long.compare(dueTime, o.dueTime) : Long.compare(sequence, o.sequence);
    }

    /** Cancel and start both go through here, so a task cannot be cancelled and run at once. */
    public synchronized boolean tryStart() {
        if (state != TaskState.WAITING) return false;
        state = TaskState.RUNNING;
        return true;
    }

    public synchronized boolean tryCancel() {
        if (state != TaskState.WAITING) return false;
        state = TaskState.CANCELLED;
        return true;
    }
    // getters, setSequence, setState, countFailure and reopen(due), which sets the due time and WAITING, are omitted
}

TaskQueue: a priority queue and a condition variable

/** The waiting room: workers sleep until the earliest task is due. No polling. */
public final class TaskQueue {
    private final PriorityQueue<Task> heap = new PriorityQueue<>();
    private final ReentrantLock lock = new ReentrantLock();
    private final Condition headChanged = lock.newCondition();
    private final LongSupplier clock;             // milliseconds; a test passes its own
    private long sequence;
    private boolean closed;

    public TaskQueue(LongSupplier clock) { this.clock = clock; }

    public void add(Task t) {
        lock.lock();
        try {
            t.setSequence(sequence++);
            heap.add(t);
            if (heap.peek() == t) headChanged.signalAll();   // a new earliest task: sleepers must look again
        } finally {
            lock.unlock();
        }
    }

    /** Blocks until the earliest task is due, then removes it. Returns null after close(). */
    public Task take() throws InterruptedException {
        lock.lock();
        try {
            while (!closed) {
                Task head = heap.peek();
                if (head == null) {
                    headChanged.await();                      // nothing waiting: sleep until add() or close()
                } else {
                    long wait = head.dueTime() - clock.getAsLong();
                    if (wait <= 0) return heap.poll();
                    headChanged.awaitNanos(TimeUnit.MILLISECONDS.toNanos(wait));   // sleep until it is due
                }
            }
            return null;
        } finally {
            lock.unlock();
        }
    }

    public void close() {
        lock.lock();
        try { closed = true; headChanged.signalAll(); } finally { lock.unlock(); }
    }
}

Scheduler: workers, cancel and what happens after a run

public final class Scheduler {
    private final Map<String, Task> tasks = new ConcurrentHashMap<>();
    private final TaskQueue queue;
    private final LongSupplier clock;
    private final Thread[] workers;

    public Scheduler(int workerCount, LongSupplier clock) {
        this.clock = clock;
        this.queue = new TaskQueue(clock);
        this.workers = new Thread[workerCount];
        for (int i = 0; i < workerCount; i++) {
            workers[i] = new Thread(this::workLoop, "worker-" + i);
            workers[i].start();
        }
    }

    public boolean schedule(Task t) {
        if (tasks.putIfAbsent(t.id(), t) != null) return false;   // an id is used once
        queue.add(t);
        return true;
    }

    public boolean cancel(String id) {
        Task t = tasks.get(id);
        return t != null && t.tryCancel();       // a cancelled task stays in the heap; a worker skips it
    }

    private void workLoop() {
        try {
            for (Task t = queue.take(); t != null; t = queue.take()) {
                if (!t.tryStart()) continue;     // cancelled while it waited
                boolean ok;
                try { ok = t.job().run(clock.getAsLong()); } catch (RuntimeException e) { ok = false; }
                finished(t, ok);
            }
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }
    }

    private void finished(Task t, boolean ok) {
        long next = t.recurrence().nextDue(t.dueTime());          // -1 for a one-off task
        if (next >= 0) { t.reopen(next); queue.add(t); return; }  // recurring: the next run, whatever happened
        if (ok) { t.setState(TaskState.DONE); return; }
        long delay = t.retry().delayAfterFailure(t.countFailure());
        if (delay < 0) { t.setState(TaskState.FAILED); return; }
        t.reopen(clock.getAsLong() + delay);                      // back to WAITING, due later
        queue.add(t);
    }
}

Deep dives

Polling, or a priority queue with a condition variable?

Weak: Each worker loops: look at the first task, run it if it is due, otherwise Thread.sleep(100). Tasks start up to 100 ms late. With 8 workers, 80 wake-ups every second find nothing to do. A shorter sleep makes the delay smaller and the waste bigger. Workers also scan an unsorted list to find the earliest task.

Good: Keep the tasks in a PriorityQueue ordered by due time, so the next task is always at the head. A worker reads the head, works out how long to wait and sleeps exactly that long. But a plain sleep cannot be cut short. If a task due in 1 second arrives while the worker sleeps for 1 hour, that task starts an hour late.

Great: Put the queue behind a lock and a Condition. A worker calls awaitNanos(wait) until the head is due. add signals the condition when the new task becomes the head, so sleepers wake up, read the head again and sleep for the new time. An idle pool uses no CPU, and a task starts when it is due. ScheduledThreadPoolExecutor works this way.

Trade-off: PriorityQueue costs O(log n) to add and O(n) to remove one task in the middle, so cancelled tasks stay in the heap and workers skip them. For millions of timers with coarse precision, a timing wheel is cheaper. For 100,000 tasks the heap is fine.

Retries and recurring tasks without holding a thread

Weak: Retry inside the job: for (int i = 0; i < 3; i++) { try { work(); break; } catch (Exception e) { Thread.sleep(300_000); } }. The worker sleeps for 5 minutes and cannot run anything else. Four failing tasks stop a pool of four. A recurring task as while (true) { run(); sleep(period); } needs one thread for each task.

Good: After each run, put the task back into the same queue with a new due time. A retry is the failed task with a due time of now plus the delay. A recurring task is the finished task with a due time of the old one plus the period. While it waits it holds no thread, only a heap entry.

Great: Hide both rules behind small interfaces: Recurrence.nextDue and RetryPolicy.delayAfterFailure. Then fixed-rate (counted from the due time), fixed-delay (counted from the end of the run) and exponential backoff are new classes, and finished stays the same. Cap the retries, and say what happens at the cap.

Trade-off: Fixed-rate keeps the schedule steady, but after a long stall it starts several late runs one after the other. Fixed-delay never overlaps runs, but it drifts. Choose by what the job needs, and write it in the contract.

Cancel against a worker that starts the task

Weak: Cancel only removes the task from the id map. The queue still holds it, and a worker runs it. The caller was told that the task is cancelled, and it ran anyway.

Good: A cancelled flag that the worker checks before it runs the task. This is better, but check and run are two steps. A cancel can arrive between them, and the task runs once after cancel returned true.

Great: One state field and two synchronized moves out of WAITING: tryStart (to RUNNING) and tryCancel (to CANCELLED). Only one of them can succeed, so the answer of cancel is always true. Keep the cancelled task in the heap and let workers skip it, because removing from the middle of a PriorityQueue is O(n). When more than half of the heap is cancelled, rebuild it.

Trade-off: A task that is already RUNNING is not stopped, and cancel says false. Stopping it needs interruption, which only works if the job checks the interrupt flag. Say this in the contract and offer a Future style result if callers need to wait for a run.

Now add one more feature

• Now add priorities: when tasks are due at the same time, the important one starts first. What changes? Add a `priority` to Task and extend `compareTo` to order by due time, then priority (higher first), then sequence. Nothing else changes.
• Now add dependencies: task B runs only after task A is DONE. What changes? Task gets a list of dependencies. `finished` checks which tasks waited for this one and adds them to the queue. Reject cycles when a task is scheduled.
• Now stop a job that runs too long. What changes? Run each job through a `Future` with a timeout and interrupt it when the time is up. Count a timeout as a failure, so the retry policy applies. Jobs must respond to interruption.
• Now run the scheduler on three servers, so each task runs once. What changes? Tasks live in a database. A worker claims the next due task with one conditional update that sets an owner and a lease time. A crashed owner's lease expires and another server takes over, so jobs must be safe to run twice.

What each level is expected to show

SDE-1: Model Scheduler, Task, Job and a TaskState enum. Keep tasks in a PriorityQueue ordered by due time and run them on an ExecutorService with N threads. Handle one-off and recurring tasks, and explain why polling every few milliseconds is wasteful.

SDE-2: Everything at SDE-1, plus workers that wait on a Condition with awaitNanos and are woken when a new earliest task arrives. Recurrence and RetryPolicy are strategies, and a retried or recurring task goes back into the queue without holding a thread. Start and cancel are guarded state moves. The clock is injected, so tests run without sleeping.

Common mistakes

• Polling with a sleep in a loop. It makes tasks late, burns CPU while idle, or both.
• Using `Thread.sleep` until the head is due. A task that arrives with an earlier due time is missed, because the sleeper cannot be woken.
• Sleeping inside a job to retry it. The worker is stuck, and a few failing tasks stop the whole pool.
• One thread per recurring task. It does not scale, and the threads sit idle almost all the time.
• Checking a cancelled flag and then running the task in two separate steps, so a cancelled task can still run once.
• Calling `System.currentTimeMillis()` in the middle of the logic. Tests then need real sleeps and become slow and flaky.
• Ignoring equal due times. Without a sequence number, the order of tasks that are due together is not defined.

Lessons behind this answer

• Double booking and producer-consumer
• Thread safety basics
• Strategy
• Command
• The java.util.concurrent toolbox