ke: optimize mutex

Optimize the mutex code in a few ways:

- Split the fast and slow paths into their own functions, inline the fast path.
- Raise IPL only when acquiring a turnstile lock.
- Avoid taking the turnstile lock if the "waiters" bit is not set.
- Do adaptive spinning only when the lock owner is running.

src/sys/kern/mutex.zig

@@ -1,37 +1,68 @@
 //! Adaptive lock implementation.
 
-const std = @import("std");
 const rtl = @import("rtl");
 const config = @import("config");
 
+const std = @import("std");
+const assert = std.debug.assert;
+
 const ke = @import("root").ke;
 const kep = ke.private;
 
-/// Number of spins before blocking.
+/// Number of spins while the lock's owner is running before blocking.
-const optimistic_spins = 100;
+// TODO: tune this!
+const optimistic_spins = 4000;
+
+/// owner | waiters (bit 0)
+const State = packed struct(usize) {
+    value: usize = 0,
+
+    const waiters: usize = 1;
+    const uncontended: State = @fromBackingInt(0);
+
+    fn from(td: *ke.Thread) State {
+        return .{ .value = @intFromPtr(td) };
+    }
+
+    inline fn owner(s: State) *ke.Thread {
+        return @ptrFromInt(s.value & ~waiters);
+    }
+
+    inline fn has_waiters(s: State) bool {
+        return s.value & waiters != 0;
+    }
+
+    inline fn set_waiters(s: State) State {
+        return .{ .value = s.value | waiters };
+    }
+
+    inline fn clear_waiters(s: State) State {
+        return .{ .value = s.value & ~waiters };
+    }
+};
 
 pub const Mutex = struct {
     /// The thread currently holding the mutex, null if unlocked.
-    owner: std.atomic.Value(?*ke.Thread),
+    state: std.atomic.Value(State),
-    wd: kep.warden.LockData,
+    wd: kep.warden.LockData = undefined,
 
     pub fn init(class: []const u8) Mutex {
         if (config.warden) {
             const cl = kep.warden.find_or_create_lock_class(class);
             return .{
-                .owner = std.atomic.Value(?*ke.Thread).init(null),
+                .state = std.atomic.Value(State).init(.{}),
                 .wd = cl,
             };
         }
 
-        return .{ .owner = std.atomic.Value(?*ke.Thread).init(null) };
+        return .{ .state = std.atomic.Value(State).init(.{}) };
     }
 
     pub fn is_locked(self: *Mutex) bool {
-        return self.owner.load(.monotonic) != null;
+        return self.state.load(.monotonic) != State.uncontended;
     }
 
-    pub fn acquire(m: *Mutex) void {
+    pub inline fn acquire(m: *Mutex) void {
         const curtd = ke.thread.current();
 
         if (config.warden) {
@@ -42,91 +73,160 @@ pub const Mutex = struct {
             kep.warden.acquired(m, m.wd, .Lock);
         };
 
-        // Very fast path: the lock is uncontended and we can acquire it
+        // Fast path: the lock is uncontended and we can acquire it
-        //immediately.
+        // immediately.
-        const cur_owner = m.owner.cmpxchgStrong(
+        _ = m.state.cmpxchgStrong(
-            null,
+            State.uncontended,
-            curtd,
+            State.from(curtd),
             .acquire,
             .monotonic,
-        );
+        ) orelse {
+            @branchHint(.likely);
 
-        if (cur_owner == null) {
+            // We got the lock.
             return;
+        };
+
+        // Nope, go to the slow path.
+        m.acquire_slow();
     }
 
-        // Fast path: Try to acquire without a turnstile by spinning a bit.
+    fn acquire_slow(m: *Mutex) void {
-        for (0..optimistic_spins) |_| {
+        @branchHint(.unlikely);
-            if (m.owner.load(.monotonic) != null) {
-                std.atomic.spinLoopHint();
-                continue;
-            }
 
-            if (m.owner.cmpxchgWeak(
+        const curtd = ke.thread.current();
-                null,
-                curtd,
-                .acquire,
-                .monotonic,
-            ) != null) {
-                continue;
-            }
 
-            return;
+        retry: while (true) {
-        }
+            var state = m.state.load(.monotonic);
 
-        const ipl = ke.ipl.raise(.Dispatch);
+            if (state == State.uncontended) {
-
+                // The lock seems free, retry acquiring it via CAS.
-        // Slow path: contended, block on a turnstile.
+                if (m.state.cmpxchgStrong(
-        while (true) {
+                    State.uncontended,
-            const ts = kep.turnstile.lookup(m);
+                    State.from(curtd),
-            const owner = m.owner.load(.monotonic);
-
-            if (owner == null) {
-                // Lock was released between lookup and here.
-                kep.turnstile.exit(m, ts);
-                if (m.owner.cmpxchgStrong(
-                    null,
-                    curtd,
                     .acquire,
                     .monotonic,
                 ) == null) {
                     break;
                 }
-                continue;
+            }
+
+            // Still contended, lock the turnstile chain.
+            var ipl = ke.ipl.raise(.Dispatch);
+            const chain = kep.turnstile.lock_chain(m);
+
+            state = m.state.load(.monotonic);
+
+            if (state == State.uncontended) {
+                // Something changed, retry acquiring.
+                kep.turnstile.unlock_chain(chain);
+                ke.ipl.lower(ipl);
+                continue :retry;
+            }
+
+            // Set the waiters bit so that release() will also take the lock.
+            waiters: while (!state.has_waiters()) {
+                if (state == State.uncontended) {
+                    kep.turnstile.unlock_chain(chain);
+                    ke.ipl.lower(ipl);
+                    continue :retry;
+                }
+
+                const marked = state.set_waiters();
+                state = m.state.cmpxchgWeak(
+                    state,
+                    marked,
+                    .monotonic,
+                    .monotonic,
+                ) orelse {
+                    state = marked;
+                    break :waiters;
+                };
+            }
+
+            // Look up the current owner, there is a guarantee that the thread
+            // pointer is still valid while we hold the lock here because
+            // upon release, whomever holds the lock right now will notice the
+            // waiters bit and try to take the same chain lock.
+            const owner = state.owner();
+
+            if (owner.state.load(.monotonic) == .Running) {
+                const cpu = owner.cpu orelse unreachable;
+                const cpu_curtd = kep.sched.current_thread.remote(cpu);
+
+                // The lock's owner is running, spin adaptively
+                // until it stops running or we run out of gas.
+                kep.turnstile.unlock_chain(chain);
+                ke.ipl.lower(ipl);
+
+                var gas: usize = 0;
+
+                while (@atomicLoad(
+                    ?*ke.Thread,
+                    cpu_curtd,
+                    .unordered,
+                ) == owner and gas < optimistic_spins) : (gas += 1) {
+                    if (m.state.load(.monotonic) == State.uncontended) {
+                        // Got released, retry acquiring it.
+                        continue :retry;
+                    }
+
+                    std.atomic.spinLoopHint();
+                }
+
+                // Couldn't get it via adaptive spinning, now actually block
+                // on a turnstile.
+
+                ipl = ke.ipl.raise(.Dispatch);
+                _ = kep.turnstile.lock_chain(m);
+
+                if (m.state.load(.monotonic) != state) {
+                    // Could've gotten released in the meantime, try again.
+                    kep.turnstile.unlock_chain(chain);
+                    ke.ipl.lower(ipl);
+                    continue :retry;
+                }
+            }
+
+            const ts = kep.turnstile.lookup_locked(m, chain);
+
+            if (m.state.load(.monotonic) != state) {
+                // Something changed between lookup and here.
+                kep.turnstile.exit(m, ts);
+                ke.ipl.lower(ipl);
+                continue :retry;
             }
 
             // Block until woken.
             kep.turnstile.block(
                 ts,
                 m,
-                .{ .single = owner orelse unreachable },
+                .{ .single = state.owner() },
                 .Exclusive,
             );
-
-            // Re-try acquisition after wakeup.
-            if (m.owner.cmpxchgStrong(
-                null,
-                curtd,
-                .acquire,
-                .monotonic,
-            ) == null) {
-                break;
         }
     }
 
-        ke.ipl.lower(ipl);
+    pub inline fn release(m: *Mutex) void {
+        const old = m.state.swap(State.uncontended, .release);
+
+        if (old.has_waiters()) {
+            m.release_slow();
         }
 
-    pub fn release(m: *Mutex) void {
+        if (config.warden) {
+            kep.warden.released(m, .Lock);
+        }
+    }
+
+    /// Called when there were waiters on the lock that we need to wake.
+    fn release_slow(m: *Mutex) void {
+        @branchHint(.unlikely);
+
+        // Look up the turnstile.
         const ipl = ke.ipl.raise(.Dispatch);
         const ts = kep.turnstile.lookup(m);
 
-        m.owner.store(null, .release);
-
-        defer if (config.warden) {
-            kep.warden.released(m, .Lock);
-        };
-
         if (ts == null) {
             kep.turnstile.exit(m, ts);
             ke.ipl.lower(ipl);
@@ -136,10 +236,10 @@ pub const Mutex = struct {
         var waiters: rtl.List = undefined;
         const turn = ts orelse unreachable;
 
-        // Note: wake up all waiters.
+        // Wake ALL the waiters. This so-called "lock barging" is meant to avoid
-        // This so-called "lock barging" (name from WTF::ParkingLot) has been
+        // lock convoys.
-        // to be better because this avoids lock convoys,
+        //
-        // see this mysterious 70s paper:
+        // See this mysterious 70s paper:
         // https://dl.acm.org/doi/pdf/10.1145/850657.850659
         kep.turnstile.signal(
             turn,