turnstile fixes

include/zinnia/rwlock.h

@@ -0,0 +1,20 @@
+// SPDX-License-Identifier: ISC
+
+#pragma once
+
+#include <zinnia/atomic.h>
+
+struct rwlock {
+    atomic_uintptr_t word;
+};
+
+void rwlock_read_lock(struct rwlock* rw);
+void rwlock_read_unlock(struct rwlock* rw);
+bool rwlock_try_read_lock(struct rwlock* rw);
+
+void rwlock_write_lock(struct rwlock* rw);
+void rwlock_write_unlock(struct rwlock* rw);
+bool rwlock_try_write_lock(struct rwlock* rw);
+
+void rwlock_downgrade(struct rwlock* rw);
+bool rwlock_try_upgrade(struct rwlock* rw);

include/zinnia/sched.h

@@ -60,6 +60,7 @@ struct task {
     atomic_bool wake_pending;
 
     struct spinlock lock;
+    _Atomic(struct spinlock*) lockp;
 
     atomic_bool on_cpu;                // Whether a CPU is still running on this task's context.
     uint8_t base_prio;                 // Set by sched_set_priority().
@@ -68,8 +69,7 @@ struct task {
     struct scheduler* last_rq;         // The most recent run queue this task was scheduled on.
 
     struct turnstile* turnstile;             // Turnstile owned by this task.
-    _Atomic(struct turnstile*) blocked_ts;   // The turnstile this task is blocked on.
-    _Atomic(const void*) blocked_obj;        // The lock that `blocked_ts` belongs to.
+    struct turnstile* blocked_ts;            // The turnstile this task is blocked on.
     tailq_head(struct turnstile) pi_lenders; // Turnstiles of locks this task owns.
     tailq_link(struct task) ts_link;         // Link in the waiters of `blocked_ts`.
     uint8_t ts_queue;
@@ -107,6 +107,8 @@ typedef void (*task_fn_t)(uintptr_t arg0);
 
 errno_t task_create(const char* name, struct vma* space, task_fn_t entry, uintptr_t arg0, struct task** out);
 void task_destroy(struct task* task);
+void task_lock(struct task* task);
+void task_unlock(struct task* task);
 [[noreturn]]
 void task_entry(task_fn_t entry, uintptr_t arg0, struct task* prev);
 
@@ -122,7 +124,7 @@ void sched_block(struct scheduler* sched);
 // Makes a blocked task runnable again, or marks the wake as pending if it hasn't blocked yet.
 void sched_wake(struct task* task);
 
-// Same as sched_wake(), but `task->lock` must be held.
+// Same as sched_wake(), but the task must be locked with task_lock().
 void sched_wake_locked(struct task* task);
 
 [[noreturn]]

include/zinnia/turnstile.h

@@ -22,6 +22,7 @@ struct turnstile {
 
     // The task that inherits the priority from the waiters.
     struct task* owner;
+    bool lending;
 
     // Priority of the most urgent waiter in any queue.
     _Atomic(uint8_t) top_priority;
@@ -55,4 +56,6 @@ struct task* turnstile_first(
 
 void turnstile_wake(struct turnstile_chain* chain, const void* obj, struct task* task, struct task* new_owner);
 
+void turnstile_waive(struct turnstile_chain* chain, const void* obj);
+
 void turnstile_reprioritize(struct task* task);

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

sched/Kconfig

@@ -9,7 +9,7 @@ choice SCHED_ALG
 config SCHED_DUMB
   bool "Dumb Scheduler"
   help
-    Very basic scheduler, without any priorities.
+    Very basic scheduler.
 
 endchoice
 

sched/core.c

@@ -28,9 +28,36 @@ void sched_to_user_context(uintptr_t ctx) {
     arch_sched_jump_to_context((struct arch_context*)ctx);
 }
 
+static struct spinlock task_cache_lock;
+static tailq_head(struct task) task_cache = {nullptr, &task_cache.tqh_first};
+
+static struct task* task_alloc() {
+    spin_lock(&task_cache_lock);
+    struct task* task = tailq_first(&task_cache);
+    if (task)
+        tailq_remove(&task_cache, task, next);
+    spin_unlock(&task_cache_lock);
+
+    if (task == nullptr)
+        return mem_alloc(sizeof(struct task), 0);
+
+    size_t keep_start = offsetof(struct task, on_cpu);
+    size_t keep_end = keep_start + sizeof(task->on_cpu);
+    memset(task, 0, keep_start);
+    memset((char*)task + keep_end, 0, sizeof(struct task) - keep_end);
+    return task;
+}
+
+static void task_release(struct task* task) {
+    spin_lock(&task_cache_lock);
+    tailq_insert_tail(&task_cache, task, next);
+    spin_unlock(&task_cache_lock);
+}
+
 static void task_init_common(struct task* task, const char* name, struct vma* space) {
     task->space = space;
     spin_init(&task->lock);
+    atomic_store(&task->lockp, &task->lock, ATOMIC_RELAXED);
     strncpy(task->name, name, sizeof(task->name));
     task->base_prio = SCHED_PRIO_DEFAULT;
     task->time_slice = SCHED_TIME_SLICE_NS;
@@ -41,13 +68,13 @@ static void task_init_common(struct task* task, const char* name, struct vma* sp
 }
 
 errno_t task_create_current(const char* name, struct vma* space, struct task** out) {
-    struct task* task = mem_alloc(sizeof(struct task), 0);
+    struct task* task = task_alloc();
     if (task == nullptr)
         return ENOMEM;
 
     task->turnstile = turnstile_alloc();
     if (task->turnstile == nullptr) {
-        mem_free(task);
+        task_release(task);
         return ENOMEM;
     }
 
@@ -62,7 +89,7 @@ errno_t task_create(const char* name, struct vma* space, task_fn_t entry, uintpt
     if (!out)
         return EINVAL;
 
-    struct task* new_task = mem_alloc(sizeof(struct task), 0);
+    struct task* new_task = task_alloc();
     if (new_task == nullptr)
         return ENOMEM;
 
@@ -70,7 +97,7 @@ errno_t task_create(const char* name, struct vma* space, task_fn_t entry, uintpt
     new_task->id = atomic_fetch_add(&last_tid, 1, ATOMIC_ACQUIRE);
     new_task->turnstile = turnstile_alloc();
     if (new_task->turnstile == nullptr) {
-        mem_free(new_task);
+        task_release(new_task);
         return ENOMEM;
     }
 
@@ -78,7 +105,7 @@ errno_t task_create(const char* name, struct vma* space, task_fn_t entry, uintpt
     errno_t s = stack_init(&new_task->kernel_stack, KERNEL_STACK_SIZE);
     if (s != 0) {
         turnstile_free(new_task->turnstile);
-        mem_free(new_task);
+        task_release(new_task);
         return s;
     }
     new_task->user_stack = 0;
@@ -88,7 +115,7 @@ errno_t task_create(const char* name, struct vma* space, task_fn_t entry, uintpt
     if (s != 0) {
         stack_destroy(&new_task->kernel_stack);
         turnstile_free(new_task->turnstile);
-        mem_free(new_task);
+        task_release(new_task);
         return s;
     }
 
@@ -102,5 +129,19 @@ void task_destroy(struct task* task) {
     turnstile_free(task->turnstile);
     arch_task_fini(&task->context);
     stack_destroy(&task->kernel_stack);
-    mem_free(task);
+    task_release(task);
+}
+
+void task_lock(struct task* task) {
+    while (true) {
+        struct spinlock* lock = atomic_load(&task->lockp, ATOMIC_ACQUIRE);
+        spin_lock(lock);
+        if (atomic_load(&task->lockp, ATOMIC_RELAXED) == lock)
+            return;
+        spin_unlock(lock);
+    }
+}
+
+void task_unlock(struct task* task) {
+    spin_unlock(atomic_load(&task->lockp, ATOMIC_RELAXED));
 }

sched/sched_dumb.c

@@ -112,6 +112,11 @@ static uint8_t effective_prio(struct task* task) {
     return MAX(prio, task->lent_prio);
 }
 
+static void refresh_prio(struct task* task) {
+    if (task->blocked_ts == nullptr)
+        atomic_store(&task->priority, effective_prio(task), ATOMIC_RELAXED);
+}
+
 static size_t quantum(struct task* task) {
     return SCHED_TIME_SLICE_NS * (task->penalty + 1);
 }
@@ -120,7 +125,7 @@ static void feedback_reset(struct scheduler* sched, struct task* task) {
     task->penalty = 0;
     task->time_slice = quantum(task);
     task->feedback_epoch = sched->boost_epoch;
-    atomic_store(&task->priority, effective_prio(task), ATOMIC_RELAXED);
+    refresh_prio(task);
 }
 
 static void boost(struct scheduler* sched) {
@@ -153,7 +158,7 @@ static void account(struct scheduler* sched, struct task* task, uint64_t now) {
         task->penalty++;
     task->time_slice = quantum(task);
     task->feedback_epoch = sched->boost_epoch;
-    atomic_store(&task->priority, effective_prio(task), ATOMIC_RELAXED);
+    refresh_prio(task);
     sched->need_resched = true;
 }
 
@@ -164,6 +169,7 @@ static void enqueue(struct scheduler* sched, struct task* task) {
     task->last_rq = sched;
     if (task->feedback_epoch < sched->boost_epoch)
         feedback_reset(sched, task);
+    refresh_prio(task);
     rq_insert(sched, task);
     bool preempt = atomic_load(&task->priority, ATOMIC_RELAXED) > atomic_load(&sched->current_priority, ATOMIC_RELAXED);
     spin_unlock(&sched->lock);
@@ -273,9 +279,9 @@ void sched_block(struct scheduler* sched) {
     irq_lock();
     struct task* current = sched->current;
 
-    spin_lock(&current->lock);
+    task_lock(current);
     if (atomic_exchange(&current->wake_pending, false, ATOMIC_ACQ_REL)) {
-        spin_unlock(&current->lock);
+        task_unlock(current);
         irq_unlock();
         return;
     }
@@ -286,7 +292,7 @@ void sched_block(struct scheduler* sched) {
     spin_unlock(&sched->lock);
 
     current->state = TASK_STATE_BLOCKED;
-    spin_unlock(&current->lock);
+    task_unlock(current);
 
     do_reschedule(sched, now);
 }
@@ -315,9 +321,9 @@ void sched_wake_locked(struct task* task) {
 }
 
 void sched_wake(struct task* task) {
-    spin_lock(&task->lock);
+    task_lock(task);
     sched_wake_locked(task);
-    spin_unlock(&task->lock);
+    task_unlock(task);
 }
 
 bool sched_update_priority_locked(struct task* task) {
@@ -357,13 +363,12 @@ bool sched_update_priority_locked(struct task* task) {
 
 void sched_set_priority(struct task* task, uint8_t priority) {
     assert(priority <= SCHED_PRIO_MAX, "Priority %u is out of range", priority);
-    spin_lock(&task->lock);
+    task_lock(task);
     task->base_prio = priority;
-    bool changed = sched_update_priority_locked(task);
-    spin_unlock(&task->lock);
-
-    if (changed)
+    if (sched_update_priority_locked(task))
         turnstile_reprioritize(task);
+    else
+        task_unlock(task);
 }
 
 void task_entry(task_fn_t entry, uintptr_t arg0, struct task* prev) {