turnstile fixes

lib/meson.build

@@ -15,6 +15,7 @@ kernel_sources += files(
     'random.c',
     'ringbuf.c',
     'runtime.c',
+    'rwlock.c',
     'spin.c',
     'stdlib.c',
     'string.c',

lib/mutex.c

@@ -68,6 +68,8 @@ static void mutex_lock_slow(mutex_t* mutex, uintptr_t self) {
         // 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;
+            arch_cpu_pause();
+            word = atomic_load(mutex, ATOMIC_RELAXED);
             continue;
         }
 

lib/rwlock.c

@@ -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);
+}

lib/turnstile.c

@@ -3,10 +3,12 @@
 #include <zinnia/alloc.h>
 #include <zinnia/assert.h>
 #include <zinnia/atomic.h>
+#include <zinnia/irq.h>
 #include <zinnia/percpu.h>
 #include <zinnia/sched.h>
 #include <zinnia/spin.h>
 #include <zinnia/turnstile.h>
+#include <zinnia/utils.h>
 #include <arch/cpu.h>
 #include <sched/internal.h>
 
@@ -20,6 +22,7 @@ struct turnstile_chain {
 };
 
 static struct turnstile_chain chains[TURNSTILE_CHAINS];
+static struct spinlock loser_lock;
 
 // Gets the chain from a lock object address.
 static struct turnstile_chain* chain_for(const void* obj) {
@@ -115,23 +118,26 @@ static bool recompute_locked(struct task* task) {
 }
 
 // Makes the waiters of `ts` lend their priority to `owner`, and recomputes the old owner's priority.
-static void set_owner(struct turnstile* ts, struct task* owner) {
+static void set_owner(struct turnstile_chain* chain, struct turnstile* ts, struct task* owner) {
     struct task* old = ts->owner;
     if (old == owner)
         return;
 
-    if (old) {
-        spin_lock(&old->lock);
+    if (old && ts->lending) {
+        assert(old == percpu_get()->sched.current, "Turnstile %p is owned by another task", ts);
+        task_lock(old);
         tailq_remove(&old->pi_lenders, ts, lender_link);
         recompute_locked(old);
-        spin_unlock(&old->lock);
+        task_unlock(old);
     }
 
     ts->owner = owner;
+    ts->lending = false;
     if (owner) {
-        spin_lock(&owner->lock);
+        assert(atomic_load(&owner->lockp, ATOMIC_RELAXED) == &chain->lock, "New owner %p isn't blocked", owner);
         tailq_insert_tail(&owner->pi_lenders, ts, lender_link);
-        spin_unlock(&owner->lock);
+        ts->lending = true;
+        recompute_locked(owner);
     }
 }
 
@@ -140,54 +146,94 @@ static void requeue(struct turnstile* ts, struct task* task) {
     waiters_insert(ts, task);
 }
 
-// Recomputes the priority of the owner of `ts` along the owner chain.
-static void propagate(struct turnstile_chain* chain, struct turnstile* ts) {
-    bool retried = false;
+static struct spinlock* lock_owner(struct task* owner, struct spinlock* root, struct spinlock* held, bool* lost) {
+    while (true) {
+        struct spinlock* lock = atomic_load(&owner->lockp, ATOMIC_ACQUIRE);
+        if (lock == root || lock == held)
+            return nullptr;
+
+        if (lock == &owner->lock || lock > MAX(root, held)) {
+            spin_lock(lock);
+        } else if (!spin_trylock(lock)) {
+            *lost = true;
+            return lock;
+        }
+
+        if (atomic_load(&owner->lockp, ATOMIC_RELAXED) == lock)
+            return lock;
+        spin_unlock(lock);
+    }
+}
+
+// Recomputes the priority of the owners along the blocking chain of `blocker`.
+static void propagate(struct turnstile_chain* root, struct task* blocker) {
+    struct turnstile_chain* chain = root;
+    struct turnstile* ts = blocker->blocked_ts;
+    size_t depth = 0;
+    size_t forced = 0;
+    bool loser = false;
     while (ts->owner) {
         struct task* owner = ts->owner;
-        spin_lock(&owner->lock);
-        bool changed = recompute_locked(owner);
-        const void* next_obj = atomic_load(&owner->blocked_obj, ATOMIC_RELAXED);
-        spin_unlock(&owner->lock);
-        if ((!changed && !retried) || next_obj == nullptr)
-            break;
-
-        struct turnstile_chain* next_chain = chain_for(next_obj);
-        if (next_chain > chain) {
-            spin_lock(&next_chain->lock);
-        } else if (next_chain < chain && !spin_trylock(&next_chain->lock)) {
-            const void* obj = ts->obj;
-            turnstile_unlock(chain);
-            chain = turnstile_lock(obj);
-            ts = chain_lookup(chain, obj);
-            if (ts == nullptr)
+        assert(owner != blocker, "Deadlock: cycle in the blocking chain of task %p", blocker);
+        bool lost = false;
+        struct spinlock* locked = lock_owner(owner, &root->lock, &chain->lock, &lost);
+        if (lost) {
+            if (chain != root)
+                turnstile_unlock(chain);
+            turnstile_unlock(root);
+            spin_lock(locked);
+            spin_unlock(locked);
+            if (!loser) {
+                spin_lock(&loser_lock);
+                loser = true;
+            }
+            spin_lock(&root->lock);
+            chain = root;
+            forced = MAX(forced, depth + 1);
+            depth = 0;
+            if (atomic_load(&blocker->lockp, ATOMIC_RELAXED) != &root->lock)
                 break;
-            retried = true;
+            ts = blocker->blocked_ts;
             continue;
         }
 
-        struct turnstile* next = chain_lookup(next_chain, next_obj);
-        bool blocked = next && atomic_load(&owner->blocked_ts, ATOMIC_RELAXED) == next;
-        if (next_chain != chain)
+        if (!ts->lending) {
+            tailq_insert_tail(&owner->pi_lenders, ts, lender_link);
+            ts->lending = true;
+        }
+        bool changed = recompute_locked(owner);
+        struct turnstile* next = owner->blocked_ts;
+        if ((!changed && depth >= forced) || next == nullptr) {
+            if (locked)
+                spin_unlock(locked);
+            break;
+        }
+
+        struct turnstile_chain* next_chain = chain_for(next->obj);
+        assert(atomic_load(&owner->lockp, ATOMIC_RELAXED) == &next_chain->lock, "Blocked task %p isn't locked", owner);
+        if (chain != root && chain != next_chain)
             turnstile_unlock(chain);
         chain = next_chain;
-        if (!blocked)
-            break;
 
         requeue(next, owner);
-        if (!update_top(next))
+        if (!update_top(next) && depth + 1 >= forced)
             break;
 
         ts = next;
-        retried = false;
+        depth++;
     }
-    turnstile_unlock(chain);
+    if (chain != root)
+        turnstile_unlock(chain);
+    turnstile_unlock(root);
+    if (loser)
+        spin_unlock(&loser_lock);
 }
 
 void turnstile_block(struct turnstile_chain* chain, const void* obj, struct task* owner, enum turnstile_queue queue) {
     struct task* self = percpu_get()->sched.current;
     assert(self->turnstile != nullptr, "Task is already blocked on a turnstile");
     assert(owner != self, "Task tried to block on a lock it owns");
+    irq_lock();
 
     // Lend our turnstile to the lock if it doesn't have one yet, otherwise donate it.
     struct turnstile* ts = chain_lookup(chain, obj);
@@ -195,6 +241,7 @@ void turnstile_block(struct turnstile_chain* chain, const void* obj, struct task
         ts = self->turnstile;
         ts->obj = obj;
         ts->owner = nullptr;
+        ts->lending = false;
         ts->spare = nullptr;
         ts->chain_next = chain->head;
         chain->head = ts;
@@ -204,24 +251,28 @@ void turnstile_block(struct turnstile_chain* chain, const void* obj, struct task
     }
     self->turnstile = nullptr;
 
+    if (owner && ts->owner == nullptr)
+        ts->owner = owner;
+    assert(owner == nullptr || ts->owner == owner, "Turnstile %p has a stale owner", ts);
+
     spin_lock(&self->lock);
     atomic_store(&self->ts_woken, false, ATOMIC_RELAXED);
-    atomic_store(&self->blocked_ts, ts, ATOMIC_RELAXED);
-    atomic_store(&self->blocked_obj, obj, ATOMIC_RELAXED);
+    self->blocked_ts = ts;
     self->ts_queue = queue;
     waiters_insert(ts, self);
     update_top(ts);
+    atomic_store(&self->lockp, &chain->lock, ATOMIC_RELEASE);
     spin_unlock(&self->lock);
 
-    set_owner(ts, owner);
-    propagate(chain, ts);
+    propagate(chain, self);
 
     // Check if we were actually woken.
     while (!atomic_load(&self->ts_woken, ATOMIC_ACQUIRE))
         sched_block(&percpu_get()->sched);
 
-    spin_lock(&self->lock);
-    spin_unlock(&self->lock);
+    task_lock(self);
+    task_unlock(self);
+    irq_unlock();
 }
 
 struct task* turnstile_first(
@@ -237,16 +288,15 @@ struct task* turnstile_first(
 }
 
 void turnstile_wake(struct turnstile_chain* chain, const void* obj, struct task* task, struct task* new_owner) {
-    struct turnstile* ts = atomic_load(&task->blocked_ts, ATOMIC_RELAXED);
+    struct turnstile* ts = task->blocked_ts;
     assert(ts && ts->obj == obj, "Task %p isn't blocked on %p", task, obj);
+    assert(atomic_load(&task->lockp, ATOMIC_RELAXED) == &chain->lock, "Task %p isn't locked by %p", task, chain);
 
     tailq_remove(&ts->waiters[task->ts_queue], task, ts_link);
-    atomic_store(&task->blocked_ts, nullptr, ATOMIC_RELAXED);
-    atomic_store(&task->blocked_obj, nullptr, ATOMIC_RELAXED);
 
     if (!has_waiters(ts)) {
         // The last waiter takes the lock's turnstile with it.
-        set_owner(ts, nullptr);
+        set_owner(chain, ts, nullptr);
         chain_remove(chain, ts);
         ts->obj = nullptr;
         task->turnstile = ts;
@@ -254,36 +304,35 @@ void turnstile_wake(struct turnstile_chain* chain, const void* obj, struct task*
         task->turnstile = ts->spare;
         ts->spare = ts->spare->spare;
         update_top(ts);
-        set_owner(ts, new_owner);
-        if (new_owner) {
-            spin_lock(&new_owner->lock);
-            recompute_locked(new_owner);
-            spin_unlock(&new_owner->lock);
-        }
+        set_owner(chain, ts, new_owner);
     }
+    task->blocked_ts = nullptr;
 
     spin_lock(&task->lock);
+    atomic_store(&task->lockp, &task->lock, ATOMIC_RELEASE);
     atomic_store(&task->ts_woken, true, ATOMIC_RELEASE);
     sched_wake_locked(task);
     spin_unlock(&task->lock);
 }
 
-void turnstile_reprioritize(struct task* task) {
-    const void* obj = atomic_load(&task->blocked_obj, ATOMIC_RELAXED);
-    if (obj == nullptr)
-        return;
-
-    struct turnstile_chain* chain = turnstile_lock(obj);
+void turnstile_waive(struct turnstile_chain* chain, const void* obj) {
     struct turnstile* ts = chain_lookup(chain, obj);
-    if (ts == nullptr || atomic_load(&task->blocked_ts, ATOMIC_RELAXED) != ts) {
-        turnstile_unlock(chain);
+    if (ts)
+        set_owner(chain, ts, nullptr);
+}
+
+void turnstile_reprioritize(struct task* task) {
+    struct turnstile* ts = task->blocked_ts;
+    if (ts == nullptr) {
+        task_unlock(task);
         return;
     }
 
+    struct turnstile_chain* chain = chain_for(ts->obj);
     requeue(ts, task);
     if (!update_top(ts)) {
         turnstile_unlock(chain);
         return;
     }
-    propagate(chain, ts);
+    propagate(chain, task);
 }