//! 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 {
self.run_fn(self.ctx, self.index);
}
};
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 "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);
}
}