From 8226a205a0d8d161b48183c7112892528d532f4a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lo=C3=AFc=20Gr=C3=A9goire?= Date: Sat, 29 Aug 2026 20:03:22 -0400 Subject: [PATCH] ke: add continuations --- src/ex/console.zig | 4 +- src/ex/init.zig | 3 +- src/ex/work.zig | 22 ++-- src/kern/amd64/impl.zig | 118 +++++++++++++---- src/kern/amd64/locore.s | 59 ++++++++- src/kern/init.zig | 24 ++-- src/kern/root.zig | 5 + src/kern/sched.zig | 175 +++++++++++++++++++------ src/kern/thread.zig | 277 ++++++++++++++++++++++++++++++++++++---- src/kern/wait.zig | 7 +- src/loader/jump.s | 1 + src/loader/main.zig | 3 +- src/platform/pc/smp.zig | 20 ++- src/proc/root.zig | 1 + src/proc/thread.zig | 55 ++++++-- src/tests/fireworks.zig | 129 ++++++++++++------- src/tests/mutex.zig | 6 +- src/tests/smr.zig | 10 +- src/tests/smr_pi.zig | 6 +- src/tests/turnstile.zig | 4 +- 20 files changed, 735 insertions(+), 194 deletions(-) diff --git a/src/ex/console.zig b/src/ex/console.zig index b513c79..91b1480 100644 --- a/src/ex/console.zig +++ b/src/ex/console.zig @@ -126,8 +126,8 @@ pub fn init() void { const td = ps.thread.create_kernel( .Default, - worker, - null, + .{ .func = worker, .arg = null }, + true, ) catch @panic("handle me"); ke.sched.enqueue(&td.kern); } diff --git a/src/ex/init.zig b/src/ex/init.zig index 5683f0b..7f4a959 100644 --- a/src/ex/init.zig +++ b/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}); diff --git a/src/ex/work.zig b/src/ex/work.zig index 6c5b86c..719ccbc 100644 --- a/src/ex/work.zig +++ b/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, - @ptrFromInt(ctx.value), + .{ .func = work_loop, .arg = @ptrFromInt(ctx.value) }, + true, ); td.kern.nice = -ke.Priority.nice_max; - if (cpu) |c| { - ke.sched.pin_on(&td.kern, c); + if (cpu != null) { + 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, - @ptrFromInt(ctx.value), + .{ .func = work_loop, .arg = @ptrFromInt(ctx.value) }, + true, ); - if (cpu) |c| { - ke.sched.pin_on(&td.kern, c); + if (cpu != null) { + 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, - null, + .{ .func = pool_manager, .arg = null }, + true, ) catch @panic("Failed to create pool manager"); ke.sched.enqueue(&td.kern); diff --git a/src/kern/amd64/impl.zig b/src/kern/amd64/impl.zig index a477ad7..302d3b4 100644 --- a/src/kern/amd64/impl.zig +++ b/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( - old: *ThreadContext, - new: *ThreadContext, +extern fn amd64_context_load(td: *ThreadContext) callconv(.c) void; + +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), - frame: ThreadFrame, + rsp: u64 align(1) = 0xDEADC0D3, 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 + stack_size; + var sp: usize = stack_top; sp = sp & ~@as(usize, 15); sp = sp - @sizeOf(ThreadFrame); @@ -60,32 +99,57 @@ 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( - self: *ThreadContext, - new: *ThreadContext, - ) callconv(.c) void { - const thread: *ke.Thread = @alignCast(@fieldParentPtr("context", self)); - do_context_switch( - self, - new, - &thread.lock.inner.locked.raw, - &thread.switching.raw, - ); - } - - pub fn load(self: *ThreadContext) callconv(.c) void { - do_context_load(self); + pub fn load(self: *ThreadContext) void { + amd64_context_load(self); } }; +pub fn switch_normal_to_normal( + 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 switch_cont_to_normal( + 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]); } diff --git a/src/kern/amd64/locore.s b/src/kern/amd64/locore.s index 5e75ade..0482263 100644 --- a/src/kern/amd64/locore.s +++ b/src/kern/amd64/locore.s @@ -133,8 +133,8 @@ __interrupt_vectors: .endr -.global do_context_switch -do_context_switch: +.global amd64_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 -do_context_load: +.global amd64_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 diff --git a/src/kern/init.zig b/src/kern/init.zig index 59a9a7d..aed14c0 100644 --- a/src/kern/init.zig +++ b/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 - // into the thread. - thread0.init( - boot_info.kernel_stack, - boot_info.kernel_stack_size, - .IdleThread, - init, - null, - ); + // our idle thread. At this point we're already running on the kernel stack. + thread0.init(.{ + .stack = boot_info.kernel_stack + boot_info.kernel_stack_size, + .priority = .IdleThread, + .entry = undefined, + // Shouldn't block + .turnstile = undefined, + }); + + thread0.continuation = null; thread0.pinned = true; - // Jump into the thread. - thread0.context.load(); + init(); } -pub fn init(_: ?*anyopaque) linksection(r.init) void { +fn init() linksection(r.init) void { ke.ncpus = 1; ki.impl.early_init(); ki.tunable.init(); diff --git a/src/kern/root.zig b/src/kern/root.zig index baad906..0aa4dde 100644 --- a/src/kern/root.zig +++ b/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 { diff --git a/src/kern/sched.zig b/src/kern/sched.zig index 97c27e8..c7f757d 100644 --- a/src/kern/sched.zig +++ b/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. - cur.context.switch_to(&next.context); + if (cur.continuation != null) { + 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 { - cpu.queues_lock.acquire_no_ipl(); - const cur = cpu.current_thread.?; + var next: *ke.Thread = undefined; - const next = cpu.next_thread orelse { + while (true) { + cpu.queues_lock.acquire_no_ipl(); + + next = cpu.next_thread orelse select_thread(cur, cpu, null) orelse { + cpu.queues_lock.release_no_ipl(); + return; + }; + + cpu.next_thread = null; cpu.queues_lock.release_no_ipl(); - return; - }; - cpu.next_thread = null; + 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; - if (next != null) { - sched_cpu.next_thread = null; - } else { - // Pick a new thread to run. - next = select_thread(null, sched_cpu, null); - } + cur.continuation = continuation; - if (next == null) { - // Nothing to run, go idle. - next = sched_cpu.idle_thread; - sched_cpu.steal_work = true; - } + 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); + } - if (next) |n| { - sched_cpu.current_thread_prio.store(n.priority, .monotonic); - } + sched_cpu.queues_lock.release_no_ipl(); - 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 != 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. - cur.?.state.store(.Running, .monotonic); - cur.?.switching.store(false, .monotonic); - cur.?.lock.release_no_ipl(); + // Same thread. + const cont = cur.continuation; + cur.continuation = null; + + 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; } diff --git a/src/kern/thread.zig b/src/kern/thread.zig index 1b96dcf..565debf 100644 --- a/src/kern/thread.zig +++ b/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(); +} diff --git a/src/kern/wait.zig b/src/kern/wait.zig index 6c38c30..dde6134 100644 --- a/src/kern/wait.zig +++ b/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; diff --git a/src/loader/jump.s b/src/loader/jump.s index ac47695..747fc0c 100644 --- a/src/loader/jump.s +++ b/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 */ diff --git a/src/loader/main.zig b/src/loader/main.zig index 467d860..a956698 100644 --- a/src/loader/main.zig +++ b/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, ); } diff --git a/src/platform/pc/smp.zig b/src/platform/pc/smp.zig index 72dd0e1..93ff072 100644 --- a/src/platform/pc/smp.zig +++ b/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_size, - stack_size, + stack_top, + @intCast(cpu_id), ); rtl.barrier.fence(.release); diff --git a/src/proc/root.zig b/src/proc/root.zig index 99532e0..a3d1e72 100644 --- a/src/proc/root.zig +++ b/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; diff --git a/src/proc/thread.zig b/src/proc/thread.zig index f1dfc51..f24a586 100644 --- a/src/proc/thread.zig +++ b/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, - arg: ?*anyopaque, + continuation: ke.Continuation, + 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); - td.kern.turnstile = try psp.turnstile_zone.create(); + const stack = if (start_with_stack) + @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); +} diff --git a/src/tests/fireworks.zig b/src/tests/fireworks.zig index bfcdd95..02eea00 100644 --- a/src/tests/fireworks.zig +++ b/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 { - var timer: ke.Timer = undefined; - timer.init(); - ke.timer.set(&timer, .init(std.time.ns_per_ms * ms), .{}); - _ = ke.wait.wait_one(&timer.hdr, "sleep", .{}) catch unreachable; +fn sleep(ms: usize, continuation: ?ke.Continuation) void { + _ = ke.wait.wait_any(&.{}, "sleep", .{ + .timeout = .from(r.Milliseconds.init(ms)), + .continuation = continuation, + }) 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; + plot_pixel(@intCast(data.x), @intCast(data.y), data.color); - while (i < expire_in) { - plot_pixel(@intCast(data.x), @intCast(data.y), data.color); + sleep(particle_delay, .{ + .func = particle_continuation, + .arg = data, + }); - const delay = @as(i32, 16) + @intFromBool(t != 0); - perform_delay(@intCast(delay)); - i += delay; - t += 1; - - if (t == 3) - 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(); + @panic("shouldn't happen"); } 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) { - 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); - } + spawning_loop(null); } diff --git a/src/tests/mutex.zig b/src/tests/mutex.zig index 80dfb69..b8471ee 100644 --- a/src/tests/mutex.zig +++ b/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); } diff --git a/src/tests/smr.zig b/src/tests/smr.zig index 5e99701..c220b27 100644 --- a/src/tests/smr.zig +++ b/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); } diff --git a/src/tests/smr_pi.zig b/src/tests/smr_pi.zig index a991022..1564d42 100644 --- a/src/tests/smr_pi.zig +++ b/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; } diff --git a/src/tests/turnstile.zig b/src/tests/turnstile.zig index 383c733..1fd640a 100644 --- a/src/tests/turnstile.zig +++ b/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); }