ke: add continuations

src/ex/console.zig

@@ -126,8 +126,8 @@ pub fn init() void {
 
     const td = ps.thread.create_kernel(
         .Default,
-        worker,
+        .{ .func = worker, .arg = null },
-        null,
+        true,
     ) catch @panic("handle me");
     ke.sched.enqueue(&td.kern);
 }

src/ex/init.zig

@@ -26,7 +26,8 @@ pub fn init() void {
         if (rtl.cmdline.get_string(cline, "test")) |tst| {
             // Run a test if specified.
             if (r.tests.tests.get(tst)) |func| {
-                const t = ps.thread.create_kernel(.Default, func, r.boot_info) catch @panic("oom");
+                const t = ps.thread.create_kernel(.Default, .{ .func = func, .arg = r.boot_info }, true) catch
+                    @panic("oom");
                 ke.sched.enqueue(&t.kern);
             } else {
                 std.log.info("invalid test: \"{s}\"", .{tst});

src/ex/work.zig

@@ -124,14 +124,15 @@ const Pool = struct {
 
             var td = try ps.thread.create_kernel(
                 WorkItem.Priority.to_sched_prio(self.prio),
-                work_loop,
+                .{ .func = work_loop, .arg = @ptrFromInt(ctx.value) },
-                @ptrFromInt(ctx.value),
+                true,
             );
 
             td.kern.nice = -ke.Priority.nice_max;
 
-            if (cpu) |c| {
+            if (cpu != null) {
-                ke.sched.pin_on(&td.kern, c);
+                td.kern.last_cpu = cpu;
+                td.kern.pinned = true;
             }
 
             ke.sched.enqueue(&td.kern);
@@ -176,12 +177,13 @@ const Pool = struct {
 
         var td = try ps.thread.create_kernel(
             WorkItem.Priority.to_sched_prio(self.prio),
-            work_loop,
+            .{ .func = work_loop, .arg = @ptrFromInt(ctx.value) },
-            @ptrFromInt(ctx.value),
+            true,
         );
 
-        if (cpu) |c| {
+        if (cpu != null) {
-            ke.sched.pin_on(&td.kern, c);
+            td.kern.last_cpu = cpu;
+            td.kern.pinned = true;
         }
 
         ke.sched.enqueue(&td.kern);
@@ -313,8 +315,8 @@ pub fn init() !void {
 
     var td = ps.thread.create_kernel(
         .Default,
-        pool_manager,
+        .{ .func = pool_manager, .arg = null },
-        null,
+        true,
     ) catch @panic("Failed to create pool manager");
 
     ke.sched.enqueue(&td.kern);

src/kern/amd64/impl.zig

@@ -1,5 +1,6 @@
 const r = @import("root");
 const ke = r.ke;
+const ki = ke.private;
 const pl = r.pl;
 const amd64 = r.arch;
 
@@ -23,19 +24,50 @@ const ThreadFrame = extern struct {
     rip: u64 align(1),
 };
 
+const Context = ke.Thread.Context;
+
 extern fn asm_thread_entry() void;
 
-extern fn do_context_switch(
+extern fn amd64_context_load(td: *ThreadContext) callconv(.c) void;
-    old: *ThreadContext,
+
-    new: *ThreadContext,
+extern fn amd64_context_switch(
+    old: *Context,
+    new: *Context,
     lock: *u8,
     switching: *bool,
 ) callconv(.c) void;
-extern fn do_context_load(td: *ThreadContext) callconv(.c) void;
+
+extern fn amd64_context_switch_cont_n(
+    old: *Context,
+    new: *Context,
+    lock: *u8,
+    switching: *bool,
+) callconv(.c) void;
+
+extern fn amd64_call_continuation(
+    td: *ke.Thread,
+    ctx: *Context,
+    func: usize,
+    arg: ?*anyopaque,
+) callconv(.c) noreturn;
+
+export fn ki_free_old_stack(stack_top: usize) callconv(.c) void {
+    ki.thread.stack_cache_free(stack_top);
+}
+
+/// Trampoline to match the zig ABI.
+export fn ki_continuation_trampoline(
+    func: usize,
+    arg: ?*anyopaque,
+) noreturn {
+    ke.ipl.lower(.Passive);
+    const f: *const fn (?*anyopaque) void = @ptrFromInt(func);
+    f(arg);
+    @panic("Continuation returned");
+}
 
 pub const ThreadContext = extern struct {
-    rsp: u64 align(1),
+    rsp: u64 align(1) = 0xDEADC0D3,
-    frame: ThreadFrame,
 
     export fn thread_entry(entry: usize, arg: usize) callconv(.c) void {
         const entry_fn: *const fn (?*anyopaque) void = @ptrFromInt(entry);
@@ -43,14 +75,21 @@ pub const ThreadContext = extern struct {
         entry_fn(@ptrFromInt(arg));
     }
 
-    pub fn init(
+    pub fn init_with_stack(
         stack: r.VAddr,
-        stack_size: usize,
+    ) @This() {
+        return .{
+            .rsp = stack,
+        };
+    }
+
+    pub fn reset(
+        self: *@This(),
+        stack_top: r.VAddr,
         entry: *const fn (?*anyopaque) void,
         arg: ?*anyopaque,
     ) @This() {
-        var ctx: @This() = undefined;
+        var sp: usize = stack_top;
-        var sp: usize = stack + stack_size;
 
         sp = sp & ~@as(usize, 15);
         sp = sp - @sizeOf(ThreadFrame);
@@ -60,31 +99,56 @@ pub const ThreadContext = extern struct {
 
         frame.* = std.mem.zeroes(ThreadFrame);
 
-        ctx.rsp = @intFromPtr(frame);
+        self.rsp = @intFromPtr(frame);
         frame.rip = @intFromPtr(&asm_thread_entry);
         frame.r12 = @intFromPtr(entry);
         frame.r13 = @intFromPtr(arg);
 
-        return ctx;
+        return self.*;
     }
 
-    pub fn switch_to(
+    pub fn load(self: *ThreadContext) void {
-        self: *ThreadContext,
+        amd64_context_load(self);
-        new: *ThreadContext,
+    }
-    ) callconv(.c) void {
+};
-        const thread: *ke.Thread = @alignCast(@fieldParentPtr("context", self));
+
-        do_context_switch(
+pub fn switch_normal_to_normal(
-            self,
+    old: *Context,
+    new: *Context,
+) void {
+    const thread: *ke.Thread = @alignCast(@fieldParentPtr("context", old));
+    amd64_context_switch(
+        old,
         new,
         &thread.lock.inner.locked.raw,
         &thread.switching.raw,
     );
-    }
+}
 
-    pub fn load(self: *ThreadContext) callconv(.c) void {
+pub fn switch_cont_to_normal(
-        do_context_load(self);
+    old: *Context,
-    }
+    new: *Context,
-};
+) void {
+    const td: *ke.Thread = @alignCast(@fieldParentPtr("context", old));
+
+    amd64_context_switch_cont_n(
+        old,
+        new,
+        &td.lock.inner.locked.raw,
+        &td.switching.raw,
+    );
+}
+
+pub fn call_continuation(ctx: *Context, continuation: ke.Continuation) noreturn {
+    const td: *ke.Thread = @alignCast(@fieldParentPtr("context", ctx));
+
+    amd64_call_continuation(
+        td,
+        ctx,
+        @intFromPtr(continuation.func),
+        continuation.arg,
+    );
+}
 
 inline fn percpu_ptr_for(variable: anytype, cpu: u32) @TypeOf(variable) {
     return @ptrFromInt(@intFromPtr(variable) +% cpu_offsets[cpu]);

src/kern/amd64/locore.s

@@ -133,8 +133,8 @@ __interrupt_vectors:
 .endr
 
 
-.global do_context_switch
+.global amd64_context_switch
-do_context_switch:
+amd64_context_switch:
     /* Save callee-saved registers. */
     push %r15
     push %r14
@@ -166,8 +166,8 @@ do_context_switch:
     pop %r15
     ret
 
-.global do_context_load
+.global amd64_context_load
-do_context_load:
+amd64_context_load:
     /* Load new stack pointer. */
     mov (%rdi), %rsp
 
@@ -180,10 +180,59 @@ do_context_load:
     pop %r15
     ret
 
+.extern ki_free_old_stack
+.global amd64_context_switch_cont_n
+amd64_context_switch_cont_n:
+    /* Load new stack pointer. */
+    movq (%rsi), %rsp
+
+    movq %rdx, %r12
+    movq %rcx, %r13
+
+    /* Free the old stack. */
+    movq 8(%rdi), %rbx
+    movq $0, 8(%rdi)
+    movq %rbx, %rdi
+    call ki_free_old_stack
+
+    /* Set switching to false. */
+    xorl %eax, %eax
+    xchgb %al, (%r13)
+
+    /* Unlock passed lock. */
+    xorl %eax, %eax
+    xchgb %al, (%r12)
+
+    /* Restore callee-saved registers. */
+    pop %rbp
+    pop %rbx
+    pop %r12
+    pop %r13
+    pop %r14
+    pop %r15
+    ret
+
+
+.extern ki_continuation_trampoline
+.global amd64_call_continuation
+amd64_call_continuation:
+    /* Reset stack pointer. */
+    movq 8(%rsi), %rsp
+
+    /* Terminate stack frame. */
+    xorq %rbp, %rbp
+
+    /* Pass func and arg to the trampoline. */
+    mov %rdx, %rdi
+    mov %rcx, %rsi
+
+    /* Jump to the function. */
+    jmp ki_continuation_trampoline
+
 .extern thread_entry
 .global asm_thread_entry
 asm_thread_entry:
     mov %r12, %rdi
     mov %r13, %rsi
-    call thread_entry
+    jmp thread_entry
     ud2

src/kern/init.zig

@@ -12,22 +12,22 @@ export fn kmain(boot_info: *r.BootInfo) callconv(.c) void {
     r.boot_info = boot_info;
 
     // Immediately initialize the context of what will eventually become
-    // our idle thread. We keep running on the loader stack up until we switch
+    // our idle thread. At this point we're already running on the kernel stack.
-    // into the thread.
+    thread0.init(.{
-    thread0.init(
+        .stack = boot_info.kernel_stack + boot_info.kernel_stack_size,
-        boot_info.kernel_stack,
+        .priority = .IdleThread,
-        boot_info.kernel_stack_size,
+        .entry = undefined,
-        .IdleThread,
+        // Shouldn't block
-        init,
+        .turnstile = undefined,
-        null,
+    });
-    );
+
+    thread0.continuation = null;
     thread0.pinned = true;
 
-    // Jump into the thread.
+    init();
-    thread0.context.load();
 }
 
-pub fn init(_: ?*anyopaque) linksection(r.init) void {
+fn init() linksection(r.init) void {
     ke.ncpus = 1;
     ki.impl.early_init();
     ki.tunable.init();

src/kern/root.zig

@@ -22,6 +22,8 @@ pub const Mutex = p.mutex.Mutex;
 pub const ShootdownState = p.shootdown.ShootdownState;
 pub const Tunable = p.tunable.Tunable;
 pub const Priority = p.thread.Priority;
+pub const Stack = p.thread.Stack;
+pub const Continuation = p.thread.Continuation;
 
 // === Exported Interfaces ===
 pub const panic = p.panic.panic;
@@ -39,6 +41,9 @@ pub const thread = struct {
     pub const exit = p.thread.exit;
     pub const current = p.thread.current;
     pub const reaper_list = &p.thread.reaper_list;
+    pub const stack_cache_update = p.thread.stack_cache_update;
+    pub const thread_stack_queue = &p.thread.thread_stack_queue;
+    pub const call_continuation = p.thread.call_continuation;
 };
 
 pub const dpc = struct {

src/kern/sched.zig

@@ -1061,6 +1061,15 @@ fn enqueue_on_cpu(c: u32, td: *ke.Thread) void {
 
     // queues_lock dropped
 }
+fn allocate_stack(td: *ke.Thread) ?usize {
+    if (ki.thread.stack_cache_pop()) |stack| {
+        return stack;
+    }
+
+    // Enqueue the thread for allocation
+    ki.thread.thread_stack_queue.insert(@ptrCast(&td.runq_link.next));
+    return null;
+}
 
 // Execute a context switch on `cpu`.
 fn do_switch(cpu: *PerCpu, cur: *ke.Thread, next: *ke.Thread) void {
@@ -1078,28 +1087,76 @@ fn do_switch(cpu: *PerCpu, cur: *ke.Thread, next: *ke.Thread) void {
         account(next, now_us, .Queued);
     }
 
+    if (next.context.stack_top == 0) {
+        std.debug.assert(next.continuation != null);
+    }
+
     next.state.store(.Running, .monotonic);
 
-    // Do machine-dependent switch.
+    if (cur.continuation != null) {
-    cur.context.switch_to(&next.context);
+        if (next.continuation) |c| {
+            // Both the incoming thread and outgoing thread have continuations
+            // set. Do stack hand-off.
+            next.context.stack_top = cur.context.stack_top;
+            cur.context.stack_top = 0;
+            next.continuation = null;
+
+            cur.switching.store(false, .release);
+            cur.lock.release_no_ipl();
+
+            ki.impl.call_continuation(&next.context, c);
+        }
+
+        ki.impl.switch_cont_to_normal(&cur.context, &next.context);
+    } else {
+        ki.impl.switch_normal_to_normal(&cur.context, &next.context);
+    }
 }
 
 /// Handle preemption.
 /// This is called in DPC dispatch when it notices `next_thread` is set.
 pub fn handle_preemption(cpu: *PerCpu) void {
+    const cur = cpu.current_thread.?;
+    var next: *ke.Thread = undefined;
+
+    while (true) {
         cpu.queues_lock.acquire_no_ipl();
 
-    const cur = cpu.current_thread.?;
+        next = cpu.next_thread orelse select_thread(cur, cpu, null) orelse {
-
-    const next = cpu.next_thread orelse {
             cpu.queues_lock.release_no_ipl();
             return;
         };
 
         cpu.next_thread = null;
+        cpu.queues_lock.release_no_ipl();
+
+        while (next.switching.load(.monotonic) == true) {
+            std.atomic.spinLoopHint();
+        }
+
+        if (next.context.stack_top == 0) {
+            // This thread has no stack, try allocating one from the stack cache.
+            // Otherwise, the thread will be waiting for its stack to get
+            // allocated and we pick another one in the meantime.
+            const new_stack = allocate_stack(next);
+
+            if (new_stack) |new| {
+                next.context.reset(new, ke.thread.call_continuation, next);
+                break;
+            }
+
+            detach_load_avg(cpu, next);
+            continue;
+        }
+
+        // Found a thread with a a stack.
+        break;
+    }
+
+    std.debug.assert(cur != next);
+
     cpu.current_thread_prio.store(next.priority, .monotonic);
 
-    cpu.queues_lock.release_no_ipl();
     cur.lock.acquire_no_ipl();
 
     cur.switching.store(true, .monotonic);
@@ -1190,7 +1247,7 @@ pub fn enqueue(td: *ke.Thread) void {
 }
 
 /// Yield on the current CPU.
-pub fn yield() void {
+pub fn yield(continuation: ?ke.Continuation) void {
     const ipl = ke.ipl.raise(.Dispatch);
     const td = percpu.local().current_thread.?;
 
@@ -1205,71 +1262,108 @@ pub fn yield() void {
         cpu.queues_lock.release_no_ipl();
     }
 
-    yield_locked();
+    yield_locked(continuation);
 
     // td lock released
     ke.ipl.lower(ipl);
 }
 
 /// Yield the current CPU with the current thread already locked at Dispatch IPL.
-pub fn yield_locked() void {
+pub fn yield_locked(continuation: ?ke.Continuation) void {
     const sched_cpu = percpu.local();
 
     sched_cpu.queues_lock.acquire_no_ipl();
 
-    const cur = sched_cpu.current_thread;
+    const cur = sched_cpu.current_thread orelse unreachable;
     var next = sched_cpu.next_thread;
 
+    cur.continuation = continuation;
+
+    while (true) {
         if (next != null) {
             sched_cpu.next_thread = null;
+            sched_cpu.preemption_reason = .None;
         } else {
             // Pick a new thread to run.
             next = select_thread(null, sched_cpu, null);
         }
 
+        sched_cpu.queues_lock.release_no_ipl();
+
         if (next == null) {
             // Nothing to run, go idle.
             next = sched_cpu.idle_thread;
             sched_cpu.steal_work = true;
+            break;
         }
 
-    if (next) |n| {
-        sched_cpu.current_thread_prio.store(n.priority, .monotonic);
-    }
-
-    sched_cpu.queues_lock.release_no_ipl();
-
-    if (next != null and cur.? != next.?) {
-        // Switch into the thread. Ensure it's not switching.
         while (next.?.switching.load(.monotonic) == true) {
             std.atomic.spinLoopHint();
         }
 
-        do_switch(sched_cpu, cur.?, next.?);
+        if (next) |n| {
+            const curstack = n.context.stack_top;
+            const can_handoff = cur.continuation != null;
+
+            if (curstack != 0 or can_handoff) {
+                break;
+            }
+
+            if (curstack == 0) {
+                const new_stack = allocate_stack(n);
+
+                if (new_stack) |new| {
+                    n.context.reset(new, ke.thread.call_continuation, n);
+                    break;
+                }
+
+                detach_load_avg(sched_cpu, n);
+                next = null;
+            }
+        }
+
+        sched_cpu.queues_lock.acquire_no_ipl();
+    }
+
+    const n = next orelse unreachable;
+
+    sched_cpu.current_thread_prio.store(n.priority, .monotonic);
+
+    if (cur != n) {
+        // Switch into the thread. Ensure it's not switching.
+        while (n.switching.load(.monotonic) == true) {
+            std.atomic.spinLoopHint();
+        }
+        do_switch(sched_cpu, cur, n);
     } else {
-        // Ensure curthread is not marked as selected.
+        // Same thread.
-        cur.?.state.store(.Running, .monotonic);
+        const cont = cur.continuation;
-        cur.?.switching.store(false, .monotonic);
+        cur.continuation = null;
-        cur.?.lock.release_no_ipl();
+
+        cur.state.store(.Running, .monotonic);
+        cur.switching.store(false, .monotonic);
+        cur.lock.release_no_ipl();
+
+        if (cont) |c| ki.impl.call_continuation(&cur.context, c);
     }
 
     // cur lock dropped
 }
 
 /// Block the currently running thread.
-pub fn block() void {
+pub fn block(cont: ?ke.Continuation) void {
     const ipl = ke.ipl.raise(.Dispatch);
     const td = percpu.local().current_thread.?;
 
     td.lock.acquire_no_ipl();
-    block_locked(td);
+    block_locked(td, cont);
 
     // td lock released
     ke.ipl.lower(ipl);
 }
 
 /// Block the currently running thread with its lock held.
-pub fn block_locked(curtd: *ke.Thread) void {
+pub fn block_locked(curtd: *ke.Thread, cont: ?ke.Continuation) void {
     curtd.state.store(.Blocked, .monotonic);
     curtd.runq = null;
 
@@ -1285,7 +1379,7 @@ pub fn block_locked(curtd: *ke.Thread) void {
         curtd.avg.est = (3 * curtd.avg.est + sample) / 4;
     }
 
-    yield_locked();
+    yield_locked(cont);
 }
 
 /// Unblock a thread.
@@ -1629,7 +1723,8 @@ pub fn idle(_: ?*anyopaque) noreturn {
         );
 
         if (cpu.load.load(.monotonic) != 0 or next != null) {
-            ke.sched.yield();
+            std.atomic.spinLoopHint();
+            ke.sched.yield(null);
             continue;
         }
 

src/kern/thread.zig

@@ -76,8 +76,35 @@ pub const Thread = struct {
         /// The thread has exited.
         Terminated,
     };
+
+    pub const Context = extern struct {
         /// Implementation-dependent context.
-    context: ki.impl.ThreadContext,
+        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,
+        opts: InitOpts,
-        stack_size: usize,
-        prio: Priority,
-        entry: *const fn (?*anyopaque) void,
-        arg: ?*anyopaque,
     ) void {
         thread.* = .{
-            .context = .init(
+            .context = .init_with_stack(opts.stack),
-                stack,
-                stack_size,
-                entry,
-                arg,
-            ),
             .lock = .init(),
             .nice = 0,
-            .priority = @intFromEnum(prio),
+            .priority = @intFromEnum(opts.priority),
-            .base_priority = @intFromEnum(prio),
+            .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();
+}

src/kern/wait.zig

@@ -101,6 +101,7 @@ pub const Status = enum(u8) {
 const Options = struct {
     timeout: ?r.Nanoseconds = null,
     waitblocks: ?[*]WaitBlock = null,
+    continuation: ?ke.Continuation = null,
 };
 
 /// Wait for the provided object to be signaled.
@@ -140,6 +141,10 @@ pub fn wait_any(objects: []*DispatchHeader, reason: []const u8, opts: Options) !
     var satisfier: ?usize = null;
     curtd.wait_status.store(.InProgress, .release);
 
+    if (has_timeout) {
+        timer.init();
+    }
+
     for (0..total_count) |i| {
         const is_timer = has_timeout and i == timer_i;
         var obj = if (is_timer) &timer.hdr else objects[i];
@@ -223,7 +228,7 @@ pub fn wait_any(objects: []*DispatchHeader, reason: []const u8, opts: Options) !
 
         curtd.wait_reason = reason;
         // We're good, now actually block.
-        ki.sched.block_locked(curtd);
+        ki.sched.block_locked(curtd, opts.continuation);
     } else {
         queue = null;
 

src/loader/jump.s

@@ -1,4 +1,5 @@
 .globl jump_to_kernel
 jump_to_kernel:
+	mov %rdx, %rsp              /* Load new stack */
 	xorq %rbp, %rbp 			/* Call chain ends here */
 	jmp *%rsi					/* rdi = loader_info */

src/loader/main.zig

@@ -4,7 +4,7 @@ const std = @import("std");
 const Entry = r.BootInfo.MemMap.Entry;
 const stack_size = 1024 * 16; // should be enough
 
-extern fn jump_to_kernel(loader_info: usize, entry: usize) void;
+extern fn jump_to_kernel(loader_info: usize, entry: usize, stack_top: usize) void;
 
 fn load_kernel(kernel: *anyopaque) usize {
     const elf: *std.elf.Elf64.Ehdr = @ptrCast(@alignCast(kernel));
@@ -103,5 +103,6 @@ pub fn loader_main(kernel: *anyopaque) void {
     jump_to_kernel(
         @intFromPtr(&r.loader_info),
         entry,
+        stack + stack_size,
     );
 }

src/platform/pc/smp.zig

@@ -38,16 +38,23 @@ fn ap_entry(cpu_id: u32) callconv(.c) noreturn {
 }
 
 fn make_thread(
-    entrypoint: *const fn (?*anyopaque) void,
     stack: usize,
-    size: usize,
+    cpu: u32,
 ) *ke.Thread {
     const td = mm.zone.gpa.create(ke.Thread) catch
         @panic("Failed to allocate thread for AP");
 
-    td.init(stack, size, .IdleThread, entrypoint, null);
+    td.init(.{
+        .stack = stack,
+        .priority = .IdleThread,
+        .entry = undefined,
+        // Shouldn't block.
+        .turnstile = undefined,
+    });
 
-    ke.sched.pin(td);
+    td.continuation = null;
+
+    ke.sched.pin_on(td, cpu);
     return td;
 }
 
@@ -139,9 +146,8 @@ pub fn init() linksection(r.init) void {
 
         start_stack.remote(@intCast(cpu_id)).* = stack_top & ~@as(usize, 15);
         start_thread.remote(@intCast(cpu_id)).* = make_thread(
-            @ptrCast(&ap_entry),
+            stack_top,
-            stack_top - stack_size,
+            @intCast(cpu_id),
-            stack_size,
         );
 
         rtl.barrier.fence(.release);

src/proc/root.zig

@@ -12,6 +12,7 @@ pub const Thread = p.thread.Thread;
 pub const thread = struct {
     pub const create_kernel = p.thread.create_kernel;
     pub const exit = p.thread.exit;
+    pub const kernel_thread_stack_size = p.thread.kernel_thread_stack_size;
 };
 
 pub const init = p.init;

src/proc/thread.zig

@@ -8,9 +8,10 @@ const ke = r.ke;
 const ki = ke.private;
 const mm = r.mm;
 
-const kernel_thread_stack_size = r.kib(16);
+pub const kernel_thread_stack_size = r.kib(16);
 
 var reaper_item: ex.WorkItem = undefined;
+var stack_item: ex.WorkItem = undefined;
 var thread_zone: mm.zone.TypedZone(Thread) = undefined;
 
 /// Higher-level wrapper over a thread.
@@ -33,7 +34,6 @@ fn reap_thread(obj: *anyopaque, _: ?*anyopaque) void {
     const td: *Thread = @fieldParentPtr("kern", ketd);
 
     psp.turnstile_zone.destroy(td.kern.turnstile);
-    mm.heap.free(td.kern.stack, kernel_thread_stack_size);
     thread_zone.destroy(td);
 }
 
@@ -41,14 +41,19 @@ fn reaper_fn(_: ?*anyopaque) void {
     ke.thread.reaper_list.process(reap_thread, null);
 }
 
-fn activation(_: *rtl.HandoffList) void {
+fn reaper_activation(_: *rtl.HandoffList) void {
     ex.work.enqueue(&reaper_item);
 }
 
+pub fn stack_activation() void {
+    ex.work.enqueue(&stack_item);
+}
+
 pub fn init() void {
     thread_zone.init("thread", .{});
-    ke.thread.reaper_list.* = .init(activation);
+    ke.thread.reaper_list.* = .init(reaper_activation);
     reaper_item.init(.High, reaper_fn, null);
+    stack_item.init(.High, stack_fn, null);
 }
 
 /// Create a kernel thread.
@@ -57,14 +62,25 @@ pub fn init() void {
 /// - `arg`: argument passed to `entry`.
 pub fn create_kernel(
     prio: ke.Priority,
-    entry: *const fn (arg: ?*anyopaque) void,
+    continuation: ke.Continuation,
-    arg: ?*anyopaque,
+    start_with_stack: bool,
 ) !*Thread {
     var td = try thread_zone.create();
-    const stack = try mm.heap.alloc(kernel_thread_stack_size, .WaitForMemory);
 
-    td.kern.init(@intFromPtr(stack), kernel_thread_stack_size, prio, entry, arg);
+    const stack = if (start_with_stack)
-    td.kern.turnstile = try psp.turnstile_zone.create();
+        @intFromPtr(try mm.heap.alloc(kernel_thread_stack_size, .WaitForMemory)) +
+            kernel_thread_stack_size
+    else
+        0;
+
+    const turnstile = try psp.turnstile_zone.create();
+
+    td.kern.init(.{
+        .stack = stack,
+        .turnstile = turnstile,
+        .entry = continuation,
+        .priority = prio,
+    });
 
     return td;
 }
@@ -74,3 +90,24 @@ pub fn create_kernel(
 pub fn exit() void {
     ke.thread.exit();
 }
+
+fn allocate_stack(obj: *anyopaque, _: ?*anyopaque) void {
+    const link: **rtl.List.Entry = @ptrCast(@alignCast(obj));
+
+    // Bleh
+    const entry: *rtl.List.Entry = @fieldParentPtr("next", link);
+    const thread: *ke.Thread = @fieldParentPtr("runq_link", entry);
+
+    const stack = mm.heap.alloc(kernel_thread_stack_size, .WaitForMemory) catch
+        unreachable;
+
+    const stack_top = @intFromPtr(stack) + kernel_thread_stack_size;
+
+    thread.context.reset(stack_top, ke.thread.call_continuation, thread);
+
+    ke.sched.enqueue(thread);
+}
+
+fn stack_fn(_: ?*anyopaque) void {
+    ke.thread.thread_stack_queue.process(allocate_stack, null);
+}

src/tests/fireworks.zig

@@ -12,6 +12,7 @@ var fb_height: usize = 0;
 var fb_pitch: usize = 0;
 var fb_bpp: usize = 0;
 var inited = std.atomic.Value(bool).init(false);
+var particle_count: std.atomic.Value(usize) = .init(0);
 
 const background_color = 0x09090F;
 
@@ -91,24 +92,61 @@ const FireworkData = struct {
     vel_x: i64,
     vel_y: i64,
     explosion_range: i32,
+    expire_in: i64,
+    i: u32,
 };
 
 fn get_random_color() u32 {
     return @intCast((rand() + 0x808080) & 0xFFFFFF); // Random pastel color
 }
 
-fn perform_delay(ms: usize) void {
+fn sleep(ms: usize, continuation: ?ke.Continuation) void {
-    var timer: ke.Timer = undefined;
+    _ = ke.wait.wait_any(&.{}, "sleep", .{
-    timer.init();
+        .timeout = .from(r.Milliseconds.init(ms)),
-    ke.timer.set(&timer, .init(std.time.ns_per_ms * ms), .{});
+        .continuation = continuation,
-    _ = ke.wait.wait_one(&timer.hdr, "sleep", .{}) catch unreachable;
+    }) catch {
+        return;
+    };
+    @panic("Not a timeout?");
+}
+
+const particle_delay = 16;
+
+fn particle_continuation(param: ?*anyopaque) void {
+    const data: *FireworkData = @ptrCast(@alignCast(param));
+    plot_pixel(@intCast(data.x), @intCast(data.y), background_color);
+
+    data.i += particle_delay;
+
+    if (data.i >= data.expire_in) {
+        mm.zone.gpa.destroy(data);
+
+        _ = particle_count.fetchSub(1, .monotonic);
+        ps.thread.exit();
+    }
+
+    data.act_x += @divTrunc(data.vel_x * particle_delay, 1000);
+    data.act_y += @divTrunc(data.vel_y * particle_delay, 1000);
+    data.x = fp_to_int(data.act_x);
+    data.y = fp_to_int(data.act_y);
+
+    data.vel_y += @divTrunc(int_to_fp(10) * particle_delay, 1000);
+
+    plot_pixel(@intCast(data.x), @intCast(data.y), data.color);
+    sleep(particle_delay, .{
+        .func = particle_continuation,
+        .arg = data,
+    });
 }
 
 fn particle(param: ?*anyopaque) void {
     const parent_data: *FireworkData = @ptrCast(@alignCast(param));
-    var data: FireworkData = undefined;
 
-    data = std.mem.zeroes(FireworkData);
+    const data = mm.zone.gpa.create(FireworkData) catch unreachable;
+
+    data.* = std.mem.zeroes(FireworkData);
+
+    _ = particle_count.fetchAdd(1, .monotonic);
 
     data.x = parent_data.x;
     data.y = parent_data.y;
@@ -124,42 +162,37 @@ fn particle(param: ?*anyopaque) void {
 
     const expire_in = 2000 + (@rem(rand(), 1000));
 
+    data.expire_in = expire_in;
+    data.i = 0;
+
     data.color = get_random_color();
 
-    var i: i32 = 0;
-    var t: i32 = 0;
-
-    while (i < expire_in) {
     plot_pixel(@intCast(data.x), @intCast(data.y), data.color);
 
-        const delay = @as(i32, 16) + @intFromBool(t != 0);
+    sleep(particle_delay, .{
-        perform_delay(@intCast(delay));
+        .func = particle_continuation,
-        i += delay;
+        .arg = data,
-        t += 1;
+    });
 
-        if (t == 3)
+    @panic("shouldn't happen");
-            t = 0;
-
-        plot_pixel(@intCast(data.x), @intCast(data.y), background_color);
-
-        data.act_x += @divTrunc(data.vel_x * delay, 1000);
-        data.act_y += @divTrunc(data.vel_y * delay, 1000);
-        data.x = fp_to_int(data.act_x);
-        data.y = fp_to_int(data.act_y);
-
-        data.vel_y += @divTrunc(int_to_fp(10) * delay, 1000);
-    }
-
-    ps.thread.exit();
 }
 
 fn spawn_particle(arg: ?*anyopaque) void {
-    const t = ps.thread.create_kernel(.Default, &particle, arg) catch @panic("OOM");
+    const t = ps.thread.create_kernel(
+        .Default,
+        .{ .func = &particle, .arg = arg },
+        false,
+    ) catch @panic("OOM");
+
     ke.sched.enqueue(&t.kern);
 }
 
 fn spawn_explodeable() void {
-    const t = ps.thread.create_kernel(.Default, &explodeable, null) catch @panic("OOM");
+    const t = ps.thread.create_kernel(
+        .Default,
+        .{ .func = &explodeable, .arg = null },
+        false,
+    ) catch @panic("OOM");
     ke.sched.enqueue(&t.kern);
 }
 
@@ -186,7 +219,7 @@ fn explodeable(_: ?*anyopaque) void {
 
         const delay: i32 = @as(u8, 16) + @intFromBool(t != 0);
 
-        perform_delay(@intCast(delay));
+        sleep(@intCast(delay), null);
 
         i += delay;
         t += 1;
@@ -216,6 +249,24 @@ fn explodeable(_: ?*anyopaque) void {
     ps.thread.exit();
 }
 
+fn spawning_loop(_: ?*anyopaque) void {
+    const spawn_count: usize = @intCast(@rem(rand(), 20) + 1);
+
+    for (0..spawn_count) |_| {
+        spawn_explodeable();
+    }
+
+    std.log.info("async: {}, sync: {}, usable memory: {} KiB, {} stacks for {} particles", .{
+        mm.private.tlb.async_shootdowns.load(.monotonic),
+        mm.private.tlb.sync_shootdowns.load(.monotonic),
+        mm.private.phys.usable_memory.load(.monotonic) / 1024,
+        ke.private.thread.stacks_count(),
+        particle_count.load(.monotonic),
+    });
+
+    sleep(2000, .{ .func = spawning_loop, .arg = null });
+}
+
 pub fn start(param: ?*anyopaque) void {
     const boot_info: *r.BootInfo = @ptrCast(@alignCast(param));
 
@@ -231,15 +282,5 @@ pub fn start(param: ?*anyopaque) void {
 
     fill_screen(0x09090F);
 
-    while (true) {
+    spawning_loop(null);
-        const spawn_count: usize = @intCast(@rem(rand(), 20) + 1);
-
-        for (0..spawn_count) |_| {
-            spawn_explodeable();
-        }
-
-        std.log.info("async: {}, sync: {}, usable memory: {} KiB", .{ mm.private.tlb.async_shootdowns.load(.monotonic), mm.private.tlb.sync_shootdowns.load(.monotonic), mm.private.phys.usable_memory.load(.monotonic) / 1024 });
-
-        perform_delay(2000);
-    }
 }

src/tests/mutex.zig

@@ -169,7 +169,11 @@ pub fn start(_: ?*anyopaque) void {
         is_parked[i] = .init(false);
 
         const prio: u8 = @intCast(2 + (i * 5) % 29);
-        const t = ps.thread.create_kernel(@enumFromInt(prio), &worker, @ptrFromInt(i + 1)) catch @panic("oom");
+        const t = ps.thread.create_kernel(
+            @enumFromInt(prio),
+            .{ .func = &worker, .arg = @ptrFromInt(i + 1) },
+            false,
+        ) catch @panic("oom");
         tds[i] = t;
         ke.sched.enqueue(&t.kern);
     }

src/tests/smr.zig

@@ -72,7 +72,7 @@ fn reader(param: ?*anyopaque) void {
 
         _ = reads.fetchAdd(done, .monotonic);
 
-        if (iter % 512 == 0) ke.sched.yield();
+        if (iter % 512 == 0) ke.sched.yield(null);
     }
 }
 
@@ -109,7 +109,7 @@ fn writer(param: ?*anyopaque) void {
 
         _ = writes.fetchAdd(1, .monotonic);
 
-        if (iter % 256 == 0) ke.sched.yield();
+        if (iter % 256 == 0) ke.sched.yield(null);
     }
 }
 
@@ -130,12 +130,14 @@ fn poller(_: ?*anyopaque) void {
 
         _ = polls.fetchAdd(1, .monotonic);
 
-        ke.sched.yield();
+        ke.sched.yield(null);
     }
 }
 
 fn spawn(entry: *const fn (?*anyopaque) void, id: usize) void {
-    const t = ps.thread.create_kernel(.Default, entry, @ptrFromInt(id + 1)) catch @panic("oom");
+    const t = ps.thread.create_kernel(.Default, .{.func = entry, .arg = @ptrFromInt(id + 1)}, true)
+ catch
+@panic("oom");
     ke.sched.enqueue(&t.kern);
 }
 

src/tests/smr_pi.zig

@@ -150,7 +150,7 @@ fn writer(param: ?*anyopaque) void {
 
         _ = writes.fetchAdd(1, .monotonic);
 
-        ke.sched.yield();
+        ke.sched.yield(null);
     }
 }
 
@@ -162,7 +162,7 @@ fn poller(_: ?*anyopaque) void {
 
         _ = polls.fetchAdd(1, .monotonic);
 
-        ke.sched.yield();
+        ke.sched.yield(null);
     }
 }
 
@@ -279,7 +279,7 @@ fn quiesce_check() void {
 }
 
 fn spawn(prio: u8, entry: *const fn (?*anyopaque) void, id: usize) *ps.Thread {
-    const t = ps.thread.create_kernel(@enumFromInt(prio), entry, @ptrFromInt(id + 1)) catch @panic("oom");
+    const t = ps.thread.create_kernel(@enumFromInt(prio), .{.func = entry, .arg= @ptrFromInt(id + 1)}, true) catch @panic("oom");
     ke.sched.enqueue(&t.kern);
     return t;
 }

src/tests/turnstile.zig

@@ -105,7 +105,7 @@ fn chain_high(_: ?*anyopaque) void {
 }
 
 fn spawn(prio: u8, entry: *const fn (?*anyopaque) void, arg: ?*anyopaque) *ps.Thread {
-    const t = ps.thread.create_kernel(@enumFromInt(prio), entry, arg) catch @panic("oom");
+    const t = ps.thread.create_kernel(@enumFromInt(prio), .{ .func = entry, .arg = arg }, true) catch @panic("oom");
     ke.sched.enqueue(&t.kern);
     return t;
 }
@@ -217,7 +217,7 @@ pub fn start(_: ?*anyopaque) void {
         const prio: u8 = @intCast(2 + (i * 5) % 29);
         stress_prio[i] = prio;
 
-        const t = ps.thread.create_kernel(@enumFromInt(prio), &stress, @ptrFromInt(i + 1)) catch @panic("oom");
+        const t = ps.thread.create_kernel(@enumFromInt(prio), .{ .func = &stress, .arg = @ptrFromInt(i + 1) }, true) catch @panic("oom");
         stress_tds[i] = t;
         ke.sched.enqueue(&t.kern);
     }