ke/turnstile: fix a bunch of bugs

src/kern/sched.zig

@@ -164,6 +164,7 @@ pub fn handle_preemption(cpu: *PerCpu) void {
     };
 
     cpu.next_thread = null;
+    cpu.current_thread_prio.store(next.priority, .monotonic);
 
     cpu.queues_lock.release_no_ipl();
     cur.lock.acquire_no_ipl();
@@ -295,6 +296,10 @@ pub fn yield_locked() void {
         sched_cpu.steal_work = true;
     }
 
+    // Publish the incoming thread's priority while the lock still
+    // orders us against priority writers.
+    sched_cpu.current_thread_prio.store(next.?.priority, .monotonic);
+
     sched_cpu.queues_lock.release_no_ipl();
 
     if (next != null and cur.? != next.?) {
@@ -340,20 +345,26 @@ pub fn update_priority_locked(td: *ke.Thread, new_prio: u8) void {
     std.debug.assert(td.lock.is_locked());
 
     const old_prio = td.priority;
-    td.priority = new_prio;
 
-    const c = td.cpu orelse return;
+    const c = td.cpu orelse {
+        td.priority = new_prio;
+        return;
+    };
+
     const cpu = percpu.remote(c);
 
     cpu.queues_lock.acquire_no_ipl();
     defer cpu.queues_lock.release_no_ipl();
 
+    td.priority = new_prio;
+
     const state = td.state.load(.monotonic);
 
     if (state == .Ready) {
         // Remove it from its queue and add it back.
         remove_from_queue(cpu, td);
         insert_in_queue(cpu, td, false);
+        td.cpu = c;
     }
 
     if (state == .Running) {
@@ -380,8 +391,14 @@ pub fn update_priority_locked(td: *ke.Thread, new_prio: u8) void {
         }
     }
 
-    // XXX: handle selected threads?
+    if (state == .Selected or cpu.next_thread == td) {
+        // Thread is committed to run on `c` but not running yet,
+        // update the hint.
+        cpu.current_thread_prio.store(new_prio, .monotonic);
+        return;
+    }
 
+    // XXX: Handle Selected threads so that they are placed properly?
 }
 
 /// Initialize a CPU for use by the scheduler.
@@ -783,7 +800,7 @@ fn enqueue_on_cpu(c: u32, td: *ke.Thread) void {
 
     defer cpu.queues_lock.release_no_ipl();
 
-    if (cpu.current_thread == null or should_preempt(td, cpu.current_thread.?)) {
+    if (should_prio_preempt(td, cpu.current_thread_prio.load(.monotonic))) {
         // We can preempt the current thread, first check whether or not there
         // is already a next thread selected, and if we can preempt it.
 
@@ -807,6 +824,7 @@ fn enqueue_on_cpu(c: u32, td: *ke.Thread) void {
             cpu.next_thread = td;
             cpu.preemption_reason = .HigherPriority;
             td.cpu = c;
+            td.state.store(.Selected, .monotonic);
 
             ki.ipl.set_softint_pending(c, .Dispatch);
 
@@ -909,9 +927,6 @@ fn do_switch(cpu: *PerCpu, cur: *ke.Thread, next: *ke.Thread) void {
     cpu.current_thread = next;
     next.last_cpu = ke.cpu.current();
 
-    // Stash thread status.
-    cpu.current_thread_prio.store(next.priority, .monotonic);
-
     // Do machine-dependent switch.
     cur.context.switch_to(&next.context);
 }

src/kern/turnstile.zig

@@ -61,7 +61,9 @@ fn highest_waiter(ts: *Turnstile) ?u8 {
         var it = q.iterator();
         while (it.next()) : (it.advance()) {
             const w: *Waiter = @fieldParentPtr("link", it.get());
+            w.thread.lock.acquire_no_ipl();
             best = @max(best orelse 0, w.thread.priority);
+            w.thread.lock.release_no_ipl();
         }
     }
     return best;
@@ -151,10 +153,6 @@ fn propagate(curtd: *ke.Thread, root: *Chain) void {
         const obj = thread.waiting_on orelse break;
         const bucket = &chains[hash_obj(obj)];
 
-        // Lock `bucket` and confirm `thread` didn't move before we got it.
-        // If the trylock fails, drop root and restart from curtd.
-        // If validation fails, restart while still holding root.
-        const restart = blk: {
         if (bucket != root) {
             if (!bucket.lock.try_acquire_no_ipl()) {
                 std.debug.assert(held == null);
@@ -174,21 +172,7 @@ fn propagate(curtd: *ke.Thread, root: *Chain) void {
             }
             held = bucket;
         }
-            break :blk thread.waiting_on != obj;
+        std.debug.assert(thread.waiting_on == obj);
-        };
-
-        // Try again.
-        if (restart) {
-            if (held) |h| h.lock.release_no_ipl();
-            held = null;
-            if (thread != curtd) {
-                thread.lock.release_no_ipl();
-                curtd.lock.acquire_no_ipl();
-                thread = curtd;
-            }
-            donate = curtd.priority;
-            continue;
-        }
 
         const ts = thread.turnstile;
         const owner = ts.owner orelse break;
@@ -200,6 +184,10 @@ fn propagate(curtd: *ke.Thread, root: *Chain) void {
             break;
         }
 
+        // The next hop blocks owner, so lend the owner's priority
+        // down the chain.
+        donate = owner.priority;
+
         // Keep the owner locked, drop everything below it.
         thread.lock.release_no_ipl();
         if (held) |h| h.lock.release_no_ipl();
@@ -265,6 +253,15 @@ pub fn block(turnstile: ?*Turnstile, obj: *anyopaque, owner: *ke.Thread, queue:
         curtd.turnstile.next_free = turn.next_free;
         turn.next_free = curtd.turnstile;
         curtd.turnstile = turn;
+
+        if (turn.owner) |cur| {
+            std.debug.assert(cur == owner);
+        } else {
+            // Owner might be null because of a partial wakeup
+            // where we cleared it, restore it here.
+            turn.owner = owner;
+            reboost(turn, owner);
+        }
     } else {
         // This is the first thread to block on this object.
         // Lend its turnstile and add it to the hash chain.
@@ -357,6 +354,9 @@ fn dequeue(ts: *Turnstile, td: *ke.Thread) void {
     std.debug.assert(td.turnstile_waiter != null);
     std.debug.assert(!ts.queues[0].is_empty() or !ts.queues[1].is_empty());
 
+    const waiter = td.turnstile_waiter.?;
+    waiter.link.remove();
+
     if (ts.next_free) |free| {
         // Steal a turnstile from the freelist.
         td.turnstile = free;
@@ -372,12 +372,12 @@ fn dequeue(ts: *Turnstile, td: *ke.Thread) void {
         ts.donated = null;
     }
 
-    const waiter = td.turnstile_waiter.?;
-    waiter.link.remove();
     td.turnstile_waiter = null;
     ts.waiters -= 1;
 
+    td.lock.acquire_no_ipl();
     td.waiting_on = null;
+    td.lock.release_no_ipl();
 
     waiter.event.signal();
 }