Home/Low-level design/Task Scheduler

Task Scheduler

Priority dispatch, worker quotas, and expiring leases

The problem

A task scheduler chooses pending work and temporarily assigns it to a worker. It owns task state, worker capacity, ordering, and the rules for accepting completion. A priority is only one part of that contract. The scheduler must also define ties, distinguish running work from completed work, and decide what happens when a worker stops reporting progress.

A lease gives an assignment a bounded lifetime. In this exercise, a lease expiring makes the task pending again; it does not prove the work finished. Workers may report completion only while their current lease is active. All time is caller-supplied logical time, so tests can exercise exact boundaries without sleeping.

Naive approach

Keep an unsorted list of tasks. To dispatch, scan for a high-priority pending task and mark it assigned. Count running tasks for the worker to check capacity. This is adequate for a small scheduler if the lifecycle is explicit and every operation uses the same rules.

Worker identity alone accepts a completion from the wrong assignment:

Python
old_attempt = {"worker": "w1", "token": 1}
current = {"worker": "w1", "token": 2, "expires": 20}
accepted = old_attempt["worker"] == current["worker"]
print(accepted)
assert accepted
assert old_attempt["token"] != current["token"]

Reassigning the same task to the same worker makes the worker comparison pass. Assignment identity must change with every lease.

Problems arise when assignment is represented by a single boolean. A task cannot then distinguish pending, running, and done. An expired assignment may remain permanently unavailable, or completed work may accidentally reenter the queue. Sorting displayed identifiers also creates subtle ties: the string "task_10" sorts before "task_2". Store numeric creation order separately from presentation.

Where the naive approach breaks

A worker with capacity one leases task 1 at time 5 until time 10. At time 9, another dispatch must fail because that worker is full. At time 10, the lease has expired, task 1 is pending again, and its capacity slot is free. The interval is [5, 10), so a completion arriving exactly at 10 fails.

Reassigning the same task, even to the same worker, requires a fresh lease token. Otherwise a delayed completion could finish the replacement attempt. Completion must match the current assignment, not just the task or worker identity.

A better design

Separate a task map, worker capacities, and assignment metadata. Tasks have pending, running, or done status. Before any valid timed operation, expire leases whose boundary is at or before that operation's timestamp. Clear their assignment metadata and return those tasks to pending. Completion changes a matching running task to done and releases its slot immediately.

Make completion eligibility a predicate over the current assignment:

Python
def may_complete(current, token, now):
    return (current["status"] == "running"
            and current["token"] == token
            and now < current["expires"])

current = {"status": "running", "token": 2, "expires": 20}
print(may_complete(current, 1, 19), may_complete(current, 2, 20))
assert not may_complete(current, 1, 19)
assert not may_complete(current, 2, 20)
assert may_complete(current, 2, 19)

A stale token and an exactly expired lease both fail. The scheduler must evaluate and apply the transition together after validating logical time. External work can still continue after expiry; this predicate cannot fence another storage system.

Choose pending work by descending priority and ascending numeric task identifier. Dispatch changes exactly one task and returns a detached snapshot with its lease token and expiration. A failed dispatch must not consume an identifier or token or assign any task. Time advancement may still expire old leases, a defined lifecycle transition independent of whether the requested dispatch succeeds. Invalid time or duration arguments are rejected before that transition.

A direct implementation scans tasks for expiration, worker load, and the best candidate, using O(T) time per timed operation for T tasks. Sorting pending candidates instead costs O(P log P). These costs are acceptable here and easier to verify than multiple mutable indexes. At larger scale, a ready heap, expiration heap, and per-worker active counts can reduce scans, but stale heap entries and count maintenance introduce additional invariants. The public contract should remain the same whichever representation is used.

Real-world usage

Leases are useful for coordinating background work, rebuilding indexes, and distributing batch jobs. A capacity limit bounds simultaneous assignments; it is not a limit on requests per minute. This in-memory scheduler neither executes tasks nor prevents a worker from continuing after lease expiry. Real workers may need idempotent effects or a destination that rejects writes carrying old lease tokens. A token checked only by this object cannot undo an external side effect.

Put it into practice

Implement the exercise in a full-screen workspace, then run its tests.

Open practice
Next lessonRecipe Manager