Add turnstiles, more sched fixes

arch/x86_64/idt.c

@@ -6,6 +6,7 @@
 #include <zinnia/pmap.h>
 #include <zinnia/print.h>
 #include <zinnia/sched.h>
+#include <zinnia/timer.h>
 #include <zinnia/usercopy.h>
 #include <zinnia/utils.h>
 #include <zinnia/vma.h>
@@ -98,6 +99,15 @@ void interrupt_handler(struct arch_context* ctx) {
         lapic_eoi(&arch_cpu->lapic);
         pmap_shootdown_handle();
         break;
+    case IDT_TIMER:
+        lapic_eoi(&arch_cpu->lapic);
+        timer_interrupt();
+        if (cpu->sched.need_resched) {
+            arch_sched_preempt_disable();
+            if (arch_sched_preempt_enable())
+                sched_reschedule(&cpu->sched);
+        }
+        break;
     default:
         if (isr < 0x20)
             panic("Unhandled exception %lu at %#lx, error: %#lx\n", isr, ctx->rip, ctx->error);

include/zinnia/mutex.h

@@ -4,7 +4,10 @@
 
 #include <stdint.h>
 
-typedef uintptr_t mutex_t;
+typedef _Atomic(uintptr_t) mutex_t;
 
 void mutex_lock(mutex_t* mutex);
+
+bool mutex_trylock(mutex_t* mutex);
+
 void mutex_unlock(mutex_t* mutex);

include/zinnia/sched.h

@@ -6,8 +6,10 @@
 #include <zinnia/errno.h>
 #include <zinnia/init.h>
 #include <zinnia/initgraph.h>
+#include <zinnia/spin.h>
 #include <zinnia/stack.h>
 #include <zinnia/tailq.h>
+#include <zinnia/timer.h>
 #include <zinnia/types.h>
 #include <bits/sched.h>
 #include <stddef.h>
@@ -23,8 +25,24 @@ enum task_state {
     TASK_STATE_DEAD,    // Task is killed.
 };
 
-#define TASK_NAME_MAX     128
+enum sched_prio : uint8_t {
-#define KERNEL_STACK_SIZE 0x4000
+    SCHED_PRIO_MIN = 0,
+
+    SCHED_PRIO_DEFAULT = 32,
+    SCHED_PRIO_WORK_HIGH = 48,
+    SCHED_PRIO_IRQ = 56,
+
+    SCHED_PRIO_MAX = 63,
+    SCHED_PRIO_COUNT = SCHED_PRIO_MAX + 1,
+};
+
+constexpr size_t SCHED_TIME_SLICE_NS = 10'000'000;
+constexpr uint8_t SCHED_FEEDBACK_DEPTH = 8;
+constexpr uint64_t SCHED_BOOST_INTERVAL_NS = 1'000'000'000;
+
+struct turnstile;
+
+constexpr size_t TASK_NAME_MAX = 128;
 
 // A task is the smallest unit of the scheduler.
 struct task {
@@ -35,11 +53,28 @@ struct task {
     struct arch_task_context context;
     struct usercopy_region* usercopy_region;
     struct stack kernel_stack;
-    uintptr_t user_stack;
+    uintptr_t user_stack; // Saved stack pointer from userspace.
     size_t time_slice;
-    int8_t priority;
+    uint8_t penalty;
+    uint64_t feedback_epoch;
     atomic_bool wake_pending;
 
+    struct spinlock lock;
+
+    atomic_bool on_cpu;                // Whether a CPU is still running on this task's context.
+    uint8_t base_prio;                 // Set by sched_set_priority().
+    uint8_t lent_prio;                 // Highest priority lent by tasks waiting on locks this task owns.
+    _Atomic(enum sched_prio) priority; // The effective priority of this 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.
+    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;
+    atomic_bool ts_woken;
+
     tailq_link(struct task) next;
 };
 
@@ -48,7 +83,24 @@ struct scheduler {
     struct task* current;
     struct task* idle_task;
     size_t preempt_level;
-    tailq_head(struct task) run_queue;
+
+    struct spinlock lock;
+
+    // Bit N is set if `run_queue[N]` is not empty.
+    uint64_t ready_mask;
+    tailq_head(struct task) run_queue[SCHED_PRIO_COUNT];
+
+    // Current priority, or negative if the CPU is idle.
+    _Atomic(int32_t) current_priority;
+
+    atomic_bool kick_pending;
+
+    struct scheduler* next_cpu;
+
+    struct timer preempt_timer;
+    bool need_resched;
+    uint64_t slice_start;
+    uint64_t boost_epoch;
 };
 
 typedef void (*task_fn_t)(uintptr_t arg0);
@@ -69,6 +121,16 @@ 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.
+void sched_wake_locked(struct task* task);
+
+[[noreturn]]
+void sched_exit();
+
+// Sets the base priority of `task`.
+void sched_set_priority(struct task* task, uint8_t priority);
+
 [[noreturn]]
 void sched_to_user(uintptr_t ctx);
 [[noreturn]]

include/zinnia/timer.h

@@ -0,0 +1,24 @@
+// SPDX-License-Identifier: ISC
+
+#pragma once
+
+#include <zinnia/atomic.h>
+#include <zinnia/tailq.h>
+#include <stdint.h>
+
+struct timer_queue;
+
+struct timer {
+    uint64_t deadline;
+    void (*func)(struct timer* timer);
+    _Atomic(struct timer_queue*) queue;
+    atomic_bool running;
+    tailq_link(struct timer) link;
+};
+
+void timer_init(struct timer* timer, void (*func)(struct timer* timer));
+void timer_arm(struct timer* timer, uint64_t deadline_ns);
+bool timer_cancel(struct timer* timer);
+
+void timer_cpu_init();
+void timer_interrupt();

include/zinnia/turnstile.h

@@ -0,0 +1,58 @@
+// SPDX-License-Identifier: ISC
+
+#pragma once
+
+#include <zinnia/tailq.h>
+#include <stdint.h>
+
+struct task;
+struct turnstile_chain;
+
+enum turnstile_queue {
+    TURNSTILE_EXCLUSIVE,
+    TURNSTILE_SHARED,
+
+    TURNSTILE_NUM_QUEUES,
+};
+
+struct turnstile {
+    // The lock object that this turnstile is assigned to.
+    // If not in use, it's set to `nullptr`.
+    const void* obj;
+
+    // The task that inherits the priority from the waiters.
+    struct task* owner;
+
+    // Priority of the most urgent waiter in any queue.
+    _Atomic(uint8_t) top_priority;
+
+    tailq_head(struct task) waiters[TURNSTILE_NUM_QUEUES];
+
+    // Spare turnstiles donated by waiters.
+    struct turnstile* spare;
+
+    // Next active turnstile in the hash chain.
+    struct turnstile* chain_next;
+
+    // Link for `task::pi_lenders`.
+    tailq_link(struct turnstile) lender_link;
+};
+
+struct turnstile* turnstile_alloc();
+void turnstile_free(struct turnstile* ts);
+
+struct turnstile_chain* turnstile_lock(const void* obj);
+void turnstile_unlock(struct turnstile_chain* chain);
+
+void turnstile_block(struct turnstile_chain* chain, const void* obj, struct task* owner, enum turnstile_queue queue);
+
+struct task* turnstile_first(
+    struct turnstile_chain* chain,
+    const void* obj,
+    enum turnstile_queue queue,
+    bool* have_more
+);
+
+void turnstile_wake(struct turnstile_chain* chain, const void* obj, struct task* task, struct task* new_owner);
+
+void turnstile_reprioritize(struct task* task);

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

sched/core.c

@@ -1 +1,106 @@
+// SPDX-License-Identifier: ISC
+
+#include <zinnia/alloc.h>
+#include <zinnia/assert.h>
+#include <zinnia/atomic.h>
+#include <zinnia/errno.h>
 #include <zinnia/sched.h>
+#include <zinnia/spin.h>
+#include <zinnia/stack.h>
+#include <zinnia/tailq.h>
+#include <zinnia/turnstile.h>
+#include <zinnia/vma.h>
+#include <arch/sched.h>
+#include <stdint.h>
+#include <string.h>
+#include "internal.h"
+
+void sched_to_user(uintptr_t info_ptr) {
+    struct task_start_info* info = (struct task_start_info*)info_ptr;
+    uintptr_t ip = info->ip;
+    uintptr_t sp = info->sp;
+    uintptr_t arg = info->arg;
+    mem_free(info);
+    arch_sched_jump_to_user(ip, sp, arg);
+}
+
+void sched_to_user_context(uintptr_t ctx) {
+    arch_sched_jump_to_context((struct arch_context*)ctx);
+}
+
+static void task_init_common(struct task* task, const char* name, struct vma* space) {
+    task->space = space;
+    spin_init(&task->lock);
+    strncpy(task->name, name, sizeof(task->name));
+    task->base_prio = SCHED_PRIO_DEFAULT;
+    task->time_slice = SCHED_TIME_SLICE_NS;
+    task->penalty = 0;
+    task->feedback_epoch = 0;
+    atomic_store(&task->priority, SCHED_PRIO_DEFAULT, ATOMIC_RELAXED);
+    tailq_init(&task->pi_lenders);
+}
+
+errno_t task_create_current(const char* name, struct vma* space, struct task** out) {
+    struct task* task = mem_alloc(sizeof(struct task), 0);
+    if (task == nullptr)
+        return ENOMEM;
+
+    task->turnstile = turnstile_alloc();
+    if (task->turnstile == nullptr) {
+        mem_free(task);
+        return ENOMEM;
+    }
+
+    task_init_common(task, name, space);
+    *out = task;
+    return 0;
+}
+
+constexpr size_t KERNEL_STACK_SIZE = 16 * 1024;
+
+errno_t task_create(const char* name, struct vma* space, task_fn_t entry, uintptr_t arg0, struct task** out) {
+    if (!out)
+        return EINVAL;
+
+    struct task* new_task = mem_alloc(sizeof(struct task), 0);
+    if (new_task == nullptr)
+        return ENOMEM;
+
+    static atomic_size_t last_tid = 0;
+    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);
+        return ENOMEM;
+    }
+
+    // Allocate a kernel stack.
+    errno_t s = stack_init(&new_task->kernel_stack, KERNEL_STACK_SIZE);
+    if (s != 0) {
+        turnstile_free(new_task->turnstile);
+        mem_free(new_task);
+        return s;
+    }
+    new_task->user_stack = 0;
+
+    // Initialize the arch-dependent parts.
+    s = arch_task_init(&new_task->context, entry, arg0, kernel_stack_top(&new_task->kernel_stack), true);
+    if (s != 0) {
+        stack_destroy(&new_task->kernel_stack);
+        turnstile_free(new_task->turnstile);
+        mem_free(new_task);
+        return s;
+    }
+
+    task_init_common(new_task, name, space);
+    *out = new_task;
+    return 0;
+}
+
+void task_destroy(struct task* task) {
+    assert(tailq_empty(&task->pi_lenders), "Task %zu exited while owning contended locks", task->id);
+    turnstile_free(task->turnstile);
+    arch_task_fini(&task->context);
+    stack_destroy(&task->kernel_stack);
+    mem_free(task);
+}

sched/internal.h

@@ -0,0 +1,9 @@
+// SPDX-License-Identifier: ISC
+
+#pragma once
+
+#include <zinnia/sched.h>
+
+errno_t task_create_current(const char* name, struct vma* space, struct task** out);
+
+bool sched_update_priority_locked(struct task* task);

sched/sched_dumb.c

@@ -1,7 +1,8 @@
-#include <zinnia/alloc.h>
+// SPDX-License-Identifier: ISC
+
 #include <zinnia/assert.h>
 #include <zinnia/atomic.h>
-#include <zinnia/errno.h>
+#include <zinnia/clock.h>
 #include <zinnia/irq.h>
 #include <zinnia/list.h>
 #include <zinnia/percpu.h>
@@ -11,9 +12,13 @@
 #include <zinnia/spin.h>
 #include <zinnia/stack.h>
 #include <zinnia/tailq.h>
+#include <zinnia/turnstile.h>
+#include <zinnia/utils.h>
 #include <zinnia/vma.h>
+#include <arch/cpu.h>
 #include <arch/sched.h>
 #include <stdint.h>
+#include "internal.h"
 
 static void idle_func(uintptr_t) {
     while (1) {
@@ -22,61 +27,194 @@ static void idle_func(uintptr_t) {
     }
 }
 
-void sched_to_user(uintptr_t info_ptr) {
+constexpr int32_t SCHED_CPU_IDLE = -1;
-    struct task_start_info* info = (struct task_start_info*)info_ptr;
+constexpr int32_t SCHED_CPU_CLAIMED = -2;
-    uintptr_t ip = info->ip;
+
-    uintptr_t sp = info->sp;
+static _Atomic(struct scheduler*) all_cpus;
-    uintptr_t arg = info->arg;
+
-    mem_free(info);
+static struct percpu* cpu_of(struct scheduler* sched) {
-    arch_sched_jump_to_user(ip, sp, arg);
+    return containerof(sched, struct percpu, sched);
 }
 
-void sched_to_user_context(uintptr_t ctx) {
+static void preempt_tick(struct timer* timer) {
-    arch_sched_jump_to_context((struct arch_context*)ctx);
+    struct scheduler* sched = containerof(timer, struct scheduler, preempt_timer);
+    sched->need_resched = true;
+}
+
+static void start_time_slice(struct scheduler* sched, struct task* task, uint64_t now) {
+    sched->need_resched = false;
+    sched->slice_start = now;
+    if (task == sched->idle_task)
+        timer_cancel(&sched->preempt_timer);
+    else
+        timer_arm(&sched->preempt_timer, now + task->time_slice);
 }
 
 void sched_init(struct scheduler* sched) {
+    spin_init(&sched->lock);
+    atomic_store(&sched->kick_pending, false, ATOMIC_RELAXED);
+    for (size_t i = 0; i < SCHED_PRIO_COUNT; i++)
+        tailq_init(&sched->run_queue[i]);
+
     struct task* idle_task;
     assert(task_create("idle", &kernel_vma, idle_func, 0, &idle_task) == 0, "Unable to create idle task!");
+    idle_task->last_rq = sched;
     sched->idle_task = idle_task;
 
-    struct task* boot_task = mem_alloc(sizeof(struct task), 0);
+    struct task* boot_task;
-    assert(boot_task != nullptr, "Unable to create boot task!");
+    assert(task_create_current("boot", &kernel_vma, &boot_task) == 0, "Unable to create boot task!");
-    strncpy(boot_task->name, "boot", sizeof(boot_task->name));
+    boot_task->state = TASK_STATE_RUNNING;
-    boot_task->space = &kernel_vma;
+    boot_task->last_rq = sched;
+    atomic_store(&boot_task->on_cpu, true, ATOMIC_RELAXED);
     sched->current = boot_task;
+    atomic_store(&sched->current_priority, SCHED_PRIO_DEFAULT, ATOMIC_RELAXED);
 
-    tailq_init(&sched->run_queue);
+    sched->next_cpu = atomic_load(&all_cpus, ATOMIC_RELAXED);
+    while (!atomic_compare_exchange_weak(&all_cpus, &sched->next_cpu, sched, ATOMIC_RELEASE, ATOMIC_RELAXED))
+        ;
+
+    timer_cpu_init();
+    timer_init(&sched->preempt_timer, preempt_tick);
+    sched->boost_epoch = 0;
+    start_time_slice(sched, boot_task, clock_get_elapsed_ns());
 
     log("sched: Scheduler initialized\n");
 }
 
-void sched_add_task(struct scheduler* sched, struct task* task) {
+// Asks the CPU of `sched` to reschedule as soon as it has interrupts enabled.
-    assert(task != nullptr, "Tried to add NULL to the run queue");
+static void kick(struct scheduler* sched) {
-    task->state = TASK_STATE_READY;
+    if (!atomic_exchange(&sched->kick_pending, true, ATOMIC_RELAXED))
-    tailq_insert_tail(&sched->run_queue, task, next);
+        arch_sched_kick(cpu_of(sched));
 }
 
-static struct task* next_task(struct scheduler* sched) {
+static void rq_insert(struct scheduler* sched, struct task* task) {
-    struct task* t = tailq_first(&sched->run_queue);
+    uint8_t prio = atomic_load(&task->priority, ATOMIC_RELAXED);
+    tailq_insert_tail(&sched->run_queue[prio], task, next);
+    sched->ready_mask |= (uint64_t)1 << prio;
+}
 
-    if (t == nullptr)
+static void rq_remove(struct scheduler* sched, struct task* task) {
+    uint8_t prio = atomic_load(&task->priority, ATOMIC_RELAXED);
+    tailq_remove(&sched->run_queue[prio], task, next);
+    if (tailq_empty(&sched->run_queue[prio]))
+        sched->ready_mask &= ~((uint64_t)1 << prio);
+}
+
+// Returns -1 if no task is ready.
+static int rq_top_priority(struct scheduler* sched) {
+    return sched->ready_mask ? 63 - __builtin_clzll(sched->ready_mask) : -1;
+}
+
+static uint8_t effective_prio(struct task* task) {
+    uint8_t prio = task->base_prio;
+    if (prio < SCHED_PRIO_WORK_HIGH)
+        prio -= MIN(task->penalty, prio);
+    return MAX(prio, task->lent_prio);
+}
+
+static size_t quantum(struct task* task) {
+    return SCHED_TIME_SLICE_NS * (task->penalty + 1);
+}
+
+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);
+}
+
+static void boost(struct scheduler* sched) {
+    for (uint8_t prio = 0; prio < SCHED_PRIO_WORK_HIGH; prio++) {
+        struct task* t = tailq_first(&sched->run_queue[prio]);
+        while (t) {
+            struct task* after = tailq_next(t, next);
+            if (t->penalty && t->feedback_epoch < sched->boost_epoch) {
+                rq_remove(sched, t);
+                feedback_reset(sched, t);
+                rq_insert(sched, t);
+            }
+            t = after;
+        }
+    }
+}
+
+static void account(struct scheduler* sched, struct task* task, uint64_t now) {
+    if (task->feedback_epoch < sched->boost_epoch)
+        feedback_reset(sched, task);
+
+    uint64_t ran = now - sched->slice_start;
+    sched->slice_start = now;
+    if (ran < task->time_slice) {
+        task->time_slice -= ran;
+        return;
+    }
+
+    if (task->base_prio < SCHED_PRIO_WORK_HIGH && task->penalty < SCHED_FEEDBACK_DEPTH)
+        task->penalty++;
+    task->time_slice = quantum(task);
+    task->feedback_epoch = sched->boost_epoch;
+    atomic_store(&task->priority, effective_prio(task), ATOMIC_RELAXED);
+    sched->need_resched = true;
+}
+
+// Queues `task` on `sched` and preempts its CPU if `task` has a higher prio than what it's running.
+static void enqueue(struct scheduler* sched, struct task* task) {
+    spin_lock(&sched->lock);
+    task->state = TASK_STATE_READY;
+    task->last_rq = sched;
+    if (task->feedback_epoch < sched->boost_epoch)
+        feedback_reset(sched, task);
+    rq_insert(sched, task);
+    bool preempt = atomic_load(&task->priority, ATOMIC_RELAXED) > atomic_load(&sched->current_priority, ATOMIC_RELAXED);
+    spin_unlock(&sched->lock);
+
+    if (preempt)
+        kick(sched);
+}
+
+void sched_add_task(struct scheduler* sched, struct task* task) {
+    assert(task != nullptr, "Tried to add NULL to the run queue");
+    enqueue(sched, task);
+}
+
+static struct task* next_task(struct scheduler* sched, uint64_t now) {
+    spin_lock(&sched->lock);
+    uint64_t epoch = now / SCHED_BOOST_INTERVAL_NS;
+    if (epoch > sched->boost_epoch) {
+        sched->boost_epoch = epoch;
+        boost(sched);
+    }
+
+    int prio = rq_top_priority(sched);
+    if (prio < 0) {
+        atomic_store(&sched->current_priority, SCHED_CPU_IDLE, ATOMIC_RELAXED);
+        spin_unlock(&sched->lock);
         return sched->idle_task;
+    }
 
-    tailq_remove(&sched->run_queue, t, next);
+    struct task* t = tailq_first(&sched->run_queue[prio]);
+    rq_remove(sched, t);
     t->state = TASK_STATE_RUNNING;
+    atomic_store(&sched->current_priority, prio, ATOMIC_RELAXED);
+    spin_unlock(&sched->lock);
 
     return t;
 }
 
 static void finish_switch(struct task* prev) {
+    // The context of `prev` is fully saved now, so other CPUs can switch to it.
+    atomic_store(&prev->on_cpu, false, ATOMIC_RELEASE);
     if (prev->state == TASK_STATE_DEAD)
         task_destroy(prev);
 }
 
-static void do_reschedule(struct scheduler* sched) {
+static void do_reschedule(struct scheduler* sched, uint64_t now) {
+    atomic_store(&sched->kick_pending, false, ATOMIC_RELAXED);
     struct task* from = sched->current;
-    struct task* to = next_task(sched);
+    struct task* to = next_task(sched, now);
+
+    if (from != to || sched->need_resched)
+        start_time_slice(sched, to, now);
 
     // We'd end up on the same task, skip switching altogether.
     if (from == to) {
@@ -84,6 +222,11 @@ static void do_reschedule(struct scheduler* sched) {
         return;
     }
 
+    // A task that was woken by another CPU might not have finished switching away on its old CPU yet.
+    while (atomic_load(&to->on_cpu, ATOMIC_ACQUIRE))
+        arch_cpu_pause();
+    atomic_store(&to->on_cpu, true, ATOMIC_RELAXED);
+
     sched->current = to;
     pmap_set(&to->space->pmap);
 
@@ -101,72 +244,126 @@ static void do_reschedule(struct scheduler* sched) {
 
 void sched_reschedule(struct scheduler* sched) {
     irq_lock();
-    if (sched->current != sched->idle_task && sched->current->state != TASK_STATE_DEAD)
+    struct task* current = sched->current;
-        sched_add_task(sched, sched->current);
+    uint64_t now = clock_get_elapsed_ns();
-    do_reschedule(sched);
+    if (current != sched->idle_task && current->state == TASK_STATE_RUNNING) {
+        spin_lock(&sched->lock);
+        account(sched, current, now);
+        current->state = TASK_STATE_READY;
+        rq_insert(sched, current);
+        spin_unlock(&sched->lock);
+    }
+    do_reschedule(sched, now);
 }
 
 void sched_yield(struct scheduler* sched) {
     irq_lock();
-    do_reschedule(sched);
+    do_reschedule(sched, clock_get_elapsed_ns());
+}
+
+void sched_exit() {
+    irq_lock();
+    struct scheduler* sched = &percpu_get()->sched;
+    sched->current->state = TASK_STATE_DEAD;
+    do_reschedule(sched, clock_get_elapsed_ns());
+    unreachable();
 }
 
 void sched_block(struct scheduler* sched) {
     irq_lock();
-    if (atomic_exchange(&sched->current->wake_pending, false, ATOMIC_ACQ_REL)) {
+    struct task* current = sched->current;
+
+    spin_lock(&current->lock);
+    if (atomic_exchange(&current->wake_pending, false, ATOMIC_ACQ_REL)) {
+        spin_unlock(&current->lock);
         irq_unlock();
         return;
     }
-    sched->current->state = TASK_STATE_BLOCKED;
+
-    do_reschedule(sched);
+    uint64_t now = clock_get_elapsed_ns();
+    spin_lock(&sched->lock);
+    account(sched, current, now);
+    spin_unlock(&sched->lock);
+
+    current->state = TASK_STATE_BLOCKED;
+    spin_unlock(&current->lock);
+
+    do_reschedule(sched, now);
+}
+
+static struct scheduler* select_cpu(struct task* task) {
+    struct scheduler* last = task->last_rq ? task->last_rq : &percpu_get()->sched;
+    int prio = atomic_load(&task->priority, ATOMIC_RELAXED);
+    if (atomic_load(&last->current_priority, ATOMIC_RELAXED) < prio)
+        return last;
+
+    // Claim an idle CPU.
+    struct scheduler* s;
+    for (s = atomic_load(&all_cpus, ATOMIC_ACQUIRE); s; s = s->next_cpu) {
+        int idle = SCHED_CPU_IDLE;
+        if (atomic_compare_exchange(&s->current_priority, &idle, SCHED_CPU_CLAIMED, ATOMIC_RELAXED, ATOMIC_RELAXED))
+            return s;
+    }
+    return last;
+}
+
+void sched_wake_locked(struct task* task) {
+    if (task->state == TASK_STATE_BLOCKED)
+        enqueue(select_cpu(task), task);
+    else
+        atomic_store(&task->wake_pending, true, ATOMIC_RELEASE);
 }
 
 void sched_wake(struct task* task) {
-    irq_lock();
+    spin_lock(&task->lock);
-    if (task->state == TASK_STATE_BLOCKED)
+    sched_wake_locked(task);
-        sched_add_task(&percpu_get()->sched, task);
+    spin_unlock(&task->lock);
-    else
-        atomic_store(&task->wake_pending, true, ATOMIC_RELEASE);
-    irq_unlock();
 }
 
-errno_t task_create(const char* name, struct vma* space, task_fn_t entry, uintptr_t arg0, struct task** out) {
+bool sched_update_priority_locked(struct task* task) {
-    if (!out)
+    struct scheduler* sched = task->last_rq;
-        return EINVAL;
+    if (sched == nullptr) {
-
+        // The task was never scheduled, so it isn't queued anywhere.
-    struct task* new_task = mem_alloc(sizeof(struct task), 0);
+        uint8_t new_prio = effective_prio(task);
-    if (new_task == nullptr)
+        return atomic_exchange(&task->priority, new_prio, ATOMIC_RELAXED) != new_prio;
-        return ENOMEM;
-
-    static atomic_size_t last_tid = 0;
-    new_task->id = atomic_fetch_add(&last_tid, 1, ATOMIC_ACQUIRE);
-    new_task->space = space;
-    strncpy(new_task->name, name, sizeof(new_task->name));
-
-    // Allocate a kernel stack.
-    errno_t s = stack_init(&new_task->kernel_stack, KERNEL_STACK_SIZE);
-    if (s != 0) {
-        mem_free(new_task);
-        return s;
-    }
-    new_task->user_stack = 0;
-
-    // Initialize the arch-dependent parts.
-    s = arch_task_init(&new_task->context, entry, arg0, kernel_stack_top(&new_task->kernel_stack), true);
-    if (s != 0) {
-        stack_destroy(&new_task->kernel_stack);
-        mem_free(new_task);
-        return s;
     }
 
-    *out = new_task;
+    bool preempt = false;
-    return 0;
+    spin_lock(&sched->lock);
+    uint8_t new_prio = effective_prio(task);
+    if (new_prio == atomic_load(&task->priority, ATOMIC_RELAXED)) {
+        spin_unlock(&sched->lock);
+        return false;
+    }
+
+    if (task->state == TASK_STATE_READY) {
+        rq_remove(sched, task);
+        atomic_store(&task->priority, new_prio, ATOMIC_RELAXED);
+        rq_insert(sched, task);
+        preempt = new_prio > atomic_load(&sched->current_priority, ATOMIC_RELAXED);
+    } else {
+        atomic_store(&task->priority, new_prio, ATOMIC_RELAXED);
+        if (task->state == TASK_STATE_RUNNING) {
+            atomic_store(&sched->current_priority, new_prio, ATOMIC_RELAXED);
+            preempt = rq_top_priority(sched) > new_prio;
+        }
+    }
+    spin_unlock(&sched->lock);
+
+    if (preempt)
+        kick(sched);
+    return true;
 }
 
-void task_destroy(struct task* task) {
+void sched_set_priority(struct task* task, uint8_t priority) {
-    arch_task_fini(&task->context);
+    assert(priority <= SCHED_PRIO_MAX, "Priority %u is out of range", priority);
-    stack_destroy(&task->kernel_stack);
+    spin_lock(&task->lock);
-    mem_free(task);
+    task->base_prio = priority;
+    bool changed = sched_update_priority_locked(task);
+    spin_unlock(&task->lock);
+
+    if (changed)
+        turnstile_reprioritize(task);
 }
 
 void task_entry(task_fn_t entry, uintptr_t arg0, struct task* prev) {
@@ -179,11 +376,7 @@ void task_entry(task_fn_t entry, uintptr_t arg0, struct task* prev) {
     irq_unlock();
 
     entry(arg0);
-
+    sched_exit();
-    struct scheduler* sched = &percpu_get()->sched;
-    sched->current->state = TASK_STATE_DEAD;
-    sched_yield(sched);
-    panic("unreachable");
 }
 
 [[__init]]