Something went wrong. Try again.
fork of gitlab.com/devnw/zig/thread-pool for zig 0.17
Something went wrong. Try again.
19 kB · 456 lines
Zig
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457//! Lock-free SPMC ring queue: per-worker local queue for the `work_stealing` and//! `lane_extension` scheduler variants.//!//! Single producer (owner thread) pushes to `tail`. The owner and any stealer take//! from `head`, so the queue is FIFO for everyone: the owner pops one item with a CAS//! on `head`, stealers take batches with the same CAS. The CAS is the linearization//! point; stealers copy slots before it and discard the copy on loss. The pool's//! per-worker `lifo_slot` only holds the FIRST worker-context submit while it is//! occupied; a task spawning N children runs child 1 from the slot and children//! 2..N oldest-first out of this queue, not last-in-first-out over the whole//! fan-out. Buffer size must be a power of 2.
const std = @import("std");
/// Lock-free SPMC ring queue modelled after Tokio's/// `runtime/scheduler/multi_thread/queue.rs` and Go's `runq*` in/// `runtime/proc.go`. Single producer (owner) pushes to `tail`. The owner/// pops one item from `head` (FIFO, like Go's `runqget` and Tokio's/// `pop`); any thread may steal from `head` in batches.////// Correctness sketch:/// - Owner is the sole writer of `tail` and of slot contents./// - Every take, owner or stealer, claims its items with a seq_cst CAS on/// `head`. The CAS is the linearization point: a claim succeeds only if/// `head` has not moved since the claimant loaded it, so no item is/// claimed twice and `head` never passes `tail`./// - Owner never writes to any index `i` where `head <= i < tail`. A/// stealer copies slots between its `head` load and its CAS; if the owner/// reuses one of those slots, `head` must have moved first, so the CAS/// fails and the copy is discarded./// - SONNY-2448: before v0.2.8 the owner popped LIFO from `tail` and/// arbitrated only the last element. A stealer's batch claim could then/// land after the owner had taken (or re-pushed into) the batch's slots:/// duplicates, lost tasks, and `head` past `tail`. Popping from `head`/// removes that window. The replays at the bottom of this file pin it.////// Buffer size must be a power of 2.pub fn RingQueue(comptime T: type) type { return struct { const Self = @This();
allocator: std.mem.Allocator, buffer: []std.atomic.Value(T), mask: u32, head: std.atomic.Value(u32) = std.atomic.Value(u32).init(0), tail: std.atomic.Value(u32) = std.atomic.Value(u32).init(0),
pub fn init(allocator: std.mem.Allocator, size: usize) !Self { std.debug.assert(size > 0 and (size & (size - 1)) == 0); std.debug.assert(size <= std.math.maxInt(u32) / 2); const buf = try allocator.alloc(std.atomic.Value(T), size); for (buf) |*slot| slot.* = std.atomic.Value(T).init(undefined); return .{ .allocator = allocator, .buffer = buf, .mask = @intCast(size - 1), }; }
pub fn deinit(self: *Self) void { self.allocator.free(self.buffer); self.* = undefined; }
pub fn capacity(self: *const Self) usize { return self.buffer.len; }
/// Approximate count. Safe from any thread; may be stale. pub fn len(self: *const Self) usize { const t = self.tail.load(.acquire); const h = self.head.load(.acquire); const diff = t -% h; // `tail` is loaded before `head`. If items were pushed and taken in // between, `head` can be ahead of the loaded `tail`, and `t -% h` // wraps to a value > capacity. Clamp to 0 so callers (stats peak // tracking, empty-victim skip) see "empty or near-empty" rather // than ~u32_max. if (diff > @as(u32, @intCast(self.buffer.len))) return 0; return diff; }
/// Owner-only. Push an item to the tail. Returns false if full. pub fn tryPush(self: *Self, value: T) bool { const t = self.tail.load(.monotonic); const h = self.head.load(.acquire); if (t -% h >= @as(u32, @intCast(self.buffer.len))) return false; self.buffer[t & self.mask].store(value, .unordered); // Release ensures the slot write is visible before tail advances. self.tail.store(t +% 1, .release); return true; }
/// Owner-only. Pop the oldest item from `head` (FIFO). Returns null /// if empty. /// /// Claims the item with the same CAS on `head` that `stealHalf` uses, /// so an in-flight batch steal whose copy covers this item fails its /// claim and retries. Uses `cmpxchgStrong` (not `Weak`) so the only /// retry source is a genuine loss to a stealer, not a spurious /// failure: the loop is bounded because a failed CAS means a /// stealer took at least one item, and only the owner (which is /// here) can add items, so it runs at most `capacity` times. pub fn pop(self: *Self) ?T { while (true) { const h = self.head.load(.acquire); // Owner is tail's only writer: its own last store is current. const t = self.tail.load(.monotonic); if (t == h) return null; const value = self.buffer[h & self.mask].load(.monotonic); if (self.head.cmpxchgStrong(h, h +% 1, .seq_cst, .monotonic) == null) { return value; } } }
/// Any thread. Steal up to `(len+1)/2` items into `dst`. Returns the /// number stolen (0..dst.len). Copies happen *before* the CAS; on /// CAS failure we retry the whole batch. No ABA because `head` is /// monotonically increasing within the u32 range we actually use. pub fn stealHalf(self: *Self, dst: []T) usize { if (dst.len == 0) return 0; // Retry budget. The transient `avail > capacity` read (a stale // `head`, see below) and CAS-contention retries resolve within a // handful of iterations; anything longer is a degenerate // contention pattern better served by moving to the next victim. // Cap rather than spin forever so shutdown cannot stall waiting // for a worker to leave stealHalf. var retries: u32 = 0; while (true) { if (retries >= 256) return 0; retries += 1; const h = self.head.load(.acquire); // Acquire on tail pairs with the owner's release on push. const t = self.tail.load(.acquire); const avail = t -% h; if (avail == 0) return 0; // `head` was loaded before `tail`. If the owner popped and // pushed in between, `avail` can exceed capacity. Treat it as // an inconsistent read and retry rather than trusting it // (Go's `runqgrab` does the same). if (avail > @as(u32, @intCast(self.buffer.len))) continue; // Steal half (round up) — matches Tokio/Go. var n: u32 = (avail + 1) / 2; if (n > @as(u32, @intCast(dst.len))) n = @intCast(dst.len);
var i: u32 = 0; while (i < n) : (i += 1) { dst[i] = self.buffer[(h +% i) & self.mask].load(.monotonic); }
if (self.head.cmpxchgStrong(h, h +% n, .seq_cst, .monotonic) == null) { return n; } // CAS failed: the owner or another stealer moved head. Retry; // the slot reads are discarded. } } };}
test "ring queue: owner push/pop is FIFO" { const Q = RingQueue(u32); var q = try Q.init(std.testing.allocator, 8); defer q.deinit();
try std.testing.expect(q.tryPush(1)); try std.testing.expect(q.tryPush(2)); try std.testing.expect(q.tryPush(3)); try std.testing.expectEqual(@as(?u32, 1), q.pop()); try std.testing.expectEqual(@as(?u32, 2), q.pop()); try std.testing.expectEqual(@as(?u32, 3), q.pop()); try std.testing.expectEqual(@as(?u32, null), q.pop());}
test "ring queue: pop and push wrap the ring" { const Q = RingQueue(u32); var q = try Q.init(std.testing.allocator, 4); defer q.deinit();
var next_push: u32 = 0; var next_pop: u32 = 0; var round: u32 = 0; while (round < 10) : (round += 1) { while (q.tryPush(next_push)) next_push += 1; try std.testing.expectEqual(@as(usize, 4), q.len()); try std.testing.expectEqual(@as(?u32, next_pop), q.pop()); try std.testing.expectEqual(@as(?u32, next_pop + 1), q.pop()); next_pop += 2; } while (q.pop()) |value| { try std.testing.expectEqual(next_pop, value); next_pop += 1; } try std.testing.expectEqual(next_push, next_pop);}
test "ring queue: full rejects push" { const Q = RingQueue(u32); var q = try Q.init(std.testing.allocator, 2); defer q.deinit();
try std.testing.expect(q.tryPush(10)); try std.testing.expect(q.tryPush(20)); try std.testing.expect(!q.tryPush(30)); try std.testing.expectEqual(@as(usize, 2), q.len());}
test "ring queue: stealHalf copies half and advances head" { const Q = RingQueue(u32); var q = try Q.init(std.testing.allocator, 16); defer q.deinit(); var i: u32 = 0; while (i < 10) : (i += 1) _ = q.tryPush(i);
var out: [16]u32 = undefined; const n = q.stealHalf(out[0..]); // (10+1)/2 = 5 try std.testing.expectEqual(@as(usize, 5), n); try std.testing.expectEqual(@as(u32, 0), out[0]); try std.testing.expectEqual(@as(u32, 4), out[4]); try std.testing.expectEqual(@as(usize, 5), q.len());}
test "ring queue: stealHalf empty returns 0" { const Q = RingQueue(u32); var q = try Q.init(std.testing.allocator, 16); defer q.deinit(); var out: [4]u32 = undefined; try std.testing.expectEqual(@as(usize, 0), q.stealHalf(out[0..]));}
test "ring queue: stealHalf caps at dst.len" { const Q = RingQueue(u32); var q = try Q.init(std.testing.allocator, 16); defer q.deinit(); var i: u32 = 0; while (i < 10) : (i += 1) _ = q.tryPush(i);
var out: [2]u32 = undefined; const n = q.stealHalf(out[0..]); try std.testing.expectEqual(@as(usize, 2), n); try std.testing.expectEqual(@as(u32, 0), out[0]); try std.testing.expectEqual(@as(u32, 1), out[1]);}
test "ring queue: wraparound via head advance" { const Q = RingQueue(u32); var q = try Q.init(std.testing.allocator, 4); defer q.deinit(); try std.testing.expect(q.tryPush(1)); try std.testing.expect(q.tryPush(2)); var out: [1]u32 = undefined; _ = q.stealHalf(out[0..]); // consumes 1 (since (2+1)/2 = 1) try std.testing.expect(q.tryPush(3)); try std.testing.expect(q.tryPush(4)); try std.testing.expect(q.tryPush(5)); try std.testing.expect(!q.tryPush(6)); // full (head=1, tail=5, capacity=4)}
// SONNY-2448: deterministic replays of the owner-pop vs batch-steal race.//// `stealHalf` loads `head` and `tail`, copies up to half the queue, and only// then claims the copy with a CAS on `head`. A stealer can be preempted// between the copy and the CAS for any length of time. These replays hold// that window open: `StealReplay.snapshot` performs stealHalf's steps up to// the claim, the owner runs `pop`/`tryPush` inside the window, and// `StealReplay.finish` performs the claim. A claim that fails is a stale// copy that stealHalf discards and retries, so `finish` then calls the real// `stealHalf`. Every value handed out, by the owner or by a claimed batch,// goes through a `Ledger`, which checks that each pushed value was taken// exactly once and that nothing else was.//// In the pool, the three outcomes these replays pin were: a panic on a// never-written slot (`index out of bounds: index 2863311530` in runTask,// i.e. 0xAAAAAAAA), a task that never ran and hung its wait group, and a// task dispatched twice. With the FIFO owner pop, every owner take inside// the window moves `head`, so the stale claim fails and the stealer retries.
/// Test-only record of every value pushed to, and taken from, a replayed/// queue.const Ledger = struct { pushed: [16]bool = @splat(false), taken: [16]u8 = @splat(0), /// Values handed out that were never pushed: a slot that was never /// written, or one whose value was already consumed and overwritten. never_pushed: usize = 0,
fn push(self: *Ledger, q: *RingQueue(u32), value: u32) !void { try std.testing.expect(q.tryPush(value)); self.pushed[value] = true; }
fn take(self: *Ledger, value: u32) void { if (value >= self.pushed.len or !self.pushed[value]) { self.never_pushed += 1; return; } self.taken[value] += 1; }
/// The owner pops until `pop` reports empty. No queue ever holds more /// than `capacity` items, so a pop that returns more has walked past /// the items that exist. fn drain(self: *Ledger, q: *RingQueue(u32)) !void { var pops: usize = 0; while (q.pop()) |value| { pops += 1; if (pops > q.capacity()) return error.PopNeverEmpties; self.take(value); } }
fn expectNothingNeverPushed(self: *const Ledger) !void { if (self.never_pushed != 0) { std.debug.print("{d} value(s) handed out were never pushed\n", .{self.never_pushed}); return error.NeverPushedValue; } }
fn expectNoDuplicate(self: *const Ledger) !void { for (self.taken, 0..) |count, value| { if (count > 1) { std.debug.print("value {d} was taken {d} times\n", .{ value, count }); return error.DuplicateTake; } } }
fn expectNoneLost(self: *const Ledger) !void { for (self.pushed, self.taken, 0..) |pushed, count, value| { if (pushed and count == 0) { std.debug.print("value {d} was pushed but never taken\n", .{value}); return error.LostValue; } } }
fn expectEachTakenOnce(self: *const Ledger) !void { try self.expectNothingNeverPushed(); try self.expectNoDuplicate(); try self.expectNoneLost(); }};
/// Test-only: `stealHalf` split at its claim.const StealReplay = struct { head: u32, n: u32, batch: [8]u32 = undefined,
/// stealHalf's steps before its CAS: load head, then tail; size the /// batch at half the queue, rounded up; copy the batch's slots. fn snapshot(q: *RingQueue(u32)) StealReplay { const h = q.head.load(.acquire); const t = q.tail.load(.acquire); var steal: StealReplay = .{ .head = h, .n = (t -% h + 1) / 2 }; std.debug.assert(steal.n <= steal.batch.len); for (steal.batch[0..steal.n], 0..) |*slot, i| { slot.* = q.buffer[(h +% @as(u32, @intCast(i))) & q.mask].load(.monotonic); } return steal; }
/// stealHalf's claim: the same CAS on `head`. On success the stealer /// owns its copied batch. On failure it discards the copy and retries, /// which here is a call to the real `stealHalf`. fn finish(self: *const StealReplay, q: *RingQueue(u32), ledger: *Ledger) void { if (q.head.cmpxchgStrong(self.head, self.head +% self.n, .seq_cst, .monotonic) == null) { for (self.batch[0..self.n]) |value| ledger.take(value); return; } var retry: [8]u32 = undefined; const n = q.stealHalf(retry[0..]); for (retry[0..n]) |value| ledger.take(value); }};
fn expectHeadNotPastTail(q: *RingQueue(u32)) !void { const h = q.head.load(.monotonic); const t = q.tail.load(.monotonic); if (t -% h > q.capacity()) { std.debug.print("head {d} is past tail {d}\n", .{ h, t }); return error.HeadPastTail; }}
test "ring queue replay (SONNY-2448): a batch claim never moves head past tail" { // Old LIFO pop: the owner took 3 and 2 from the tail, the claim (head // still 0) landed anyway, and head = 2 was past tail = 1. Later pops // returned 1 again, then slots never written, while tail wrapped. var q = try RingQueue(u32).init(std.testing.allocator, 8); defer q.deinit(); var ledger: Ledger = .{}; for ([_]u32{ 1, 2, 3 }) |value| try ledger.push(&q, value);
const steal = StealReplay.snapshot(&q); // head 0, tail 3: batch {1, 2} try std.testing.expectEqual(@as(u32, 2), steal.n); ledger.take(q.pop().?); ledger.take(q.pop().?); steal.finish(&q, &ledger);
try expectHeadNotPastTail(&q); try ledger.drain(&q); try ledger.expectEachTakenOnce(); try std.testing.expectEqual(@as(usize, 0), q.len());
// The queue is still usable. try ledger.push(&q, 4); try ledger.drain(&q); try ledger.expectEachTakenOnce();}
test "ring queue replay (SONNY-2448): a push inside an in-flight batch is never lost" { // Old LIFO pop: the owner took 3 and 2, then pushed 4 into slot 1, // inside the stealer's batch. The claim landed (head 2 == tail 2), so // 4 sat below head and no pop or steal could reach it; in the pool // that task never ran and its wait group never completed. var q = try RingQueue(u32).init(std.testing.allocator, 8); defer q.deinit(); var ledger: Ledger = .{}; for ([_]u32{ 1, 2, 3 }) |value| try ledger.push(&q, value);
const steal = StealReplay.snapshot(&q); // batch {1, 2} ledger.take(q.pop().?); ledger.take(q.pop().?); try ledger.push(&q, 4); steal.finish(&q, &ledger);
try ledger.drain(&q); try ledger.expectNoneLost(); try expectHeadNotPastTail(&q); try ledger.expectEachTakenOnce();}
test "ring queue replay (SONNY-2448): the owner and a stealer never take the same item" { // Old LIFO pop: the owner took 4 and 3, then 2 from inside the // stealer's batch {1, 2}. The claim landed, so 2 was handed out twice; // in the pool, the second dispatch of that entry reached runTask. var q = try RingQueue(u32).init(std.testing.allocator, 8); defer q.deinit(); var ledger: Ledger = .{}; for ([_]u32{ 1, 2, 3, 4 }) |value| try ledger.push(&q, value);
const steal = StealReplay.snapshot(&q); // head 0, tail 4: batch {1, 2} try std.testing.expectEqual(@as(u32, 2), steal.n); ledger.take(q.pop().?); ledger.take(q.pop().?); ledger.take(q.pop().?); steal.finish(&q, &ledger);
try ledger.expectNoDuplicate(); try expectHeadNotPastTail(&q); try ledger.drain(&q); try ledger.expectEachTakenOnce();}