ke: add continuations

src/kern/thread.zig

@@ -76,8 +76,35 @@ pub const Thread = struct {
         /// The thread has exited.
         Terminated,
     };
-    /// Implementation-dependent context.
-    context: ki.impl.ThreadContext,
+
+    pub const Context = extern struct {
+        /// Implementation-dependent context.
+        impl: ki.impl.ThreadContext,
+
+        stack_top: usize,
+
+        pub fn init_with_stack(
+            stack: r.VAddr,
+        ) Context {
+            return .{
+                .impl = .init_with_stack(
+                    stack,
+                ),
+                .stack_top = stack,
+            };
+        }
+
+        pub fn reset(
+            self: *@This(),
+            stack_top: r.VAddr,
+            entry: *const fn (?*anyopaque) void,
+            arg: ?*anyopaque,
+        ) void {
+            self.stack_top = stack_top;
+            _ = self.impl.reset(stack_top, entry, arg);
+        }
+    };
+    context: Context,
     /// Thread lock.
     lock: ke.SpinLock,
     /// Niceness value.
@@ -120,8 +147,6 @@ pub const Thread = struct {
     /// Queue this thread is associated with.
     queue: ?*ke.Queue,
     queue_item: ?*rtl.List.Entry,
-    /// Base of the stack.
-    stack: r.VAddr,
     /// PELT load average,
     avg: ki.sched.Average,
     /// Accounting statistics.
@@ -132,32 +157,30 @@ pub const Thread = struct {
     switching: std.atomic.Value(bool),
     smr_sections: rtl.List,
     hard_affinity: ke.CpuMask,
+    continuation: ?Continuation,
+
+    const InitOpts = struct {
+        /// Entry point of the thread.
+        entry: ke.Continuation,
+        /// Top of the stack for the thread.
+        stack: r.VAddr,
+        /// Turnstile associated with the thread.
+        turnstile: *ki.turnstile.Turnstile,
+        /// Thread's priority.
+        priority: Priority,
+    };
 
     /// Initialize a thread.
-    /// - `stack`: Address of the base of the stack used by the thread.
-    /// - `stack_size`: Size of the stack
-    /// - `prio`: Base priority of the thread
-    /// - `entry`: Entry point of the thread
-    /// - `arg`: Extraneous argument to be passed to `entry`
     pub fn init(
         thread: *Thread,
-        stack: r.VAddr,
-        stack_size: usize,
-        prio: Priority,
-        entry: *const fn (?*anyopaque) void,
-        arg: ?*anyopaque,
+        opts: InitOpts,
     ) void {
         thread.* = .{
-            .context = .init(
-                stack,
-                stack_size,
-                entry,
-                arg,
-            ),
+            .context = .init_with_stack(opts.stack),
             .lock = .init(),
             .nice = 0,
-            .priority = @intFromEnum(prio),
-            .base_priority = @intFromEnum(prio),
+            .priority = @intFromEnum(opts.priority),
+            .base_priority = @intFromEnum(opts.priority),
             .inherited_prio = 0,
             .pinned = false,
             .state = .init(.Ready),
@@ -170,13 +193,13 @@ pub const Thread = struct {
             .waitblocks = undefined,
             .wait_reason = null,
             .timer = undefined,
-            .turnstile = undefined,
+            .turnstile = opts.turnstile,
             .turnstile_waiter = null,
             .turnstiles_owned = undefined,
             .waiting_on = null,
             .queue = null,
             .queue_item = null,
-            .stack = stack,
+            .continuation = opts.entry,
             .switching = .init(false),
             .avg = .{},
             .acct = .{},
@@ -187,6 +210,14 @@ pub const Thread = struct {
         thread.turnstiles_owned.init();
         thread.timer.init();
         thread.smr_sections.init();
+
+        if (opts.stack != 0) {
+            thread.context.reset(
+                opts.stack,
+                call_continuation,
+                thread,
+            );
+        }
     }
 
     pub fn priority_class(self: *Thread) Priority.Class {
@@ -226,10 +257,206 @@ pub fn exit() void {
     reaper_list.insert(@ptrCast(&curtd.runq_link.next));
 
     ki.sched.detach_load_avg(ki.sched.percpu.local(), curtd);
-    ki.sched.yield_locked();
+    ki.sched.yield_locked(Continuation.dummy);
 }
 
 /// Return the currently running thread.
 pub fn current() *ke.Thread {
     return ki.sched.percpu.local().current_thread.?;
 }
+
+// -- Stack allocation and continuations ---------------------------------------
+pub const Continuation = struct {
+    func: *const fn (_: ?*anyopaque) void,
+    arg: ?*anyopaque,
+
+    pub const dummy: Continuation = .{
+        .func = dummy_fn,
+        .arg = null,
+    };
+
+    fn dummy_fn(_: ?*anyopaque) void {}
+};
+
+const Depot = struct {
+    stacks: ?*Stack = null,
+    count: u32 = 0,
+    min: u32 = 0,
+    wma: u32 = 0,
+};
+
+pub const Stack = struct {
+    next: ?*Stack,
+    prev: ?*Stack,
+};
+
+const Cpu = struct {
+    stack_alloc_dpc: ke.Dpc,
+    depot: Depot,
+};
+
+/// Number of excess stacks allowed in the global stack depot before trimming.
+/// Note that this is in addition to the global minimum of `ncpus` stacks.
+/// Keep this low for minimal memory overhead, but potentially increased latency.
+const excess_stacks = ke.Tunable(u32, 1, "ke.thread.excess_stacks");
+
+var global_depot_lock: ke.SpinLock = .init();
+var global_depot: Depot = .{};
+
+const percpu = ke.CpuLocal(Cpu, undefined);
+const percpu_count_max = 2;
+
+const wma_unit = 256;
+
+pub var thread_stack_queue: rtl.HandoffList = .init(stack_activation);
+
+inline fn wma_mix(old: u32, new: u32) u32 {
+    // Keep 75% of old sample and 25% of new.
+    return (3 * old + new * wma_unit) / 4;
+}
+
+fn init_cpu() linksection(r.init) void {
+    const cpu = percpu.local();
+    cpu.* = .{
+        .stack_alloc_dpc = .init(stack_alloc_handler),
+        .depot = .{},
+    };
+}
+
+comptime {
+    _ = r.percpu_init_set.insert(&init_cpu);
+}
+
+pub fn call_continuation(ptr: ?*anyopaque) noreturn {
+    const td: *ke.Thread = @ptrCast(@alignCast(ptr));
+    std.debug.assert(td.continuation != null);
+
+    const cont = td.continuation.?;
+    td.continuation = null;
+
+    ki.impl.call_continuation(&td.context, cont);
+}
+
+/// Try to get a stack from the stack cache.
+pub fn stack_cache_pop() ?usize {
+    const ipl = ke.ipl.raise(.Dispatch);
+    defer ke.ipl.lower(ipl);
+
+    const cpu = &percpu.local().depot;
+
+    // Try to get a stack from the per-cpu depot.
+    if (cpu.count != 0) {
+        const entry = cpu.stacks.?;
+        cpu.stacks = entry.next;
+        cpu.count -= 1;
+        return @intFromPtr(entry) + @sizeOf(ke.Stack);
+    }
+
+    // Otherwise, try and get it from the global depot.
+    global_depot_lock.acquire_no_ipl();
+    defer global_depot_lock.release_no_ipl();
+
+    if (global_depot.count != 0) {
+        const entry = global_depot.stacks.?;
+        global_depot.stacks = entry.next;
+        global_depot.count -= 1;
+        global_depot.min = @min(global_depot.min, global_depot.count);
+        return @intFromPtr(entry) + @sizeOf(ke.Stack);
+    }
+
+    return null;
+}
+
+/// Free a stack back to the depot.
+pub fn stack_cache_free(stack: usize) void {
+    const entry: *Stack = @ptrFromInt(stack - @sizeOf(Stack));
+
+    const ipl = ke.ipl.raise(.Dispatch);
+    defer ke.ipl.lower(ipl);
+
+    const cpu = &percpu.local().depot;
+
+    if (cpu.count < percpu_count_max) {
+        entry.next = cpu.stacks;
+        cpu.stacks = entry;
+        cpu.count += 1;
+        return;
+    }
+
+    global_depot_lock.acquire_no_ipl();
+    defer global_depot_lock.release_no_ipl();
+
+    entry.next = global_depot.stacks;
+    global_depot.stacks = entry;
+    global_depot.count += 1;
+}
+
+/// Periodic updates to the stack depot.
+/// Keeps track of the number of alloc / free that happened
+/// in the last time span and reaps the excess stacks, if necessary.
+pub fn stack_cache_update() ?*Stack {
+    const ipl = global_depot_lock.acquire();
+    defer global_depot_lock.release(ipl);
+
+    global_depot.wma = wma_mix(global_depot.wma, global_depot.min);
+
+    var excess = @min(
+        global_depot.min * wma_unit,
+        global_depot.wma,
+    ) / wma_unit;
+
+    global_depot.min = global_depot.count;
+
+    // If we have more than `excess_stacks` sitting unused, then trim.
+    // Ensure there are still `ncpus` stacks sitting in the depot.
+    if (excess > excess_stacks.load()) {
+        std.debug.assert(global_depot.count >= excess);
+
+        // We want every CPU to have at least one stack sitting in the depot.
+        const floor: u32 = @truncate(ke.ncpus);
+        excess = @min(excess, global_depot.count -| floor);
+
+        const new_count = global_depot.count - excess;
+        global_depot.min = new_count;
+
+        if (excess == 0) return null;
+
+        // Splice the list from the tail to get the coldest stacks.
+        var prev = global_depot.stacks.?;
+        for (1..new_count) |_| prev = prev.next.?;
+
+        const head = prev.next;
+        prev.next = null;
+
+        global_depot.count = new_count;
+        global_depot.wma -= excess * wma_unit;
+
+        return head;
+    }
+
+    return null;
+}
+
+pub fn stacks_count() usize {
+    var ret: usize = 0;
+
+    for (0..ke.ncpus) |i| {
+        const cpu = percpu.remote(@truncate(i));
+        ret += @atomicLoad(u32, &cpu.depot.count, .monotonic);
+    }
+
+    const ipl = global_depot_lock.acquire();
+    ret += global_depot.count;
+    global_depot_lock.release(ipl);
+
+    return ret;
+}
+
+fn stack_activation(_: ?*rtl.HandoffList) void {
+    const cpu = percpu.local();
+    ke.dpc.enqueue(&cpu.stack_alloc_dpc, null);
+}
+
+fn stack_alloc_handler(_: *ke.Dpc, _: ?*anyopaque) void {
+    r.ps.private.thread.stack_activation();
+}