fix(ke/sched): fix various bugs

src/kern/amd64/impl.zig

@@ -40,11 +40,11 @@ pub const ThreadContext = extern struct {
         var ctx: @This() = undefined;
         var sp: usize = stack + stack_size;
 
-        if (sp % 16 != 0) {
+        sp = sp & ~@as(usize, 15);
-            sp -= (sp % 16);
+        sp = sp - @sizeOf(ThreadFrame);
-        }
+        sp = (sp & ~@as(usize, 15)) - 8;
 
-        const frame: *ThreadFrame = @ptrFromInt(sp - @sizeOf(ThreadFrame));
+        const frame: *ThreadFrame = @ptrFromInt(sp);
 
         frame.* = std.mem.zeroes(ThreadFrame);
 

src/kern/sched.zig

@@ -185,6 +185,9 @@ pub fn handle_preemption(cpu: *PerCpu) void {
         cpu.queues_lock.release_no_ipl();
     }
 
+    next.lock.acquire_no_ipl();
+    next.lock.release_no_ipl();
+
     do_switch(cpu, cur, next);
 
     // cur.lock dropped
@@ -244,21 +247,27 @@ pub fn block() void {
     const td = percpu.local().current_thread.?;
 
     td.lock.acquire_no_ipl();
-
+    block_locked(td);
-    td.state = .Blocked;
-    td.sleep_start = ke.time.read_time();
-    td.runq = null;
-
-    yield(ke.cpu.current());
 
     // td lock released
     ke.ipl.lower(ipl);
 }
 
+/// Block the currently running thread with its lock held.
+pub fn block_locked(curtd: *ke.Thread) void {
+    curtd.state = .Blocked;
+    curtd.sleep_start = ke.time.read_time();
+    curtd.runq = null;
+
+    yield(ke.cpu.current());
+}
+
 /// Unblock a thread.
 pub fn unblock(td: *ke.Thread) void {
     const ipl = td.lock.acquire();
 
+    std.debug.assert(td.state == .Blocked);
+
     td.sleep_time = (ke.time.read_time() - td.sleep_start) / std.time.ns_per_ms;
 
     if (td.sleep_time >= config.sched_timeslice and td.priority_class == .Batch) {
@@ -356,10 +365,10 @@ fn pick_realtime_thread(cpu: *PerCpu, minprio: u8, migrate: bool) ?*ke.Thread {
             while (td.pinned) {
                 const entry = td.runq_link.next;
 
-                td = @fieldParentPtr("runq_link", entry);
-
                 // Every thread is pinned, go through another queue
                 if (entry == &curr_queue.head) continue :qloop;
+
+                td = @fieldParentPtr("runq_link", entry);
             }
         }
 
@@ -410,10 +419,10 @@ fn pick_batch_thread(cpu: *PerCpu, migrate: bool) ?*ke.Thread {
             while (td.pinned) {
                 const entry = td.runq_link.next;
 
-                td = @fieldParentPtr("runq_link", entry);
-
                 // Every thread is pinned, go through another queue
                 if (entry == &curr_queue.head) continue :qloop;
+
+                td = @fieldParentPtr("runq_link", entry);
             }
         }
 
@@ -442,10 +451,12 @@ fn pick_idle_thread(cpu: *PerCpu, migrate: bool) ?*ke.Thread {
     // Just get the first thread from the queue.
     const td: *ke.Thread = @fieldParentPtr("runq_link", cpu.idle_queue.first());
 
-    if (!(migrate and td.pinned)) {
+    if (migrate and td.pinned) {
+        return null;
+    }
+
     // We can safely remove it.
     td.runq_link.remove();
-    }
 
     return td;
 }
@@ -655,6 +666,8 @@ fn should_preempt_cpu(td: *ke.Thread, cpu: *PerCpu) bool {
 
 // Enqueue a thread on a CPU.
 fn enqueue_on_cpu(c: u32, td: *ke.Thread) void {
+    std.debug.assert(td.lock.is_locked());
+
     const cpu = percpu.remote(c);
     cpu.queues_lock.acquire_no_ipl();
 
@@ -813,8 +826,15 @@ fn yield(cpu: u32) void {
 
     sched_cpu.queues_lock.release_no_ipl();
 
-    // Switch into the thread.
+    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();
+
         do_switch(sched_cpu, cur.?, next.?);
+    } else {
+        cur.?.lock.release_no_ipl();
+    }
 
     // cur lock dropped
 }
@@ -934,7 +954,9 @@ pub fn idle(_: ?*anyopaque) noreturn {
             // Steal a thread.
             if (steal_work(cpu_id)) |td| {
                 // Enqueue it on this CPU.
+                td.lock.acquire_no_ipl();
                 enqueue_on_cpu(cpu_id, td);
+                td.lock.release_no_ipl();
             }
 
             ke.ipl.lower(ipl);

src/kern/um/impl.zig

@@ -61,8 +61,13 @@ pub fn send_resched_ipi(cpu: u32) void {
 
 pub const ThreadContext = struct {
     ucontext: c.ucontext_t,
+    old_lock: ?*ke.SpinLock,
 
     fn real_entry(entry_int: c_ulong, arg_int: c_ulong) callconv(.c) void {
+        if (ki.sched.percpu.local().current_thread.?.context.old_lock) |lock| {
+            lock.release_no_ipl();
+            ki.sched.percpu.local().current_thread.?.context.old_lock = null;
+        }
         const entry: *const fn (?*anyopaque) void = @ptrFromInt(entry_int);
         const arg: ?*anyopaque = @ptrFromInt(arg_int);
         ke.ipl.lower(.Passive);
@@ -70,7 +75,7 @@ pub const ThreadContext = struct {
     }
 
     pub fn init(stack: r.VAddr, stack_size: usize, entry: *const fn (?*anyopaque) void, arg: ?*anyopaque) ThreadContext {
-        var new: ThreadContext = .{ .ucontext = undefined };
+        var new: ThreadContext = .{ .ucontext = undefined, .old_lock = null };
 
         if (c.getcontext(&new.ucontext) == -1) {
             @panic("getcontext() failed");
@@ -101,12 +106,16 @@ pub const ThreadContext = struct {
 
     pub fn switch_to(self: *ThreadContext, new: *ThreadContext) void {
         const thread: *ke.Thread = @fieldParentPtr("context", self);
-
+        new.old_lock = &thread.lock;
-        thread.lock.release_no_ipl();
 
         if (c.swapcontext(&self.ucontext, &new.ucontext) == -1) {
             @panic("swapcontext() failed");
         }
+
+        if (self.old_lock) |lock| {
+            self.old_lock = null;
+            lock.release_no_ipl();
+        }
     }
 };
 

src/kern/wait.zig

@@ -132,18 +132,20 @@ pub fn wait_any(objects: []*DispatchHeader, timeout: ?r.Nanoseconds, waitblocks:
     if (satisfier != null or (has_timeout and timeout.? == 0)) {
         std.debug.assert(curtd.wait_status.load(.acquire) == .Satisfied);
 
-        if (satisfier) |sat| {
+        const cleanup_count = satisfier orelse total_count;
-            // Remove any wait block we might've installed
+
-            for (0..sat) |i| {
+        // Remove any wait block we might've installed.
+        for (0..cleanup_count) |i| {
             const is_timer = has_timeout and i == timer_i;
             var obj = if (is_timer) &timer.hdr else objects[i];
             var wb = if (is_timer) timer_block else &blocks[i];
 
             obj.lock.acquire_no_ipl();
+            if (wb.status == .Active) {
                 wb.status = .Inactive;
                 wb.link.remove();
-                obj.lock.release_no_ipl();
             }
+            obj.lock.release_no_ipl();
         }
 
         return satisfier orelse error.Timeout;
@@ -153,12 +155,18 @@ pub fn wait_any(objects: []*DispatchHeader, timeout: ?r.Nanoseconds, waitblocks:
         ke.timer.set(timer, timeout.?, null);
     }
 
+    curtd.lock.acquire_no_ipl();
+
     // Now try committing the wait.
     // While we're trying to commit the wait, the object locks have been released, and the state could therefore change.
     // We need to re-check the state of the wait before actually blocking.
     if (curtd.wait_status.cmpxchgStrong(.InProgress, .Committed, .acq_rel, .monotonic) == null) {
         // We're good, now actually block.
-        ke.sched.block();
+        //std.log.debug("Blocking {*} on {*}", .{ curtd, objects[0] });
+        ki.sched.block_locked(curtd);
+    } else {
+        // Could not commit.
+        curtd.lock.release_no_ipl();
     }
 
     // We're back!
@@ -193,6 +201,8 @@ pub fn wait_any(objects: []*DispatchHeader, timeout: ?r.Nanoseconds, waitblocks:
 
 /// Satisfy a wait on an object.
 pub fn satisfy_wait(obj: *DispatchHeader) void {
+    std.debug.assert(obj.lock.is_locked());
+
     const all = obj.type == .Notification;
 
     // Go through (potentially) all wait blocks and try to satisfy them.
@@ -200,6 +210,8 @@ pub fn satisfy_wait(obj: *DispatchHeader) void {
         // Get the first waitblock from the list.
         var wb: *WaitBlock = @fieldParentPtr("link", obj.waitblocks.first());
         var td = wb.thread;
+        std.debug.assert(wb.status == .Active);
+        std.debug.assert(wb.object == obj);
 
         // Remove it.
         wb.link.remove();