Add turnstiles, more sched fixes

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) {
-    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);
+constexpr int32_t SCHED_CPU_IDLE = -1;
+constexpr int32_t SCHED_CPU_CLAIMED = -2;
+
+static _Atomic(struct scheduler*) all_cpus;
+
+static struct percpu* cpu_of(struct scheduler* sched) {
+    return containerof(sched, struct percpu, sched);
 }
 
-void sched_to_user_context(uintptr_t ctx) {
-    arch_sched_jump_to_context((struct arch_context*)ctx);
+static void preempt_tick(struct timer* timer) {
+    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);
-    assert(boot_task != nullptr, "Unable to create boot task!");
-    strncpy(boot_task->name, "boot", sizeof(boot_task->name));
-    boot_task->space = &kernel_vma;
+    struct task* boot_task;
+    assert(task_create_current("boot", &kernel_vma, &boot_task) == 0, "Unable to create boot task!");
+    boot_task->state = TASK_STATE_RUNNING;
+    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) {
-    assert(task != nullptr, "Tried to add NULL to the run queue");
-    task->state = TASK_STATE_READY;
-    tailq_insert_tail(&sched->run_queue, task, next);
+// Asks the CPU of `sched` to reschedule as soon as it has interrupts enabled.
+static void kick(struct scheduler* sched) {
+    if (!atomic_exchange(&sched->kick_pending, true, ATOMIC_RELAXED))
+        arch_sched_kick(cpu_of(sched));
 }
 
-static struct task* next_task(struct scheduler* sched) {
-    struct task* t = tailq_first(&sched->run_queue);
+static void rq_insert(struct scheduler* sched, struct task* task) {
+    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)
-        sched_add_task(sched, sched->current);
-    do_reschedule(sched);
+    struct task* current = sched->current;
+    uint64_t now = clock_get_elapsed_ns();
+    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();
-    if (task->state == TASK_STATE_BLOCKED)
-        sched_add_task(&percpu_get()->sched, task);
-    else
-        atomic_store(&task->wake_pending, true, ATOMIC_RELEASE);
-    irq_unlock();
+    spin_lock(&task->lock);
+    sched_wake_locked(task);
+    spin_unlock(&task->lock);
 }
 
-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->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;
+bool sched_update_priority_locked(struct task* task) {
+    struct scheduler* sched = task->last_rq;
+    if (sched == nullptr) {
+        // The task was never scheduled, so it isn't queued anywhere.
+        uint8_t new_prio = effective_prio(task);
+        return atomic_exchange(&task->priority, new_prio, ATOMIC_RELAXED) != new_prio;
     }
 
-    *out = new_task;
-    return 0;
+    bool preempt = false;
+    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) {
-    arch_task_fini(&task->context);
-    stack_destroy(&task->kernel_stack);
-    mem_free(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->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);
-
-    struct scheduler* sched = &percpu_get()->sched;
-    sched->current->state = TASK_STATE_DEAD;
-    sched_yield(sched);
-    panic("unreachable");
+    sched_exit();
 }
 
 [[__init]]