Add turnstiles, more sched fixes

2026-10-20 23:13:29 +02:00
parent a8bf7d7a28
commit 05c0920998
12 changed files with 1072 additions and 83 deletions
@@ -0,0 +1,289 @@
+// SPDX-License-Identifier: ISC
+
+#include <zinnia/alloc.h>
+#include <zinnia/assert.h>
+#include <zinnia/atomic.h>
+#include <zinnia/percpu.h>
+#include <zinnia/sched.h>
+#include <zinnia/spin.h>
+#include <zinnia/turnstile.h>
+#include <arch/cpu.h>
+#include <sched/internal.h>
+
+constexpr size_t TURNSTILE_CHAIN_BITS = 7;
+constexpr size_t TURNSTILE_CHAINS = (1 << TURNSTILE_CHAIN_BITS);
+
+struct turnstile_chain {
+    // Aligned to cache line to avoid false sharing.
+    alignas(CACHE_LINE_SIZE) struct spinlock lock;
+    struct turnstile* head;
+};
+
+static struct turnstile_chain chains[TURNSTILE_CHAINS];
+
+// Gets the chain from a lock object address.
+static struct turnstile_chain* chain_for(const void* obj) {
+    uint64_t hash = ((uintptr_t)obj >> 3) * 0x9E3779B97F4A7C15ull;
+    return &chains[hash >> (64 - TURNSTILE_CHAIN_BITS)];
+}
+
+static struct turnstile* chain_lookup(struct turnstile_chain* chain, const void* obj) {
+    struct turnstile* ts = chain->head;
+    while (ts && ts->obj != obj)
+        ts = ts->chain_next;
+    return ts;
+}
+
+static void chain_remove(struct turnstile_chain* chain, struct turnstile* ts) {
+    struct turnstile** link = &chain->head;
+    while (*link != ts)
+        link = &(*link)->chain_next;
+    *link = ts->chain_next;
+}
+
+struct turnstile* turnstile_alloc() {
+    struct turnstile* ts = mem_alloc(sizeof(struct turnstile), 0);
+    if (ts) {
+        for (size_t q = 0; q < TURNSTILE_NUM_QUEUES; q++)
+            tailq_init(&ts->waiters[q]);
+    }
+    return ts;
+}
+
+static bool has_waiters(struct turnstile* ts) {
+    for (size_t q = 0; q < TURNSTILE_NUM_QUEUES; q++) {
+        if (!tailq_empty(&ts->waiters[q]))
+            return true;
+    }
+    return false;
+}
+
+void turnstile_free(struct turnstile* ts) {
+    assert(ts->obj == nullptr, "Freeing turnstile %p that is still in use", ts);
+    mem_free(ts);
+}
+
+struct turnstile_chain* turnstile_lock(const void* obj) {
+    struct turnstile_chain* chain = chain_for(obj);
+    spin_lock(&chain->lock);
+    return chain;
+}
+
+void turnstile_unlock(struct turnstile_chain* chain) {
+    spin_unlock(&chain->lock);
+}
+
+static uint8_t task_priority(struct task* task) {
+    return atomic_load(&task->priority, ATOMIC_RELAXED);
+}
+
+// Inserts `task` into the waiters of `ts`, behind all waiters that are at least as urgent.
+static void waiters_insert(struct turnstile* ts, struct task* task) {
+    uint8_t prio = task_priority(task);
+    struct task* pos = tailq_first(&ts->waiters[task->ts_queue]);
+    while (pos && task_priority(pos) >= prio)
+        pos = tailq_next(pos, ts_link);
+
+    if (pos)
+        tailq_insert_before(pos, task, ts_link);
+    else
+        tailq_insert_tail(&ts->waiters[task->ts_queue], task, ts_link);
+}
+
+// Updates the cached priority of the first waiter. Returns true if it changed.
+static bool update_top(struct turnstile* ts) {
+    uint8_t top = SCHED_PRIO_MIN;
+    for (size_t q = 0; q < TURNSTILE_NUM_QUEUES; q++) {
+        struct task* first = tailq_first(&ts->waiters[q]);
+        if (first && task_priority(first) > top)
+            top = task_priority(first);
+    }
+    return atomic_exchange(&ts->top_priority, top, ATOMIC_RELAXED) != top;
+}
+
+// Recomputes the inherited priority of `task`. Returns true if its effective priority changed.
+static bool recompute_locked(struct task* task) {
+    uint8_t inherited = SCHED_PRIO_MIN;
+    struct turnstile* ts;
+    for (ts = tailq_first(&task->pi_lenders); ts; ts = tailq_next(ts, lender_link)) {
+        uint8_t top = atomic_load(&ts->top_priority, ATOMIC_RELAXED);
+        if (top > inherited)
+            inherited = top;
+    }
+    task->lent_prio = inherited;
+    return sched_update_priority_locked(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) {
+    struct task* old = ts->owner;
+    if (old == owner)
+        return;
+
+    if (old) {
+        spin_lock(&old->lock);
+        tailq_remove(&old->pi_lenders, ts, lender_link);
+        recompute_locked(old);
+        spin_unlock(&old->lock);
+    }
+
+    ts->owner = owner;
+    if (owner) {
+        spin_lock(&owner->lock);
+        tailq_insert_tail(&owner->pi_lenders, ts, lender_link);
+        spin_unlock(&owner->lock);
+    }
+}
+
+static void requeue(struct turnstile* ts, struct task* task) {
+    tailq_remove(&ts->waiters[task->ts_queue], task, ts_link);
+    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;
+    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)
+                break;
+            retried = true;
+            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)
+            turnstile_unlock(chain);
+        chain = next_chain;
+        if (!blocked)
+            break;
+
+        requeue(next, owner);
+        if (!update_top(next))
+            break;
+
+        ts = next;
+        retried = false;
+    }
+    turnstile_unlock(chain);
+}
+
+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");
+
+    // Lend our turnstile to the lock if it doesn't have one yet, otherwise donate it.
+    struct turnstile* ts = chain_lookup(chain, obj);
+    if (ts == nullptr) {
+        ts = self->turnstile;
+        ts->obj = obj;
+        ts->owner = nullptr;
+        ts->spare = nullptr;
+        ts->chain_next = chain->head;
+        chain->head = ts;
+    } else {
+        self->turnstile->spare = ts->spare;
+        ts->spare = self->turnstile;
+    }
+    self->turnstile = nullptr;
+
+    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->ts_queue = queue;
+    waiters_insert(ts, self);
+    update_top(ts);
+    spin_unlock(&self->lock);
+
+    set_owner(ts, owner);
+    propagate(chain, ts);
+
+    // 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);
+}
+
+struct task* turnstile_first(
+    struct turnstile_chain* chain,
+    const void* obj,
+    enum turnstile_queue queue,
+    bool* have_more
+) {
+    struct turnstile* ts = chain_lookup(chain, obj);
+    struct task* first = ts ? tailq_first(&ts->waiters[queue]) : nullptr;
+    *have_more = first && tailq_next(first, ts_link) != nullptr;
+    return 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);
+    assert(ts && ts->obj == obj, "Task %p isn't blocked on %p", task, obj);
+
+    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);
+        chain_remove(chain, ts);
+        ts->obj = nullptr;
+        task->turnstile = ts;
+    } else {
+        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);
+        }
+    }
+
+    spin_lock(&task->lock);
+    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);
+    struct turnstile* ts = chain_lookup(chain, obj);
+    if (ts == nullptr || atomic_load(&task->blocked_ts, ATOMIC_RELAXED) != ts) {
+        turnstile_unlock(chain);
+        return;
+    }
+
+    requeue(ts, task);
+    if (!update_top(ts)) {
+        turnstile_unlock(chain);
+        return;
+    }
+    propagate(chain, ts);
+}