Add turnstiles, more sched fixes

lib/meson.build

@@ -7,6 +7,7 @@ kernel_sources += files(
     'irq.c',
     'ksyms.c',
     'module.c',
+    'mutex.c',
     'panic.c',
     'percpu.c',
     'print.c',
@@ -17,4 +18,6 @@ kernel_sources += files(
     'spin.c',
     'stdlib.c',
     'string.c',
+    'timer.c',
+    'turnstile.c',
 )

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

lib/timer.c

@@ -0,0 +1,115 @@
+// SPDX-License-Identifier: ISC
+
+#include <zinnia/assert.h>
+#include <zinnia/atomic.h>
+#include <zinnia/clock.h>
+#include <zinnia/irq.h>
+#include <zinnia/percpu.h>
+#include <zinnia/spin.h>
+#include <zinnia/timer.h>
+#include <arch/cpu.h>
+#include <arch/timer.h>
+
+struct timer_queue {
+    struct spinlock lock;
+    tailq_head(struct timer) timers;
+};
+
+[[__percpu]]
+static struct timer_queue timer_queue;
+
+void timer_init(struct timer* timer, void (*func)(struct timer* timer)) {
+    timer->func = func;
+    timer->deadline = 0;
+    atomic_init(&timer->queue, nullptr);
+    atomic_init(&timer->running, false);
+}
+
+void timer_cpu_init() {
+    struct timer_queue* q = percpu_ref(timer_queue);
+    spin_init(&q->lock);
+    tailq_init(&q->timers);
+}
+
+static void program_next(struct timer_queue* q) {
+    struct timer* first = tailq_first(&q->timers);
+    if (first)
+        arch_timer_program(first->deadline);
+}
+
+static void remove_locked(struct timer_queue* q, struct timer* timer) {
+    tailq_remove(&q->timers, timer, link);
+    atomic_store(&timer->queue, nullptr, ATOMIC_RELEASE);
+}
+
+static struct timer_queue* lock_queue_of(struct timer* timer) {
+    while (true) {
+        struct timer_queue* q = atomic_load(&timer->queue, ATOMIC_ACQUIRE);
+        if (!q)
+            return nullptr;
+        spin_lock(&q->lock);
+        if (atomic_load(&timer->queue, ATOMIC_RELAXED) == q)
+            return q;
+        spin_unlock(&q->lock);
+    }
+}
+
+static bool dequeue(struct timer* timer) {
+    struct timer_queue* q = lock_queue_of(timer);
+    if (!q)
+        return false;
+    remove_locked(q, timer);
+    spin_unlock(&q->lock);
+    return true;
+}
+
+void timer_arm(struct timer* timer, uint64_t deadline_ns) {
+    dequeue(timer);
+
+    irq_lock();
+    struct timer_queue* q = percpu_ref(timer_queue);
+    spin_lock(&q->lock);
+
+    timer->deadline = deadline_ns;
+    struct timer* pos = tailq_first(&q->timers);
+    while (pos && pos->deadline <= deadline_ns)
+        pos = tailq_next(pos, link);
+    if (pos)
+        tailq_insert_before(pos, timer, link);
+    else
+        tailq_insert_tail(&q->timers, timer, link);
+    atomic_store(&timer->queue, q, ATOMIC_RELEASE);
+
+    if (tailq_first(&q->timers) == timer)
+        arch_timer_program(deadline_ns);
+
+    spin_unlock(&q->lock);
+    irq_unlock();
+}
+
+bool timer_cancel(struct timer* timer) {
+    bool removed = dequeue(timer);
+    while (atomic_load(&timer->running, ATOMIC_ACQUIRE))
+        arch_cpu_pause();
+    return removed;
+}
+
+void timer_interrupt() {
+    struct timer_queue* q = percpu_ref(timer_queue);
+    const uint64_t now = clock_get_elapsed_ns();
+
+    spin_lock(&q->lock);
+    struct timer* timer;
+    while ((timer = tailq_first(&q->timers)) && timer->deadline <= now) {
+        atomic_store(&timer->running, true, ATOMIC_RELAXED);
+        remove_locked(q, timer);
+        spin_unlock(&q->lock);
+
+        timer->func(timer);
+        atomic_store(&timer->running, false, ATOMIC_RELEASE);
+
+        spin_lock(&q->lock);
+    }
+    program_next(q);
+    spin_unlock(&q->lock);
+}

lib/turnstile.c

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