From 878be31c7a11a2a79a02ddf597f77f2dd8818614 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lo=C3=AFc=20Gr=C3=A9goire?= <loic@gregoire.dev>
Date: Sun, 19 Jul 2026 20:40:50 -0400
Subject: [PATCH] ke/sched: fix lock inversion in scheduler

Taking next's lock can be a lock inversion with the turnstile code.
Use a switching field instead and spin on that.
---
 src/kern/amd64/impl.zig |  4 ++--
 src/kern/amd64/locore.s |  8 ++++++--
 src/kern/sched.zig      | 16 +++++++++++-----
 src/kern/thread.zig     |  5 +++++
 4 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/src/kern/amd64/impl.zig b/src/kern/amd64/impl.zig
index db3e35b..2656131 100644
--- a/src/kern/amd64/impl.zig
+++ b/src/kern/amd64/impl.zig
@@ -25,7 +25,7 @@ const ThreadFrame = extern struct {
 
 extern fn asm_thread_entry() void;
 
-extern fn do_context_switch(old: *ThreadContext, new: *ThreadContext, lock: *u8) callconv(.c) void;
+extern fn do_context_switch(old: *ThreadContext, new: *ThreadContext, lock: *u8, switching: *bool) callconv(.c) void;
 
 pub const ThreadContext = extern struct {
     rsp: u64 align(1),
@@ -60,7 +60,7 @@ pub const ThreadContext = extern struct {
 
     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);
+        do_context_switch(self, new, &thread.lock.inner.locked.raw, &thread.switching.raw);
     }
 };
 
diff --git a/src/kern/amd64/locore.s b/src/kern/amd64/locore.s
index 249a46a..a9181c4 100644
--- a/src/kern/amd64/locore.s
+++ b/src/kern/amd64/locore.s
@@ -149,9 +149,13 @@ do_context_switch:
     /* Load new stack pointer. */
     mov (%rsi), %rsp
 
+	/* Set switching to false. */
+    xorl %eax, %eax
+    xchgb %al, (%rcx)
+
     /* Unlock passed lock. */
-    xorl %ecx, %ecx
-    xchgb %cl, (%rdx)
+    xorl %eax, %eax
+    xchgb %al, (%rdx)
 
     /* Restore callee-saved registers. */
     pop %rbp
diff --git a/src/kern/sched.zig b/src/kern/sched.zig
index f5b3464..17e0423 100644
--- a/src/kern/sched.zig
+++ b/src/kern/sched.zig
@@ -355,6 +355,8 @@ pub fn handle_preemption(cpu: *PerCpu) void {
     cpu.queues_lock.release_no_ipl();
     cur.lock.acquire_no_ipl();
 
+    cur.switching.store(true, .monotonic);
+
     if (cur != cpu.idle_thread and cur.state.load(.monotonic) != .Blocked) {
         // Put it back in the queue.
         cpu.queues_lock.acquire_no_ipl();
@@ -363,8 +365,9 @@ pub fn handle_preemption(cpu: *PerCpu) void {
         cpu.queues_lock.release_no_ipl();
     }
 
-    next.lock.acquire_no_ipl();
-    next.lock.release_no_ipl();
+    while (next.switching.load(.monotonic) == true) {
+        std.atomic.spinLoopHint();
+    }
 
     do_switch(cpu, cur, next);
 
@@ -435,6 +438,7 @@ pub fn yield() void {
     const td = percpu.local().current_thread.?;
 
     td.lock.acquire_no_ipl();
+    td.switching.store(true, .monotonic);
 
     const cpu = percpu.local();
     if (td != cpu.idle_thread and td.state.load(.monotonic) == .Running) {
@@ -512,14 +516,16 @@ pub fn yield_locked() void {
     sched_cpu.queues_lock.release_no_ipl();
 
     if (next != null and cur.? != next.?) {
-        // Switch into the thread. Take its lock here to ensure it is not still on its stack.
-        next.?.lock.acquire_no_ipl();
-        next.?.lock.release_no_ipl();
+        // Switch into the thread. Ensure it's not switching.
+        while (next.?.switching.load(.monotonic) == true) {
+            std.atomic.spinLoopHint();
+        }
 
         do_switch(sched_cpu, cur.?, next.?);
     } else {
         // Ensure curthread is not marked as selected.
         cur.?.state.store(.Running, .monotonic);
+        cur.?.switching.store(false, .monotonic);
         cur.?.lock.release_no_ipl();
     }
 
diff --git a/src/kern/thread.zig b/src/kern/thread.zig
index bcf96ed..8a11fd3 100644
--- a/src/kern/thread.zig
+++ b/src/kern/thread.zig
@@ -114,6 +114,10 @@ pub const Thread = struct {
     stack: r.VAddr,
     /// PELT load average,
     avg: ki.sched.Average,
+    /// Set whenever the thread is switching off its stack.
+    /// This is used to avoid taking thread next lock to wait for switch off
+    /// to complete.
+    switching: std.atomic.Value(bool),
 
     /// Initialize a thread.
     /// - `stack`: Address of the **base** of the stack on which the initial context for the thread is built
@@ -149,6 +153,7 @@ pub const Thread = struct {
             .queue = null,
             .queue_item = null,
             .stack = stack,
+            .switching = .init(false),
             .avg = .{
                 .last_update = 0,
                 // New threads are considered heavy until they prove themselves,