Skip to content
zit Artifact [fccdb1fbd1]
Login

Artifact [fccdb1fbd1]

Artifact fccdb1fbd106efeab2928207218b88779eea9a22177e72d7353bf19f3ba9ae43:


//! 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.
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);
        }
    }

    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));
}