workstealpool

package module
v0.1.1 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Aug 28, 2026 License: MIT Imports: 5 Imported by: 0

README

worksteal

A lock-free work-stealing scheduler for recursive divide-and-conquer workloads in Go, implementing the Chase-Lev deque (Lê, Pop, Cohen & Nardelli, PPoPP 2013).

  • deque.go : LFdeque[T], a growable, array-based work-stealing deque. One owner goroutine pushes/pops from the bottom; any number of thieves steal from the top.
  • work_pool.go : WorkerPool[T, R], a pool of workers, each owning one LFdeque[T]. Workers run local work LIFO and steal from other workers' deques FIFO (in half-batches) when they run dry.
  • primecount.go : a divide-and-conquer prime-counting workload used to exercise the pool end-to-end, and as the correctness/benchmark harness for the rest of the package.

Install

go get github.com/PAKIWASI/work_steal_pool

Usage

pool := worksteal.NewWorkerPool[T, R](ctx, poolSize, initialCap, resultBuf, task)
pool.Submit(initialItem)
for r := range pool.Run() {
    // consume results as they arrive
}
if err := pool.Wait(); err != nil {
    // first error from any worker
}

A Task[T, R] either returns a result (leaf) or calls spawn to schedule more work of the same type onto the calling worker's own deque (internal node). See primecount.go for a worked example.

Writing a Task function

type Task[T, R any] func(ctx context.Context, item T, spawn func(T)) (result R, ok bool, err error)

Every call to your Task is one node in the recursion tree, and it returns three distinct signals:

  • Leaf: do the real work for item and return (result, true, nil). The result is emitted to the results channel.
  • Internal node: decide item is still too big, call spawn(child) one or more times to hand off smaller pieces, and return (zero, false, nil). Nothing is emitted.
  • Fatal error: return (zero, false, err). The pool aborts and Wait() reports this error.

Because R is returned by value without indirection (*R), R can be any type (struct, pointer, interface, or scalar) without forcing a heap allocation for leaf returns.

Worked example, from primecount.go (bisect a range until it's small enough to count directly):

func countPrimesTask(threshold int) Task[primeRange, int] {
    return func(ctx context.Context, item primeRange, spawn func(primeRange)) (int, bool, error) {
        width := item.Hi - item.Lo
        if width <= threshold {
            // Leaf: do the real work, return a result.
            count := countPrimesSequential(item.Lo, item.Hi)
            return count, true, nil
        }

        // Internal node: split and hand off both halves.
        mid := item.Lo + width/2
        spawn(primeRange{Lo: item.Lo, Hi: mid})
        spawn(primeRange{Lo: mid, Hi: item.Hi})
        return 0, false, nil
    }
}

Rules to keep in mind:

  • Call spawn synchronously, from inside the Task call itself. It schedules onto the calling worker's own deque — don't stash it and call it later, or call it from another goroutine.
  • Every result you want must come from a return result, true, nil, not from a side channel. The pool collects results only through the return value; combining/summing them (like the range-counting loop in CountPrimesParallel) is the caller's job after draining Run(), not the task tree's.
  • Return a non-nil error to abort the whole pool. The first error from any worker cancels every other worker and is what Wait() returns — don't use it for expected/recoverable conditions, only real failures.
  • Check ctx in long-running leaf work if you want it to be interruptible when the pool cancels (e.g. after another worker errors). Cheap leaves (like countPrimesSequential here) usually don't need to.
  • Pick a leaf-size threshold deliberately — see Performance below. Too fine and CAS/spawn overhead dominates; too coarse and there isn't enough work to steal.

User control over discovered work: recoverable vs. fatal errors

Task's error return is intentionally blunt: any non-nil error is treated as the failure, cancels every other worker, and is what Wait() returns. That's the right behavior for "the pool itself broke," and the wrong behavior for "this one item didn't work out." Those are different situations and the pool only gives you a tool for the first one — the second is on you to build, using the result type.

A directory walker is the clearest example. Walking a tree with spawn(subdir) per directory, you will hit EACCES on some subdirectory somewhere. Returning that from Task as error would cancel the entire walk over one unreadable folder, which is almost never what you want.

Instead, fold expected, recoverable failures into R and return (result, true, nil):

type WalkResult struct {
    Path string
    Info os.FileInfo // nil if Err != nil
    Err  error        // non-nil = recoverable per-item failure
}

func walkTask(ctx context.Context, dir string, spawn func(string)) (WalkResult, bool, error) {
    entries, err := os.ReadDir(dir)
    if err != nil {
        if errors.Is(err, fs.ErrPermission) || errors.Is(err, fs.ErrNotExist) {
            // Expected: report it as a leaf result, don't abort the pool.
            return WalkResult{Path: dir, Err: err}, true, nil
        }
        // Unexpected (corrupted mount, unusual I/O failure, ...): abort.
        return WalkResult{}, false, err
    }

    for _, e := range entries {
        spawn(filepath.Join(dir, e.Name()))
    }
    return WalkResult{Path: dir}, true, nil
}

The consumer sorts results after draining Run(), the same way CountPrimesParallel sums leaf counts:

for r := range pool.Run() {
    if r.Err != nil {
        denied = append(denied, r) // collect, log, retry later — your call
        continue
    }
    files = append(files, r)
}

Guidelines this generalizes to:

  • Use errors.Is/sentinel checks inside the Task to classify an error, not the pool. WorkerPool doesn't know or care what fs.ErrPermission is — keeping that judgment call in your Task keeps the scheduler generic instead of leaking walker-specific semantics into it.
  • Recoverable failures are data, not control flow. They travel out through return result, true, nil, same as any other leaf value, because results are the only per-item channel out of the pool.
  • Reserve return zero, false, err for the true first case: something the running Task can't classify or recover from, where continuing to process other items isn't safe or meaningful (out of file descriptors, a bug, a context you should have checked but didn't). That's the one bucket the pool will actually abort everything for.
  • A zero-value R isn't the same as an error. If R is a plain struct like WalkResult above, callers should check Err first, not infer failure from a nil/zero field — nothing in the pool enforces that convention for you, it's a discipline you're opting into on the result type.

What this is useful for

This pool is built for one specific shape of problem: recursive divide-and-conquer work where you don't know the shape of the tree up front — each item, once you look at it, may produce more items, and the resulting subtrees can be wildly uneven in size. That's exactly the case a fixed, evenly-partitioned worker pool handles badly: partition a lopsided tree evenly across N workers up front and some of them finish early and sit idle while one worker is still chewing through the big branch. Work stealing fixes that by letting idle workers pull work from busy ones at runtime, so load balances itself regardless of how the tree turns out.

Good fits:

  • Parallel recursive sorts — quicksort/mergesort, where partitions split unevenly depending on the data.
  • Tree/graph search — game tree search (e.g. alpha-beta), N-Queens, puzzle solvers, where branches prune to very different depths.
  • Recursive numeric subdivision — this repo's own prime-counting example: bisect a range, recurse until small enough, sum leaf results.
  • Spatial/recursive structures — ray tracing scene traversal, k-d tree or octree construction and queries, fractal rendering.

Poor fits:

  • Uniform, embarrassingly parallel batches (e.g. "resize these 10,000 images") — the work is already evenly sized and known up front, so a plain fixed-size worker pool with a shared channel is simpler and just as fast; there's nothing to steal because there's no imbalance.
  • I/O-bound work — stealing balances CPU work across cores; a blocked-on-network task doesn't benefit from a lock-free deque, use a bounded goroutine pool or semaphore instead.
  • Work that must be processed in order — Run()'s results arrive as leaves finish, not in submission order.

How it works

Each worker pops its own deque LIFO (cache-friendly, and it means a worker keeps working on the subtree it just spawned rather than jumping around). When a worker's deque empties, it steals half of a random victim's deque from the top (FIFO), amortizing the cost of a steal over many items instead of stealing one item at a time. Workers spin briefly on a failed steal, then park until woken by a PushBottom/spawn elsewhere in the pool.

Performance

Full suite (go test ./...) passes, including the concurrent owner/thief/parking tests. Benchmarks below are from go test -bench=. -benchmem, run on:

CPU:    11th Gen Intel(R) Core(TM) i5-1135G7 @ 2.40GHz
Cores:  4 physical cores, 8 threads (GOMAXPROCS=8)
GOOS/GOARCH: linux/amd64
Deque primitives (LFdeque, uncontended and contended)
PushBottom              46.14 ns/op    0 B/op   0 allocs/op
PopBottom               31.31 ns/op    0 B/op   0 allocs/op
Steal                   25.61 ns/op    0 B/op   0 allocs/op
StealHalf (from 1024)  6224    ns/op  682 B/op   0 allocs/op   (~1000 elements/steal)

ConcurrentSteal, 1024 pre-filled, N thieves racing for them concurrently:
  thieves=1    19.63 ns/op
  thieves=2    19.65 ns/op
  thieves=4    19.66 ns/op
  thieves=8    19.63 ns/op
  thieves=16   19.74 ns/op

The concurrent-steal number is the one worth noting: per-steal cost is flat from 1 to 16 concurrent thieves (19.6–19.7 ns/op throughout). That's the lock-free deque doing its job — thieves aren't queueing up behind each other or degrading under contention, each Steal call costs the same regardless of how many other goroutines are hammering the same deque at once.

Full workload: threshold (leaf granularity) is the dominant knob

Counting primes below 200,000, pool size 8, sweeping leaf-size threshold:

threshold=1        94,972,866 ns/op  14,508,950 B/op  602,791 allocs/op
threshold=10       20,260,268 ns/op   2,392,607 B/op   99,196 allocs/op
threshold=50        4,795,156 ns/op     309,412 B/op   12,630 allocs/op
threshold=200       3,287,177 ns/op      83,704 B/op    3,260 allocs/op
threshold=1,000     3,043,703 ns/op      26,275 B/op      875 allocs/op   ← fastest
threshold=5,000     3,134,107 ns/op      11,624 B/op      266 allocs/op
threshold=200,000  12,998,140 ns/op       6,656 B/op       57 allocs/op   (1 leaf, no stealing)
sequential baseline 12,199,920 ns/op          0 B/op        0 allocs/op

Threshold=1 is ~31x slower than the threshold=1,000 sweet spot here, almost entirely from spawn/steal/CAS bookkeeping (600k+ allocations vs. 875) rather than real work — bisecting to single numbers generates far more tree nodes than the leaf work justifies. At the other extreme, threshold=200,000 forces a single leaf with zero parallelism, landing right back near the sequential baseline. The best threshold sits where leaves are cheap enough to steal in useful chunks but not so fine that overhead swamps the work — sweep it for your own workload, this number won't transfer directly.

Pool size: speedup tracks physical core count, then plateaus

Same workload, threshold fixed at 500:

workers=1   12,503,690 ns/op   (1.0x — pool overhead roughly cancels out vs. sequential)
workers=2    6,778,972 ns/op   (1.8x)
workers=4    4,142,564 ns/op   (2.9x)
workers=8    3,100,820 ns/op   (3.9x)
workers=16   3,051,032 ns/op   (4.0x — no further gain)
workers=32   3,145,615 ns/op   (3.9x — slightly worse: oversubscription)

Speedup climbs cleanly through the physical core count (4) and continues a bit further into hyperthreading territory, then flattens right around 8 workers — this CPU's thread count — and going well past that (32 workers) costs a little rather than gaining anything, from scheduling more goroutines than there's parallelism to run them. Matching pool size to runtime.NumCPU() (or close to it) is the right default; there's rarely a reason to go far beyond your thread count.

Threshold × pool size interact

PoolSizeXThreshold makes the interaction explicit: at a too-fine threshold (20), more workers barely help (17.4ms → 11.5ms → 12.6ms going 1 → 4 → 16 workers, actually regressing at 16) because the bottleneck is per-node overhead, not available parallelism. At a reasonable threshold (500), the same pool-size increase scales cleanly (12.6ms → 4.1ms → 3.1ms). Tuning one without the other leaves performance on the table either way.

Pool size, deque capacity, result buffer — mostly a memory knob past a point

Deque capacity and result-buffer size had a much smaller effect than threshold or pool size on this workload. A few hundred µs across the whole sweep — but memory scales directly with whatever you ask for (e.g. InitialWorkerCap=512 uses ~107 KB/op vs. ~43 KB/op at the default), so oversizing them costs memory for little to no speed benefit once they're past the point of avoiding resize churn.

BenchmarkCountPrimes_PoolSize, _InitialWorkerCap, _ResultBuffSize, _Threshold, and _PoolSizeXThreshold in primecount_bench_test.go sweep all of the above; rerun them on your target hardware and workload before picking production values — these numbers are a starting point, not a guarantee.

Testing

go test ./...            # normal run
go test -race ./...      # may intermittently report the known race below;
                          # any other race, or a wrong count/value, is a real bug
go test -bench=. ./...   # benchmarks sweep pool size, deque cap, result
                          # buffer size, leaf threshold, and range size

Known limitation: benign race under -race

Running the tests with -race will occasionally report a race between LFdeque.PushBottom's array write and LFdeque.Steal/StealHalf's array read (deque.go). This is expected and does not affect correctness.

Steal reads the array slot before CASing top to claim it — matching the Chase-Lev paper. If the CAS fails (a thief lost the race), the value it just read is discarded via ok == false, but the read itself already happened, unsynchronized, against whatever the owner does next. So a losing thief's read is a genuine data race on paper — the value can even be torn for multi-word T, but since it's always thrown away, no caller ever observes it. -race is correctly flagging an unsynchronized access, not producing a false positive; it's just one that provably can't corrupt a result.

Two real fixes exist if you're adapting this for production, neither implemented here on purpose: boxing each element as atomic.Pointer[T] (real atomic access, costs one allocation per push), or epoch/hazard-pointer reclamation (will attempt this)

TestCountPrimesParallel_MatchesSequential and TestCountPrimesParallel_Repeated are the tests most likely to trigger it, by design. They drive real concurrent push/steal traffic through a struct-typed T.

Documentation

Overview

Package workstealpool implements concurrent work stealing for worker pools.

Work-stealing pools exist for a specific shape of problem: recursive divide-and-conquer. The classic example is parallel quicksort or a parallel tree walk. You don't know the full list of work upfront. Each piece of work, when you look at it, discovers more work.

A worker pool consists of multiple worker goroutines. Each worker owns a lock-free deque. Workers execute their own work from the bottom of the deque and steal work from the top of other workers' deques when they run out of local work.

Index

Constants

This section is empty.

Variables

This section is empty.

Functions

func CountPrimesParallel 

func CountPrimesParallel(ctx context.Context, lo, hi int, cfg poolConfig) (int, error)

CountPrimesParallel counts primes in [lo, hi) using a WorkerPool configured by cfg. It's the divide-and-conquer entry point: Submit seeds worker 0's deque with the whole range, Run starts every worker pulling and spawning, and the loop over results here does the "conquer" step - summing independent leaf counts as they arrive.

Types

type LFdeque 

type LFdeque[T any] struct {
	// contains filtered or unexported fields
}

LFdeque is a lock-free double-ended queue.

  • top and bottom are ever-increasing int64 counters indices into the backing array are always `counter % cap`.
  • The invarient is `top <= bottom `and the size is just `bottom - top`
  • Because they only ever increase, there is no ABA problem on the CAS below: a value top once held can never recur later.
  • The owner works the bottom end (LIFO: PushBottom/PopBottom).
  • Thieves work the top end (FIFO: Steal), racing each other, and resolved by CAS
  • Thieves only race for the owner's PopBottom for the very last element,
  • This race is resolved by a CAS by both the thief and the owner

TODO: experminet with cache line padding

func NewLFdeque 

func NewLFdeque[T any](capacity int) *LFdeque[T]

func (*LFdeque[T]) Len 

func (d *LFdeque[T]) Len() int64

func (*LFdeque[T]) PopBottom 

func (d *LFdeque[T]) PopBottom() (v T, ok bool)

PopBottom removes and returns the value at the bottom (owner-only). ok is false if the deque was empty, or if a concurrent thief won the race for the last remaining element.

func (*LFdeque[T]) PushBottom 

func (d *LFdeque[T]) PushBottom(v T)

PushBottom adds v to the bottom (owner-only).

func (*LFdeque[T]) PushSliceBottom 

func (d *LFdeque[T]) PushSliceBottom(v []T)

PushSliceBottom pushes all the elements of the slice `v` into the owner's queue at bottom (LIFO). A thief calls this to store the values it stole

func (*LFdeque[T]) Steal 

func (d *LFdeque[T]) Steal() (v T, ok bool)

Steal removes and returns the value at the top (thief-safe: any number of goroutines may call this concurrently, including concurrently with the owner's PushBottom/PopBottom).

KNOWN LIMITATION: the a.get(t) read below can race with a concurrent PushBottom's array write under `go test -race`. This is expected and does not affect correctness - see README.md, "Known limitation: benign data race under -race", for the full explanation.

func (*LFdeque[T]) StealHalf 

func (d *LFdeque[T]) StealHalf() (v []T, ok bool)

StealHalf removes approximately half of the victim's current work from the top and returns it as a batch.

The operation is thief-safe: any number of thieves may call StealHalf concurrently, and it may also race with the owner's PushBottom/PopBottom.

type Task 

type Task[T, R any] func(ctx context.Context, item T, spawn func(T)) (result R, ok bool, err error)

Task is the unit of work a WorkerPool executes.

ctx should be checked by long-running tasks that want to be interruptible. spawn schedules a child item of work onto the calling worker's local deque; it must only be called synchronously, from within this Task invocation.

The three return values encode three distinct outcomes:

  • err != nil: fatal — the pool cancels and records this as its terminal error.
  • ok == true: leaf — result is emitted on the results channel.
  • ok == false: internal node — the task only spawned children; nothing is emitted.

R is unconstrained (any), so it can be a value type, pointer, interface, or struct. Returning by value avoids any forced heap allocation. T is the input type and R is the result type.

type Worker 

type Worker[T any] struct {
	// contains filtered or unexported fields
}

Worker owns a local work-stealing deque.

The worker's normal path is to pop work from the bottom of its deque and push newly spawned work onto the bottom. This keeps the common path local to the worker.

Worker does not know about other workers. The WorkerPool coordinates stealing between workers.

type WorkerPool 

type WorkerPool[T, R any] struct {
	// contains filtered or unexported fields
}

WorkerPool manages a collection of workers and schedules work between them.

The pool does not care what T represents. It only moves T between worker deques. execute defines how a worker executes a T.

R is the result type expected from each execute call of the worker

func NewWorkerPool 

func NewWorkerPool[T, R any](
	ctx context.Context,
	poolSize, initialWorkerCap, resultBuffSize int,
	execute Task[T, R],
) *WorkerPool[T, R]

NewWorkerPool creates a pool of poolSize workers, each with its own deque of initial capacity initialWorkerCap.

execute defines the work each worker performs for a given item. See Task for the contract around ctx, spawn, and error handling.

The pool does not start running until Submit is called and workers begin pulling from their deques; there is no separate "Start" step, workers run as soon as they're constructed, watching ctx and their deques.

func (*WorkerPool[T, R]) Run 

func (p *WorkerPool[T, R]) Run() <-chan R

Run: Result channel generator. Starts all workers and returns the results channel. The channel closes once every worker has exited, either because there's no work left anywhere or because a task returned an error.

Call Wait afterward (or concurrently, while draining results in another goroutine) to get the terminal error, if any.

func (*WorkerPool[T, R]) StealHalf 

func (p *WorkerPool[T, R]) StealHalf(thiefIdx int) (ok bool)

StealHalf attempts to steal work for the given worker from a randomly chosen victim among the other workers in the pool. It tries up to len(workers)-1 distinct victims before giving up.

func (*WorkerPool[T, R]) Submit 

func (p *WorkerPool[T, R]) Submit(item T)

Submit adds initial work to the pool. Call before Run Submit does not itself start any workers.

func (*WorkerPool[T, R]) Wait 

func (p *WorkerPool[T, R]) Wait() error

Wait blocks until every worker has exited and returns the first error encountered (nil on normal completion). Safe to call while another goroutine drains the results channel returned by Run, since results only closes once workers have exited too.

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.