turnstile fixes

2026-10-22 01:34:45 +02:00
parent dbb543de66
commit b83f1fe6d9
10 changed files with 419 additions and 82 deletions
@@ -0,0 +1,214 @@
+// SPDX-License-Identifier: ISC
+
+#include <zinnia/assert.h>
+#include <zinnia/atomic.h>
+#include <zinnia/percpu.h>
+#include <zinnia/rwlock.h>
+#include <zinnia/sched.h>
+#include <zinnia/turnstile.h>
+#include <arch/cpu.h>
+
+constexpr uintptr_t RW_WAITERS = 1 << 0;
+constexpr uintptr_t RW_WRITE_WANT = 1 << 1;
+constexpr uintptr_t RW_WRITE_LOCK = 1 << 2;
+constexpr uintptr_t RW_FLAGS = RW_WAITERS | RW_WRITE_WANT | RW_WRITE_LOCK;
+constexpr uintptr_t RW_READ_INCR = 1 << 3;
+constexpr uintptr_t RW_ANON_OWNER = 1 << 3;
+constexpr size_t RW_SPIN_MAX = 2048;
+
+static_assert(alignof(struct task) > RW_WRITE_LOCK, "Task pointers need free low bits");
+
+static inline uintptr_t self_token() {
+    struct task* current = percpu_get()->sched.current;
+    return current ? (uintptr_t)current : RW_ANON_OWNER;
+}
+
+static inline uintptr_t holders(uintptr_t word) {
+    return word & ~RW_FLAGS;
+}
+
+static inline uintptr_t writer_of(uintptr_t word) {
+    return (word & RW_WRITE_LOCK) ? holders(word) : 0;
+}
+
+static inline bool writer_running(uintptr_t owner) {
+    if (owner == RW_ANON_OWNER)
+        return true;
+    return atomic_load(&((struct task*)owner)->on_cpu, ATOMIC_RELAXED);
+}
+
+static inline uint8_t priority_of(struct task* task) {
+    return atomic_load(&task->priority, ATOMIC_RELAXED);
+}
+
+bool rwlock_try_read_lock(struct rwlock* rw) {
+    uintptr_t word = atomic_load(&rw->word, ATOMIC_RELAXED);
+    while (!(word & (RW_WRITE_LOCK | RW_WRITE_WANT))) {
+        if (atomic_compare_exchange_weak(&rw->word, &word, word + RW_READ_INCR, ATOMIC_ACQUIRE, ATOMIC_RELAXED))
+            return true;
+    }
+    return false;
+}
+
+bool rwlock_try_write_lock(struct rwlock* rw) {
+    uintptr_t self = self_token();
+    uintptr_t word = atomic_load(&rw->word, ATOMIC_RELAXED);
+    while (holders(word) == 0 && !(word & RW_WRITE_LOCK)) {
+        if (atomic_compare_exchange_weak(&rw->word, &word, word | self | RW_WRITE_LOCK, ATOMIC_ACQUIRE, ATOMIC_RELAXED))
+            return true;
+    }
+    return false;
+}
+
+static void lock_slow(struct rwlock* rw, bool write) {
+    uintptr_t self = self_token();
+    enum turnstile_queue queue = write ? TURNSTILE_EXCLUSIVE : TURNSTILE_SHARED;
+    uintptr_t need = write ? (RW_WAITERS | RW_WRITE_WANT) : RW_WAITERS;
+    uintptr_t blocked_by = write ? 0 : (RW_WRITE_LOCK | RW_WRITE_WANT);
+    size_t spins = 0;
+    uintptr_t word = atomic_load(&rw->word, ATOMIC_RELAXED);
+    assert(writer_of(word) != self || self == RW_ANON_OWNER, "Recursive lock of rwlock %p", rw);
+
+    while (true) {
+        bool held = holders(word) != 0;
+        if (write ? !held : !(word & blocked_by)) {
+            uintptr_t new_word = write ? (word | self | RW_WRITE_LOCK) : (word + RW_READ_INCR);
+            if (atomic_compare_exchange_weak(&rw->word, &word, new_word, ATOMIC_ACQUIRE, ATOMIC_RELAXED))
+                return;
+            continue;
+        }
+
+        uintptr_t owner = writer_of(word);
+        if (owner && spins < RW_SPIN_MAX && writer_running(owner)) {
+            spins++;
+            arch_cpu_pause();
+            word = atomic_load(&rw->word, ATOMIC_RELAXED);
+            continue;
+        }
+
+        if (self == RW_ANON_OWNER || owner == RW_ANON_OWNER) {
+            spins = 0;
+            arch_cpu_pause();
+            word = atomic_load(&rw->word, ATOMIC_RELAXED);
+            continue;
+        }
+
+        if ((word & need) != need) {
+            if (!atomic_compare_exchange_weak(&rw->word, &word, word | need, ATOMIC_RELAXED, ATOMIC_RELAXED))
+                continue;
+        }
+
+        struct turnstile_chain* chain = turnstile_lock(rw);
+        word = atomic_load(&rw->word, ATOMIC_RELAXED);
+        held = holders(word) != 0;
+        if (!held || (word & need) != need || (!write && !(word & blocked_by))) {
+            turnstile_unlock(chain);
+            spins = 0;
+            continue;
+        }
+
+        turnstile_block(chain, rw, (struct task*)writer_of(word), queue);
+        return;
+    }
+}
+
+void rwlock_read_lock(struct rwlock* rw) {
+    uintptr_t word = atomic_load(&rw->word, ATOMIC_RELAXED);
+    if (!(word & (RW_WRITE_LOCK | RW_WRITE_WANT)) &&
+        atomic_compare_exchange_weak(&rw->word, &word, word + RW_READ_INCR, ATOMIC_ACQUIRE, ATOMIC_RELAXED))
+        return;
+    lock_slow(rw, false);
+}
+
+void rwlock_write_lock(struct rwlock* rw) {
+    uintptr_t expected = 0;
+    if (!atomic_compare_exchange(&rw->word, &expected, self_token() | RW_WRITE_LOCK, ATOMIC_ACQUIRE, ATOMIC_RELAXED))
+        lock_slow(rw, true);
+}
+
+static uintptr_t readers_to_wake(struct task* reader, struct task* writer) {
+    uint8_t min = writer ? priority_of(writer) : SCHED_PRIO_MIN;
+    uintptr_t count = 0;
+    for (struct task* t = reader; t && priority_of(t) >= min; t = tailq_next(t, ts_link))
+        count++;
+    return count;
+}
+
+static void release_slow(struct rwlock* rw, bool write, uintptr_t keep) {
+    struct turnstile_chain* chain = turnstile_lock(rw);
+    uintptr_t word = atomic_load(&rw->word, ATOMIC_RELAXED);
+
+    if (!write) {
+        while (holders(word) > RW_READ_INCR) {
+            if (atomic_compare_exchange_weak(&rw->word, &word, word - RW_READ_INCR, ATOMIC_RELEASE, ATOMIC_RELAXED)) {
+                turnstile_unlock(chain);
+                return;
+            }
+        }
+    }
+
+    bool more_writers, more_readers;
+    struct task* writer = turnstile_first(chain, rw, TURNSTILE_EXCLUSIVE, &more_writers);
+    struct task* reader = turnstile_first(chain, rw, TURNSTILE_SHARED, &more_readers);
+    uintptr_t count = (write || !writer) ? readers_to_wake(reader, writer) : 0;
+
+    if (count == 0 && writer && keep == 0) {
+        uintptr_t new_word = (uintptr_t)writer | RW_WRITE_LOCK;
+        if (more_writers || reader)
+            new_word |= RW_WAITERS;
+        if (more_writers)
+            new_word |= RW_WRITE_WANT;
+        atomic_store(&rw->word, new_word, ATOMIC_RELEASE);
+        turnstile_wake(chain, rw, writer, writer);
+    } else {
+        uintptr_t new_word = keep + count * RW_READ_INCR;
+        if (writer)
+            new_word |= RW_WAITERS | RW_WRITE_WANT;
+        atomic_store(&rw->word, new_word, ATOMIC_RELEASE);
+        if (count == 0)
+            turnstile_waive(chain, rw);
+        while (count--)
+            turnstile_wake(chain, rw, turnstile_first(chain, rw, TURNSTILE_SHARED, &more_readers), nullptr);
+    }
+    turnstile_unlock(chain);
+}
+
+void rwlock_read_unlock(struct rwlock* rw) {
+    uintptr_t word = atomic_load(&rw->word, ATOMIC_RELAXED);
+    while (true) {
+        assert(holders(word) != 0 && !(word & RW_WRITE_LOCK), "Read unlock of rwlock %p that isn't read-held", rw);
+        if (holders(word) == RW_READ_INCR && (word & RW_WAITERS)) {
+            release_slow(rw, false, 0);
+            return;
+        }
+        if (atomic_compare_exchange_weak(&rw->word, &word, word - RW_READ_INCR, ATOMIC_RELEASE, ATOMIC_RELAXED))
+            return;
+    }
+}
+
+void rwlock_write_unlock(struct rwlock* rw) {
+    uintptr_t self = self_token();
+    uintptr_t expected = self | RW_WRITE_LOCK;
+    if (atomic_compare_exchange(&rw->word, &expected, 0, ATOMIC_RELEASE, ATOMIC_RELAXED))
+        return;
+
+    assert(writer_of(expected) == self, "Rwlock %p unlocked by a task that doesn't own it", rw);
+    release_slow(rw, true, 0);
+}
+
+void rwlock_downgrade(struct rwlock* rw) {
+    uintptr_t self = self_token();
+    uintptr_t expected = self | RW_WRITE_LOCK;
+    if (atomic_compare_exchange(&rw->word, &expected, RW_READ_INCR, ATOMIC_RELEASE, ATOMIC_RELAXED))
+        return;
+
+    assert(writer_of(expected) == self, "Rwlock %p downgraded by a task that doesn't own it", rw);
+    release_slow(rw, true, RW_READ_INCR);
+}
+
+bool rwlock_try_upgrade(struct rwlock* rw) {
+    uintptr_t word = atomic_load(&rw->word, ATOMIC_RELAXED);
+    assert(holders(word) != 0 && !(word & RW_WRITE_LOCK), "Upgrade of rwlock %p that isn't read-held", rw);
+    uintptr_t expected = RW_READ_INCR;
+    return atomic_compare_exchange(&rw->word, &expected, self_token() | RW_WRITE_LOCK, ATOMIC_ACQUIRE, ATOMIC_RELAXED);
+}