ke: make getting the current thread atomic

src/sys/kern/amd64/impl.zig

@@ -246,3 +246,10 @@ pub inline fn lower_software_ipl(new: ke.Ipl, old: *ke.Ipl) bool {
           [old_ptr] "r" (old),
         : .{ .rcx = true, .memory = true });
 }
+
+pub inline fn current_thread() ?*ke.Thread {
+    return @ptrFromInt(asm volatile (
+        \\mov %%gs:current_thread, %[ret]
+        : [ret] "=r" (-> u64),
+    ));
+}

src/sys/kern/amd64/init.zig

@@ -95,8 +95,8 @@ pub fn ap_entry(cpu_id: u32, booted: *std.atomic.Value(usize)) noreturn {
 
     const sched = kep.sched.percpu.local();
 
-    sched.current_thread = pl.impl.smp.start_thread.local().*;
-    sched.idle_thread = sched.current_thread;
+    kep.sched.current_thread.local().* = pl.impl.smp.start_thread.local().*;
+    sched.idle_thread = kep.thread.current();
 
     amd64.sti();
 

src/sys/kern/dpc.zig

@@ -141,7 +141,7 @@ fn dispatch_queue(cpu: u32) void {
     kep.shootdown.process_shootdowns();
 
     if (sched_cpu.next_thread != null) {
-        const curtd = sched_cpu.current_thread orelse unreachable;
+        const curtd = ke.thread.current();
 
         if (!curtd.smr_sections.is_empty()) {
             kep.smr.mark_thread_stalled(curtd);

src/sys/kern/init.zig

@@ -54,7 +54,7 @@ fn init() linksection(r.init) void {
     kep.cpu.init_cpu(0);
     kep.turnstile.init_turnstiles();
 
-    kep.sched.percpu.local().current_thread = &thread0;
+    kep.sched.current_thread.local().* = &thread0;
     kep.sched.percpu.local().idle_thread = &thread0;
 
     std.log.info("Zag for {s} ({s}), cmdline is \"{?s}\"", .{

src/sys/kern/mutex.zig

@@ -32,9 +32,7 @@ pub const Mutex = struct {
     }
 
     pub fn acquire(m: *Mutex) void {
-        const ipl = ke.ipl.raise(.Dispatch);
-        const curtd = kep.sched.percpu.local().current_thread orelse
-            unreachable;
+        const curtd = ke.thread.current();
 
         if (config.warden) {
             kep.warden.check(m.wd, .Lock);
@@ -54,7 +52,6 @@ pub const Mutex = struct {
         );
 
         if (cur_owner == null) {
-            ke.ipl.lower(ipl);
             return;
         }
 
@@ -74,10 +71,11 @@ pub const Mutex = struct {
                 continue;
             }
 
-            ke.ipl.lower(ipl);
             return;
         }
 
+        const ipl = ke.ipl.raise(.Dispatch);
+
         // Slow path: contended, block on a turnstile.
         while (true) {
             const ts = kep.turnstile.lookup(m);

src/sys/kern/private.zig

@@ -74,6 +74,13 @@ const ImplSchema = struct {
         _ = old;
         unreachable;
     }
+
+    /// Return a pointer to the current thread.
+    /// This must be atomic with regards to migration (i.e this should be a
+    /// consistent pointer no matter what).
+    pub fn current_thread() ?*ke.Thread {
+        unreachable;
+    }
 };
 
 comptime {

src/sys/kern/queue.zig

@@ -60,7 +60,7 @@ pub const Queue = struct {
     /// This blocks until an item is actually popped.
     pub fn remove(self: *Queue, timeout: ?rtl.Duration) !*rtl.List.Entry {
         const ipl = self.hdr.lock.acquire();
-        const td = kep.sched.percpu.local().current_thread orelse unreachable;
+        const td = ke.thread.current();
 
         if (td.queue) |q| {
             assert(q == self);

src/sys/kern/sched.zig

@@ -208,8 +208,6 @@ pub const PerCpu = struct {
     start_timer: bool,
     /// The reason why the current thread was preempted.
     preemption_reason: PreemptionReason,
-    /// Currently running thread.
-    current_thread: ?*ke.Thread,
     /// Per-CPU idle thread.
     idle_thread: ?*ke.Thread,
     /// Thread selected for preemption.
@@ -242,6 +240,11 @@ pub const Accounting = struct {
 };
 
 pub const percpu = ke.CpuLocalType(PerCpu, undefined);
+pub const current_thread = ke.ExportedCpuLocalType(
+    ?*ke.Thread,
+    null,
+    "current_thread",
+);
 
 var balance_timer: ke.Timer = undefined;
 var balance_dpc = ke.Dpc.init(balance);
@@ -276,7 +279,6 @@ fn init_cpu() linksection(r.init) void {
         .resched_timer = undefined,
         .start_timer = false,
         .preemption_reason = .None,
-        .current_thread = null,
         .idle_thread = null,
         .next_thread = null,
         .pick_offset = 0,
@@ -1022,7 +1024,7 @@ fn enqueue_on_cpu(c: u32, td: *ke.Thread) void {
 
     attach_load_avg(cpu, td);
 
-    if (cpu.current_thread != null and
+    if (current_thread.local().* != null and
         should_preempt_cpu(td, cpu, c != ke.cpu.current()))
     {
         // We can preempt the current thread, first check whether or not there
@@ -1088,8 +1090,8 @@ fn allocate_stack(td: *ke.Thread) ?usize {
 }
 
 // Execute a context switch on `cpu`.
-fn do_switch(cpu: *PerCpu, cur: *ke.Thread, next: *ke.Thread) void {
-    cpu.current_thread = next;
+fn do_switch(cur: *ke.Thread, next: *ke.Thread) void {
+    current_thread.local().* = next;
     next.last_cpu = ke.cpu.current();
 
     const time = ke.time.read_time();
@@ -1153,7 +1155,7 @@ fn do_switch(cpu: *PerCpu, cur: *ke.Thread, next: *ke.Thread) void {
 /// 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 orelse unreachable;
+    const cur = ke.thread.current();
     var next: *ke.Thread = undefined;
 
     while (true) {
@@ -1220,7 +1222,7 @@ pub fn handle_preemption(cpu: *PerCpu) void {
         std.atomic.spinLoopHint();
     }
 
-    do_switch(cpu, cur, next);
+    do_switch(cur, next);
 
     // cur.lock dropped
 }
@@ -1230,7 +1232,7 @@ pub fn clock(_: *ke.Dpc, _: ?*anyopaque) void {
     const cpu = percpu.local();
     assert(kep.ipl.current() == .Dispatch);
 
-    const curtd = cpu.current_thread orelse return;
+    const curtd = current_thread.local().* orelse return;
 
     const now = ke.time.read_time();
     const now_us = now.to_us();
@@ -1292,7 +1294,7 @@ pub fn enqueue(td: *ke.Thread) void {
 /// Yield on the current CPU.
 pub fn yield(continuation: ?ke.Continuation) void {
     const ipl = ke.ipl.raise(.Dispatch);
-    const td = percpu.local().current_thread orelse unreachable;
+    const td = ke.thread.current();
 
     td.lock.acquire_no_ipl();
     td.switching.store(true, .monotonic);
@@ -1317,7 +1319,7 @@ pub fn yield_locked(continuation: ?ke.Continuation) void {
 
     sched_cpu.queues_lock.acquire_no_ipl();
 
-    const cur = sched_cpu.current_thread orelse unreachable;
+    const cur = ke.thread.current();
     var next = sched_cpu.next_thread;
 
     cur.continuation = continuation;
@@ -1386,7 +1388,7 @@ pub fn yield_locked(continuation: ?ke.Continuation) void {
         while (n.switching.load(.monotonic) == true) {
             std.atomic.spinLoopHint();
         }
-        do_switch(sched_cpu, cur, n);
+        do_switch(cur, n);
     } else {
         // Same thread.
         cur.state.store(.Running, .monotonic);

src/sys/kern/thread.zig

@@ -316,7 +316,7 @@ pub fn exit() void {
 
 /// Return the currently running thread.
 pub fn current() *ke.Thread {
-    return kep.sched.percpu.local().current_thread orelse unreachable;
+    return kep.impl.current_thread() orelse unreachable;
 }
 
 // -- Stack allocation and continuations ---------------------------------------

src/sys/kern/turnstile.zig

@@ -443,7 +443,7 @@ pub fn block(
     queue: Queue,
 ) void {
     const chain = chain_for(obj);
-    const curtd = kep.sched.percpu.local().current_thread.?;
+    const curtd = ke.thread.current();
     const ipl = ke.ipl.current();
 
     var ts = turnstile;