//! A fixed set of worker threads for the pure half of a store write: delta
//! search and compression. The writer thread hands out tasks and helps run
//! them while it waits, so a pool of zero threads runs every task on the
//! writer itself, through the same code. Which thread ran a task never
//! reaches the store: a task computes a value from inputs fixed before it was
//! submitted, and the writer consumes the values in its own order.
//!
//! The import's verifier reuses the same threads once the stores are closed,
//! for the reconstruct-and-rehash sweeps over the frozen files; the rule is
//! the same there: tasks compute, the submitter decides in its own order.
const std = @import("std");
/// The most threads an operator may ask for.
pub const max_threads: usize = 16;
/// The most the default takes. Measured on a 30k-commit kernel slice, 8
/// threads ingest as fast as 14: past that the importing thread's own share
/// (parsing, tree edits, the writes) bounds the run, and more threads only
/// add memory and contention.
pub const default_max_threads: usize = 8;
pub const Task = struct {
run_fn: *const fn (ctx: *anyopaque, index: usize) void,
ctx: *anyopaque,
index: usize,
next: ?*Task = null,
fn run(self: *Task) void {
const outer = in_task;
in_task = true;
defer in_task = outer;
self.run_fn(self.ctx, self.index);
}
};
/// Whether this thread is running a task right now, whichever thread it is.
///
/// The process allocator keeps one free list per thread and never unmaps a
/// slab, so a buffer allocated on one side of the pool and freed on the other
/// lands on a free list its allocator never draws from. Done once per
/// artifact, that strands memory in proportion to the store, not to anything
/// in flight. Pooled code therefore frees a buffer on the side that allocated
/// it -- in a task, or outside one -- and `SideCheck` holds it to that.
pub threadlocal var in_task: bool = false;
/// A test allocator that records which side of the pool allocated each
/// buffer and counts the frees made on the other side.
///
/// A proxy for the stranding it guards against, not a measure of it: the
/// submitter also runs tasks while it waits (`helpUntil`), so "in a task" is
/// not "on a worker thread", and the process allocator moves a contended
/// thread to another thread's free list. What it does catch exactly is a
/// buffer that crosses on every artifact, which is the shape that grows.
pub const SideCheck = struct {
child: std.mem.Allocator,
lock: std.atomic.Mutex = .unlocked,
sides: std.AutoHashMapUnmanaged(usize, bool) = .empty,
crossed: std.atomic.Value(usize) = .init(0),
pub fn allocator(self: *SideCheck) std.mem.Allocator {
return .{ .ptr = self, .vtable = &.{ .alloc = alloc, .resize = resize, .remap = remap, .free = free } };
}
pub fn deinit(self: *SideCheck) void {
self.sides.deinit(self.child);
self.* = undefined;
}
fn acquire(self: *SideCheck) void {
while (!self.lock.tryLock()) std.atomic.spinLoopHint();
}
fn record(self: *SideCheck, ptr: [*]u8) bool {
self.acquire();
defer self.lock.unlock();
self.sides.put(self.child, @intFromPtr(ptr), in_task) catch return false;
return true;
}
fn forget(self: *SideCheck, ptr: [*]u8) void {
self.acquire();
defer self.lock.unlock();
const side = self.sides.fetchRemove(@intFromPtr(ptr)) orelse return;
if (side.value != in_task) _ = self.crossed.fetchAdd(1, .monotonic);
}
fn alloc(ctx: *anyopaque, len: usize, alignment: std.mem.Alignment, ra: usize) ?[*]u8 {
const self: *SideCheck = @ptrCast(@alignCast(ctx));
const ptr = self.child.rawAlloc(len, alignment, ra) orelse return null;
if (!self.record(ptr)) {
self.child.rawFree(ptr[0..len], alignment, ra);
return null;
}
return ptr;
}
fn resize(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ra: usize) bool {
const self: *SideCheck = @ptrCast(@alignCast(ctx));
return self.child.rawResize(memory, alignment, new_len, ra);
}
fn remap(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ra: usize) ?[*]u8 {
const self: *SideCheck = @ptrCast(@alignCast(ctx));
// Refused: a move would need the record rekeyed under the lock
// around the child's call; the caller falls back to alloc and free.
_ = self;
_ = memory;
_ = alignment;
_ = new_len;
_ = ra;
return null;
}
fn free(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, ra: usize) void {
const self: *SideCheck = @ptrCast(@alignCast(ctx));
self.forget(memory.ptr);
self.child.rawFree(memory, alignment, ra);
}
};
pub const Pool = struct {
io: std.Io,
mutex: std.Io.Mutex = .init,
/// Signalled once per submitted task; workers wait on it.
work_ready: std.Io.Condition = .init,
/// Broadcast after every finished task, whichever thread ran it, so any
/// thread waiting in `helpUntil` rechecks its flag.
progress: std.Io.Condition = .init,
head: ?*Task = null,
tail: ?*Task = null,
shutdown: bool = false,
threads: [max_threads - 1]std.Thread = undefined,
spawned: usize = 0,
/// Start `threads - 1` workers; the caller's thread is the last one. A
/// spawn that fails leaves a smaller pool, never a broken one: every task
/// still runs, on the threads that did start or on the writer.
pub fn start(self: *Pool, threads: usize) void {
std.debug.assert(threads >= 1 and threads <= max_threads);
std.debug.assert(self.spawned == 0);
while (self.spawned + 1 < threads) : (self.spawned += 1) {
self.threads[self.spawned] = std.Thread.spawn(.{}, workerLoop, .{self}) catch break;
}
}
/// The number of threads tasks run on, the caller's included.
pub fn threadCount(self: *const Pool) usize {
return self.spawned + 1;
}
pub fn deinit(self: *Pool) void {
self.mutex.lockUncancelable(self.io);
// Every submitter waits for its own tasks before it returns, so
// nothing can still be queued here.
std.debug.assert(self.head == null);
self.shutdown = true;
self.work_ready.broadcast(self.io);
self.mutex.unlock(self.io);
for (self.threads[0..self.spawned]) |t| t.join();
self.* = undefined;
}
pub fn submit(self: *Pool, task: *Task) void {
task.next = null;
self.mutex.lockUncancelable(self.io);
defer self.mutex.unlock(self.io);
if (self.tail) |tail| tail.next = task else self.head = task;
self.tail = task;
self.work_ready.signal(self.io);
}
/// Run queued tasks on this thread until `done` reads true. A task's
/// final step sets its flag before the pool's mutex is taken to announce
/// progress, so a flag read false under the mutex is always followed by
/// a wakeup. Today one thread, the store's writer, submits and waits;
/// the broadcast after a task run here keeps a second waiter correct too.
pub fn helpUntil(self: *Pool, done: *const std.atomic.Value(bool)) void {
self.mutex.lockUncancelable(self.io);
defer self.mutex.unlock(self.io);
while (!done.load(.acquire)) {
if (self.pop()) |task| {
self.mutex.unlock(self.io);
task.run();
self.mutex.lockUncancelable(self.io);
self.progress.broadcast(self.io);
continue;
}
self.progress.waitUncancelable(self.io, &self.mutex);
}
}
/// Call `run_fn(ctx, i)` once for every `i` below `count`, spread over
/// the pool with this thread helping, and return once every call has
/// finished. The calls must be independent of each other: which thread
/// runs which index, and in what order, is not fixed.
pub fn runEach(
self: *Pool,
gpa: std.mem.Allocator,
count: usize,
ctx: *anyopaque,
run_fn: *const fn (ctx: *anyopaque, index: usize) void,
) error{OutOfMemory}!void {
if (count == 0) return;
const tasks = try gpa.alloc(Task, count);
defer gpa.free(tasks);
var batch: Batch = .{ .ctx = ctx, .run_fn = run_fn, .left = .init(@intCast(count)) };
for (tasks, 0..) |*task, i| {
task.* = .{ .run_fn = Batch.run, .ctx = &batch, .index = i };
self.submit(task);
}
self.helpUntil(&batch.done);
}
const Batch = struct {
ctx: *anyopaque,
run_fn: *const fn (ctx: *anyopaque, index: usize) void,
left: std.atomic.Value(u32),
done: std.atomic.Value(bool) = .init(false),
// The last call to finish sets `done`, and nothing here touches the
// batch after that store: the submitter may return and free it.
fn run(ctx: *anyopaque, index: usize) void {
const self: *Batch = @ptrCast(@alignCast(ctx));
self.run_fn(self.ctx, index);
if (self.left.fetchSub(1, .acq_rel) == 1) self.done.store(true, .release);
}
};
fn pop(self: *Pool) ?*Task {
const task = self.head orelse return null;
self.head = task.next;
if (self.head == null) self.tail = null;
return task;
}
fn workerLoop(self: *Pool) void {
self.mutex.lockUncancelable(self.io);
defer self.mutex.unlock(self.io);
while (true) {
if (self.pop()) |task| {
self.mutex.unlock(self.io);
task.run();
self.mutex.lockUncancelable(self.io);
self.progress.broadcast(self.io);
continue;
}
if (self.shutdown) return;
self.work_ready.waitUncancelable(self.io, &self.mutex);
}
}
};
/// The default thread count: one per core, bounded, and 1 when the count
/// cannot be read.
pub fn defaultThreads() usize {
const cpus = std.Thread.getCpuCount() catch return 1;
return @max(1, @min(cpus, default_max_threads));
}
test "SideCheck counts a buffer allocated outside a task and freed inside one" {
var threaded = std.Io.Threaded.init(std.testing.allocator, .{});
defer threaded.deinit();
var pool: Pool = .{ .io = threaded.io() };
pool.start(2);
defer pool.deinit();
var side: SideCheck = .{ .child = std.testing.allocator };
defer side.deinit();
const gpa = side.allocator();
const Frees = struct {
gpa: std.mem.Allocator,
outside: []u8,
fn run(ctx: *anyopaque, index: usize) void {
const self: *@This() = @ptrCast(@alignCast(ctx));
_ = index;
self.gpa.free(self.outside);
// Allocated and freed in the same task: not a crossing.
const own = self.gpa.alloc(u8, 16) catch return;
self.gpa.free(own);
}
};
var frees: Frees = .{ .gpa = gpa, .outside = try gpa.alloc(u8, 16) };
try pool.runEach(std.testing.allocator, 1, &frees, Frees.run);
try std.testing.expectEqual(@as(usize, 1), side.crossed.load(.monotonic));
}
test "a pool of one thread runs every task on the caller" {
var threaded = std.Io.Threaded.init(std.testing.allocator, .{});
defer threaded.deinit();
var pool: Pool = .{ .io = threaded.io() };
pool.start(1);
defer pool.deinit();
try std.testing.expectEqual(@as(usize, 1), pool.threadCount());
const Counter = struct {
hits: [8]u32 = @splat(0),
left: std.atomic.Value(u32) = .init(8),
done: std.atomic.Value(bool) = .init(false),
fn run(ctx: *anyopaque, index: usize) void {
const self: *@This() = @ptrCast(@alignCast(ctx));
self.hits[index] += 1;
if (self.left.fetchSub(1, .acq_rel) == 1) self.done.store(true, .release);
}
};
var counter: Counter = .{};
var tasks: [8]Task = undefined;
for (&tasks, 0..) |*task, i| {
task.* = .{ .run_fn = Counter.run, .ctx = &counter, .index = i };
pool.submit(task);
}
pool.helpUntil(&counter.done);
for (counter.hits) |h| try std.testing.expectEqual(@as(u32, 1), h);
}
test "a pool of several threads runs every task exactly once" {
var threaded = std.Io.Threaded.init(std.testing.allocator, .{});
defer threaded.deinit();
var pool: Pool = .{ .io = threaded.io() };
pool.start(4);
defer pool.deinit();
const Counter = struct {
hits: [256]std.atomic.Value(u32) = @splat(.init(0)),
left: std.atomic.Value(u32) = .init(256),
done: std.atomic.Value(bool) = .init(false),
fn run(ctx: *anyopaque, index: usize) void {
const self: *@This() = @ptrCast(@alignCast(ctx));
_ = self.hits[index].fetchAdd(1, .monotonic);
if (self.left.fetchSub(1, .acq_rel) == 1) self.done.store(true, .release);
}
};
var counter: Counter = .{};
var tasks: [256]Task = undefined;
for (&tasks, 0..) |*task, i| {
task.* = .{ .run_fn = Counter.run, .ctx = &counter, .index = i };
pool.submit(task);
}
pool.helpUntil(&counter.done);
for (&counter.hits) |*h| try std.testing.expectEqual(@as(u32, 1), h.load(.monotonic));
}
test "runEach calls every index once and returns only after the last" {
var threaded = std.Io.Threaded.init(std.testing.allocator, .{});
defer threaded.deinit();
for ([_]usize{ 1, 4 }) |threads| {
var pool: Pool = .{ .io = threaded.io() };
pool.start(threads);
defer pool.deinit();
const Counter = struct {
hits: [64]std.atomic.Value(u32) = @splat(.init(0)),
fn run(ctx: *anyopaque, index: usize) void {
const self: *@This() = @ptrCast(@alignCast(ctx));
_ = self.hits[index].fetchAdd(1, .monotonic);
}
};
var counter: Counter = .{};
try pool.runEach(std.testing.allocator, counter.hits.len, &counter, Counter.run);
for (&counter.hits) |*h| try std.testing.expectEqual(@as(u32, 1), h.load(.monotonic));
try pool.runEach(std.testing.allocator, 0, &counter, Counter.run);
}
}