Add turnstiles, more sched fixes

lib/mutex.c

@@ -0,0 +1,118 @@
+// SPDX-License-Identifier: ISC
+
+#include <zinnia/assert.h>
+#include <zinnia/atomic.h>
+#include <zinnia/mutex.h>
+#include <zinnia/percpu.h>
+#include <zinnia/sched.h>
+#include <zinnia/turnstile.h>
+#include <arch/cpu.h>
+
+// Set while at least one task might be blocked on the mutex.
+constexpr uintptr_t MUTEX_PARKED = 1 << 0;
+constexpr uintptr_t MUTEX_FLAGS = (MUTEX_PARKED);
+constexpr uintptr_t MUTEX_ANON_OWNER = 1 << 1;
+constexpr size_t MUTEX_SPIN_MAX = 2048;
+
+static_assert(alignof(struct task) > MUTEX_ANON_OWNER, "Task pointers need free low bits");
+
+static inline uintptr_t self_token() {
+    struct task* current = percpu_get()->sched.current;
+    return current ? (uintptr_t)current : MUTEX_ANON_OWNER;
+}
+
+static inline uintptr_t owner_of(uintptr_t word) {
+    return word & ~MUTEX_FLAGS;
+}
+
+// Whether it's worth spinning for `owner` to release the lock.
+static inline bool owner_is_running(uintptr_t owner) {
+    if (owner == MUTEX_ANON_OWNER)
+        return true; // No scheduler yet, so there is nothing to sleep on.
+    return atomic_load(&((struct task*)owner)->on_cpu, ATOMIC_RELAXED);
+}
+
+bool mutex_trylock(mutex_t* mutex) {
+    uintptr_t self = self_token();
+    uintptr_t word = atomic_load(mutex, ATOMIC_RELAXED);
+    while (owner_of(word) == 0) {
+        if (atomic_compare_exchange_weak(mutex, &word, word | self, ATOMIC_ACQUIRE, ATOMIC_RELAXED))
+            return true;
+    }
+    return false;
+}
+
+static void mutex_lock_slow(mutex_t* mutex, uintptr_t self) {
+    size_t spins = 0;
+    uintptr_t word = atomic_load(mutex, ATOMIC_RELAXED);
+    assert(owner_of(word) != self || self == MUTEX_ANON_OWNER, "Recursive lock of mutex %p", mutex);
+
+    while (true) {
+        uintptr_t owner = owner_of(word);
+
+        // The mutex is free. Take it, even if others are parked.
+        if (owner == 0) {
+            if (atomic_compare_exchange_weak(mutex, &word, word | self, ATOMIC_ACQUIRE, ATOMIC_RELAXED))
+                return;
+            continue;
+        }
+
+        // The owner is running and will probably release the lock soon, so just wait for it do so.
+        if (spins < MUTEX_SPIN_MAX && owner_is_running(owner)) {
+            spins++;
+            arch_cpu_pause();
+            word = atomic_load(mutex, ATOMIC_RELAXED);
+            continue;
+        }
+
+        // Before the scheduler is up, nobody can be preempted while holding a lock. Just keep spinning.
+        if (self == MUTEX_ANON_OWNER || owner == MUTEX_ANON_OWNER) {
+            spins = 0;
+            continue;
+        }
+
+        // Tell the owner that it has to wake us.
+        if (!(word & MUTEX_PARKED)) {
+            if (!atomic_compare_exchange_weak(mutex, &word, word | MUTEX_PARKED, ATOMIC_RELAXED, ATOMIC_RELAXED))
+                continue;
+        }
+
+        // Only block if the mutex is still locked and the owner will see our flags.
+        struct turnstile_chain* chain = turnstile_lock(mutex);
+        word = atomic_load(mutex, ATOMIC_RELAXED);
+        owner = owner_of(word);
+        if (owner == 0 || !(word & MUTEX_PARKED)) {
+            turnstile_unlock(chain);
+            spins = 0;
+            continue;
+        }
+
+        turnstile_block(chain, mutex, (struct task*)owner, TURNSTILE_EXCLUSIVE);
+        spins = 0;
+        word = atomic_load(mutex, ATOMIC_RELAXED);
+    }
+}
+
+void mutex_lock(mutex_t* mutex) {
+    uintptr_t self = self_token();
+    uintptr_t expected = 0;
+    if (!atomic_compare_exchange(mutex, &expected, self, ATOMIC_ACQUIRE, ATOMIC_RELAXED))
+        mutex_lock_slow(mutex, self);
+}
+
+void mutex_unlock(mutex_t* mutex) {
+    uintptr_t self = self_token();
+    uintptr_t expected = self;
+    if (atomic_compare_exchange(mutex, &expected, 0, ATOMIC_RELEASE, ATOMIC_RELAXED))
+        return;
+
+    assert(owner_of(expected) == self, "Mutex %p unlocked by a task that doesn't own it", mutex);
+
+    struct turnstile_chain* chain = turnstile_lock(mutex);
+    atomic_store(mutex, 0, ATOMIC_RELEASE);
+    bool have_more;
+    struct task* waiter;
+    while ((waiter = turnstile_first(chain, mutex, TURNSTILE_EXCLUSIVE, &have_more)))
+        turnstile_wake(chain, mutex, waiter, nullptr);
+    turnstile_unlock(chain);
+}