The Functional Programmer
0.00
recursion26tagscollectionsarchivewikifeed

A Work-Stealing Scheduler in OCaml

2026-09-26 · 10 min

ocamlconcurrencydomains

domainslib's task pool is exactly this: every domain owns a deque, pushes and pops its own end, and idle domains steal from the opposite end of someone else's. This builds the deque that makes that safe without a lock on the common path, then a pool of domains driving it, then a fork/join layer on top that turns it into something you actually call parallel_for against.

§ 01

Why a plain queue is not enough

A shared work queue behind one mutex is the obvious first design - correct, and the wrong one. Every push and every pop takes the lock, so the owning domain pays contention on its own overwhelmingly common operation just because some other domain might occasionally want to steal.

(* the design this post replaces *)
type naive_pool = { queue : (unit -> unit) Queue.t; mutex : Mutex.t }
 
let push p task =
  Mutex.lock p.mutex;
  Queue.push task p.queue;
  Mutex.unlock p.mutex
 
let pop p =
  Mutex.lock p.mutex;
  let r = if Queue.is_empty p.queue then None else Some (Queue.pop p.queue) in
  Mutex.unlock p.mutex;
  r

The fix that actually matters: give every domain its own deque. The owner pushes and pops the bottom, uncontended, in the common case; a thief pops the top, and top/bottom operations only collide when the deque is down to its last one or two items - rare, and cheap to handle correctly when it happens.

owner:  push_bottom, pop_bottom   -- LIFO, uncontended, no other domain touches this end
thief:  steal (= pop_top)          -- FIFO from the other end, contended only near-empty
§ 02

The Chase-Lev deque: layout

A growable circular buffer plus two atomic indices. top only ever increases via a CAS from a thief; bottom is only ever written by the owner - that asymmetry is the whole basis for the algorithm needing no lock on the owner's fast path.

(* lib/deque.ml *)
 
type 'a t = {
  mutable buffer : 'a option array;
  top : int Atomic.t;      (* written only by thieves, via CAS *)
  bottom : int Atomic.t;   (* written only by the owner *)
}
 
let create (initial_capacity : int) : 'a t =
  { buffer = Array.make initial_capacity None;
    top = Atomic.make 0;
    bottom = Atomic.make 0 }

Indices wrap via modulo into the circular buffer - size is always a power of two so a mask replaces a division on every access.

let mask (d : 'a t) : int = Array.length d.buffer - 1
 
let get (d : 'a t) (i : int) : 'a = Option.get d.buffer.(i land mask d)
let set (d : 'a t) (i : int) (v : 'a) : unit = d.buffer.(i land mask d) <- Some v
§ 03

push_bottom: owner-only, no CAS needed

The owner is the only writer of bottom, so this needs no atomic read-modify-write at all - a plain load, a store into the buffer, and a plain store back to bottom with a release fence so a thief reading top afterward sees the new item. Growing the buffer when full is the one wrinkle.

let grow (d : 'a t) (b : int) (t : int) : unit =
  let old_cap = Array.length d.buffer in
  let new_buf = Array.make (old_cap * 2) None in
  for i = t to b - 1 do
    new_buf.(i land (old_cap * 2 - 1)) <- Some (get d i)
  done;
  d.buffer <- new_buf
 
let push_bottom (d : 'a t) (item : 'a) : unit =
  let b = Atomic.get d.bottom in
  let t = Atomic.get d.top in
  if b - t >= Array.length d.buffer - 1 then grow d b t;
  set d b item;
  Atomic.set d.bottom (b + 1)

OCaml's Atomic.set already implies a release fence on the runtimes this targets (5.x with the multicore GC) - a plain store on top of a boxed atomic ref would not be safe here on a weaker memory model, which is exactly why bottom and top are Atomic.t rather than plain mutable fields even though only one side ever writes each.

(* the correctness argument depends on this store being visible to other
   domains in the right order relative to prior writes into the buffer -
   Atomic.set is the primitive that guarantees that on OCaml 5's memory model *)
§ 04

pop_bottom: owner-only, but racing thieves near empty

This is the subtle one. The owner tentatively removes the bottom item by decrementing bottom first, then checks whether a thief might have concurrently taken the very same last item from the top - if so, a compare-and-swap on top settles who actually won it.

Speculatively claim the slot by decrementing bottom before checking anything - this ordering, not the more obvious check-then-decrement, is what Chase-Lev actually requires for the algorithm to be linearizable.

let pop_bottom (d : 'a t) : 'a option =
  let b = Atomic.get d.bottom - 1 in
  Atomic.set d.bottom b;
  let t = Atomic.get d.top in
  if t > b then begin
    (* deque was already empty - undo the speculative decrement *)
    Atomic.set d.bottom (b + 1);
    None
  end
  else begin
    let item = get d b in
    if t = b then begin
      (* exactly one item left: race a thief for it via CAS on top *)
      if Atomic.compare_and_set d.top t (t + 1) then begin
        Atomic.set d.bottom (b + 1);
        Some item
      end else begin
        (* a thief won the race first *)
        Atomic.set d.bottom (b + 1);
        None
      end
    end else
      (* more than one item was left: no thief could reach bottom, safe outright *)
      Some item
  end

Tracing the one-item-left race explicitly: owner and a thief both reach for the same last item. Whichever wins the CAS on top gets it; the loser's bottom decrement is harmless since bottom is restored to b+1 either way, leaving the deque correctly empty afterward regardless of who won.

owner:  b := bottom-1; sees t = b (one item left)
thief:  concurrently CAS top from t to t+1, trying to steal the same item
 
case A: owner's CAS succeeds first -> owner gets the item, thief's later CAS fails, steals nothing
case B: thief's CAS succeeds first -> owner's CAS fails, owner gets None, thief has the item
either way: exactly one of them ends up with the item, never both, never neither
§ 05

steal: the thief's side

A thief only ever reads bottom (never writes it) and only ever advances top via CAS - symmetrical to the tricky branch of pop_bottom above, but simpler, since a thief never has to worry about racing itself.

let steal (d : 'a t) : 'a option =
  let t = Atomic.get d.top in
  let b = Atomic.get d.bottom in
  if t >= b then None            (* empty, from this thief's point of view *)
  else begin
    let item = get d t in
    if Atomic.compare_and_set d.top t (t + 1) then Some item
    else None                     (* lost a race - either the owner popped it, or another thief did *)
  end

A losing steal returns None rather than retrying internally - deliberately. The caller (the scheduler's idle loop, next) decides whether to retry the same victim or move on to a different one, which is a scheduling policy question, not a data-structure one.

(* steal returning None means "try again or try someone else", not "empty forever" -
   the deque itself makes no promises about what an idle domain should do next *)
§ 06

A worker per domain

Each worker owns one deque and runs a loop: drain its own bottom first, and only steal when its own deque is empty. n_domains and a shared array of every worker's deque is all workers need to see of each other.

(* lib/scheduler.ml *)
 
type task = unit -> unit
 
type worker = {
  id : int;
  deque : task Deque.t;
}
 
type pool = {
  workers : worker array;
  running : bool Atomic.t;
}

Random victim selection - picking a fixed order (always steal from worker 0 first) would let every idle domain hammer the same one under load; a random pick spreads contention evenly across the pool.

let random_victim (self_id : int) (n_workers : int) : int =
  let v = Random.int (n_workers - 1) in
  if v >= self_id then v + 1 else v

The idle loop: try every other worker's deque, in a random order refreshed on each pass, and only actually park (yield the domain) once a full sweep finds nothing to steal - spinning briefly before yielding trades a little CPU for much lower latency picking up freshly pushed work.

let try_steal (pool : pool) (self_id : int) : task option =
  let n = Array.length pool.workers in
  let rec sweep attempts_left =
    if attempts_left = 0 then None
    else
      let victim = random_victim self_id n in
      match Deque.steal pool.workers.(victim).deque with
      | Some task -> Some task
      | None -> sweep (attempts_left - 1)
  in
  sweep (n - 1)

The per-domain run loop: own work first, then stolen work, then a brief spin-then-yield before checking pool.running - the flag one call sets to stop every domain cleanly at shutdown.

let run_worker (pool : pool) (self : worker) : unit =
  let idle_spins = ref 0 in
  while Atomic.get pool.running do
    match Deque.pop_bottom self.deque with
    | Some task -> idle_spins := 0; task ()
    | None ->
      (match try_steal pool self.id with
       | Some task -> idle_spins := 0; task ()
       | None ->
         incr idle_spins;
         if !idle_spins > 1000 then Domain.cpu_relax ()
         else ())
  done

Spawning the pool: one OCaml Domain per worker, each running the loop above against its own deque. Domain.spawn returns a handle you must eventually join, the same as Thread.create does for threads.

let create (n_domains : int) : pool =
  let workers =
    Array.init n_domains (fun id -> { id; deque = Deque.create 256 })
  in
  let pool = { workers; running = Atomic.make true } in
  let handles =
    Array.map
      (fun w -> Domain.spawn (fun () -> run_worker pool w))
      (Array.sub workers 1 (n_domains - 1))
  in
  (pool, handles)

Submitting work from outside any worker (the caller's own domain, domain 0, which is never spawned above - it runs its own worker loop directly instead of being handed a Domain.spawn handle).

let submit (pool : pool) (task : task) : unit =
  Deque.push_bottom pool.workers.(0).deque task
 
let shutdown (pool : pool) (handles : unit Domain.t array) : unit =
  Atomic.set pool.running false;
  Array.iter Domain.join handles
§ 07

Fork/join on top: futures

A raw task queue is not yet the API anyone wants to program against. Fork/join wraps it: fork submits a task and immediately returns a future; join blocks until that future is filled - except when the calling domain is itself an idle worker, in which case it should help steal and run other work instead of blocking uselessly.

A future as a single-assignment cell, backed by a condition variable for the case where the waiter has nothing else useful to do.

(* lib/future.ml *)
 
type 'a state = Pending | Done of 'a
 
type 'a t = {
  state : 'a state Atomic.t;
  mutex : Mutex.t;
  cond : Condition.t;
}
 
let create () : 'a t =
  { state = Atomic.make Pending; mutex = Mutex.create (); cond = Condition.create () }
 
let fill (f : 'a t) (v : 'a) : unit =
  Mutex.lock f.mutex;
  Atomic.set f.state (Done v);
  Condition.broadcast f.cond;
  Mutex.unlock f.mutex

fork: wrap the task so it fills the future on completion, submit it to the current worker's own deque, and hand back the future immediately - the caller does not wait here.

(* lib/fork_join.ml *)
 
let current_worker : worker option ref = ref None   (* set by run_worker on entry *)
 
let fork (pool : pool) (f : unit -> 'a) : 'a Future.t =
  let fut = Future.create () in
  let task () = Future.fill fut (f ()) in
  (match !current_worker with
   | Some w -> Deque.push_bottom w.deque task
   | None -> submit pool task);
  fut

join: the part that makes this a real scheduler rather than a thread pool with extra steps. A worker blocking on Condition.wait while idle would starve the pool under deep recursion - forking many small tasks and immediately joining each. Instead, while waiting, help: pop or steal other work and run it, and only fall back to a real blocking wait if there is truly nothing left anywhere.

let rec join (pool : pool) (fut : 'a Future.t) : 'a =
  match Atomic.get fut.Future.state with
  | Done v -> v
  | Pending ->
    let did_work =
      match !current_worker with
      | Some w ->
        (match Deque.pop_bottom w.deque with
         | Some task -> task (); true
         | None ->
           (match try_steal pool w.id with
            | Some task -> task (); true
            | None -> false))
      | None -> false
    in
    if did_work then join pool fut
    else begin
      Mutex.lock fut.Future.mutex;
      while Atomic.get fut.Future.state = Pending do
        Condition.wait fut.Future.cond fut.Future.mutex
      done;
      Mutex.unlock fut.Future.mutex;
      join pool fut
    end
§ 08

parallel_for on top of fork/join

A parallel loop as a binary fork/join split, bottoming out at a sequential chunk once the range is small enough that spawning more tasks would cost more than it saves - the standard divide-and-conquer shape every fork/join workload reduces to.

(* lib/parallel.ml *)
 
let sequential_threshold = 1024
 
let rec parallel_for (pool : Scheduler.pool) (lo : int) (hi : int) (f : int -> unit) : unit =
  if hi - lo <= sequential_threshold then
    for i = lo to hi - 1 do f i done
  else begin
    let mid = lo + (hi - lo) / 2 in
    let left = Fork_join.fork pool (fun () -> parallel_for pool lo mid f) in
    parallel_for pool mid hi f;   (* run the right half inline, not forked - halves task count *)
    Fork_join.join pool left
  end

Running only the left half as a forked task and the right half inline is a deliberate, standard trick - it halves the number of futures created versus forking both halves, at no cost to parallelism, since the calling domain was going to do work either way.

(* forking both halves: 2 futures per split, this domain then joins both *)
(* forking only the left: 1 future per split, this domain does the right
   half directly instead of idling on a second join - strictly less overhead *)

Using it: summing an array in parallel, with each leaf task writing into its own slice of a results array rather than contending on one shared accumulator.

let parallel_sum (pool : Scheduler.pool) (xs : int array) : int =
  let n = Array.length xs in
  let partials = Array.make n 0 in
  Parallel.parallel_for pool 0 n (fun i -> partials.(i) <- xs.(i));
  Array.fold_left ( + ) 0 partials
§ 09

What real work-stealing schedulers add on top of this

Not implemented: work-first vs help-first scheduling policy (this post's join is help-first - it looks for other work before blocking, which is the right default for fine-grained fork/join but not the only valid policy); a bounded-size fast path avoiding the CAS entirely when a deque has never been stolen from; NUMA-aware victim selection, preferring to steal from a domain on the same physical socket; and locality hints so a task tends to be picked up by the domain that is already warm on its data, rather than a uniformly random one. None of these change correctness - they are exactly the kind of throughput work the SAT series named as out of scope for its own solver, in the same spirit here.

work-first vs help-first     which side of a fork the forking domain continues on
CAS-free fast path           skip the atomic entirely on an unstolen deque
NUMA-aware stealing          bias toward physically nearby victims
locality hints                keep a task's data warm on the domain that touches it most

related

← older

Lenses and Prisms in OCaml

2026-09-25

newer →

LiquidHaskell

2026-09-27