ex/work: dynamic thread spawning

src/ex/work.zig

@@ -1,11 +1,18 @@
-//! Work queues and work items.
-//! Work queues are used to enqueue work to be done in a thread context asynchronously.
-//! There are three global work queues for each priority:
+//! Worker pools and work items.
+//!
+//! Work items are used to enqueue work to be done in a thread context asynchronously.
+//! There are three worker priorities:
 //! - High
 //! - Normal
 //! - Low
-//! High and Normal work queues are per-cpu, and there is one global queue and thread for low-priority work.
-//! Each CPU has one worker thread, until that thread blocks, then another thread is spawned to pick up work on it.
+//! High and Normal worker pools are per-cpu, and there is one global worker pool
+//! for low-priority work.
+//!
+//! Each CPU worker pool comprises of at least 2 threads, and the low pool
+//! has 4 threads by default. An additional thread will get spawned by the
+//! pool manager thread if a queue has a backlog, up to 4 additional dynamic threads
+//! (tunable) per pool.
+
 const std = @import("std");
 const rtl = @import("rtl");
 const r = @import("root");
@@ -13,6 +20,19 @@ const ke = r.ke;
 const ps = r.ps;
 const mm = r.mm;
 
+/// Maximum number of additional threads created per pool.
+const max_dynamic_threads = ke.Tunable(u8, 4, "ex.work.max_dynamic");
+
+var pool_manager_event: ke.Event = undefined;
+
+fn min_threads_for_prio(prio: WorkItem.Priority) usize {
+    return if (prio == .Low) 4 else 2;
+}
+
+fn max_threads_for_prio(prio: WorkItem.Priority) usize {
+    return min_threads_for_prio(prio) + max_dynamic_threads.load();
+}
+
 /// Represents a single work item.
 pub const WorkItem = struct {
     func: *const fn (arg: ?*anyopaque) void,
@@ -55,36 +75,34 @@ pub const WorkItem = struct {
     }
 };
 
+const Context = rtl.TaggedPtr(Pool);
+
 /// Represents a thread worker pool and its queue.
 const Pool = struct {
-    const Worker = struct {
-        link: rtl.List.Entry,
-        td: *ps.Thread,
-    };
     /// The enqueued items.
     queue: ke.Queue,
     /// Priority of the pool.
     prio: WorkItem.Priority,
     /// CPU this pool is on.
     cpu: ?u32,
-    lock: ke.SpinLock,
+    total_threads: std.atomic.Value(usize),
 
-    pub fn init(self: *Pool, prio: WorkItem.Priority, cpu: ?u32) !void {
-        self.queue.init(1);
+    fn init(self: *Pool, prio: WorkItem.Priority, cpu: ?u32) !void {
+        self.queue.init(if (prio == .Low) ke.ncpus else 1);
         self.prio = prio;
         self.cpu = cpu;
-        self.lock = .init();
 
-        // TODO: dynamic worker pool.
-        // We don't support thread exits yet, so add that first.
-        // Do linux-style where worker threads themselves spawn new threads.
-        const how_many: usize = if (prio == .Low) 4 else 2;
+        const how_many = min_threads_for_prio(prio);
+
+        self.total_threads = .init(how_many);
 
         for (0..how_many) |_| {
+            const ctx: Context = .init(self, 0);
+
             var td = try ps.thread.create_kernel(
                 WorkItem.Priority.to_sched_prio(self.prio),
                 work_loop,
-                self,
+                @ptrFromInt(ctx.value),
             );
 
             if (cpu != null) {
@@ -95,6 +113,53 @@ const Pool = struct {
             ke.sched.enqueue(&td.kern);
         }
     }
+
+    /// Return whether or not a new worker thread should be spawned on this pool.
+    fn should_grow(self: *Pool) bool {
+        // NOTE: lockless racy check whether or not the queue is empty
+        // and its active count.
+        // This is fine as this is best effort anyway and is a heuristic.
+        const is_empty = @atomicLoad(
+            *rtl.List.Entry,
+            &self.queue.items.head.next,
+            .monotonic,
+        ) == &self.queue.items.head;
+
+        const active = @atomicLoad(usize, &self.queue.active, .monotonic);
+        const threads = self.total_threads.load(.monotonic);
+
+        if (threads < max_threads_for_prio(self.prio) and !is_empty and active < self.queue.max_active) {
+            // We have room to grow, the list still has items and
+            // not enough runnable threads are present, create an
+            // additional thread to ease the backlog.
+            return true;
+        }
+
+        return false;
+    }
+
+    /// Spawn a dynamic worker thread on this pool.
+    fn grow(self: *Pool, cpu: ?u32) !void {
+        if (self.total_threads.load(.monotonic) >= max_threads_for_prio(self.prio))
+            return;
+
+        _ = self.total_threads.fetchAdd(1, .monotonic);
+
+        const ctx: Context = .init(self, 1);
+
+        var td = try ps.thread.create_kernel(
+            WorkItem.Priority.to_sched_prio(self.prio),
+            work_loop,
+            @ptrFromInt(ctx.value),
+        );
+
+        if (cpu != null) {
+            td.kern.last_cpu = cpu;
+            td.kern.pinned = true;
+        }
+
+        ke.sched.enqueue(&td.kern);
+    }
 };
 
 /// Per-CPU work data.
@@ -111,10 +176,19 @@ const percpu = ke.CpuLocal(PerCpu, undefined);
 var low: Pool = undefined;
 
 fn work_loop(p: ?*anyopaque) void {
-    const pool: *Pool = @ptrCast(@alignCast(p));
+    const ctx: Context = .{ .value = @intFromPtr(p) };
+    const dynamic = ctx.get_tag() == 1;
 
     while (true) {
-        const item = pool.queue.remove();
+        const timeout: ?r.Nanoseconds = if (dynamic) std.time.ns_per_s * 5 else null;
+        const item = ctx.get_ptr().queue.remove(timeout) catch {
+            _ = ctx.get_ptr().total_threads.fetchSub(1, .monotonic);
+
+            // We haven't had work in a while, exit.
+            ps.thread.exit();
+            return;
+        };
+
         const work: *WorkItem = @fieldParentPtr("link", item);
 
         work.enqueued.store(false, .release);
@@ -158,10 +232,42 @@ pub fn enqueue(item: *WorkItem) void {
     }
 
     pool.queue.insert(&item.link, .Tail);
+
+    if (pool.should_grow()) {
+        pool_manager_event.signal();
+    }
+}
+
+fn pool_manager(_: ?*anyopaque) void {
+    while (true) {
+        _ = ke.wait.wait_one(&pool_manager_event.hdr, null) catch unreachable;
+
+        // We have been signaled, go through every pool and grow it if needed.
+        // Failure to grow is ignored, it is only a latency problem and should be fixed
+        // eventually as memory is freed.
+        for (0..ke.ncpus) |i| {
+            const cpu: u32 = @intCast(i);
+            const c = percpu.remote(cpu);
+
+            if (c.high.should_grow()) {
+                _ = c.high.grow(cpu) catch {};
+            }
+
+            if (c.normal.should_grow()) {
+                _ = c.normal.grow(cpu) catch {};
+            }
+        }
+
+        if (low.should_grow()) {
+            _ = low.grow(null) catch {};
+        }
+    }
 }
 
 /// Initialize the work subsystem.
 pub fn init() !void {
+    pool_manager_event.init(.Synchronization);
+
     for (0..ke.ncpus) |cpu| {
         const c = percpu.remote(@intCast(cpu));
 
@@ -170,4 +276,12 @@ pub fn init() !void {
     }
 
     try low.init(.Low, null);
+
+    var td = ps.thread.create_kernel(
+        ke.Thread.Priority.default,
+        pool_manager,
+        null,
+    ) catch @panic("Failed to create pool manager");
+
+    ke.sched.enqueue(&td.kern);
 }

src/kern/queue.zig

@@ -54,7 +54,7 @@ pub const Queue = struct {
 
     /// Remove the item at the head.
     /// This blocks until an item is actually popped.
-    pub fn remove(self: *Queue) *rtl.List.Entry {
+    pub fn remove(self: *Queue, timeout: ?r.Nanoseconds) !*rtl.List.Entry {
         const ipl = self.hdr.lock.acquire();
         const td = ki.sched.percpu.local().current_thread.?;
 
@@ -75,7 +75,7 @@ pub const Queue = struct {
         self.hdr.lock.release(ipl);
 
         // Wait until the queue has something for us.
-        _ = ke.wait.wait_one(&self.hdr, null) catch unreachable;
+        _ = try ke.wait.wait_one(&self.hdr, timeout);
 
         // Grab the queue item and set it to null.
         const ret = td.queue_item orelse unreachable;