feat: implement O(1) Priority Bitmap Scheduling Class and Scheduler Core

2 files changed, 473 insertions, 0 deletions

diff --git a/na/kernel/sched_class_o1.cc b/na/kernel/sched_class_o1.cc
new file mode 100644
index 0000000..a8ba5df
--- /dev/null
+++ b/na/kernel/sched_class_o1.cc
@@ -0,0 +1,76 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later */
+/* Copyright (C) 2026 Richard Qin */
+
+/*
+ * O(1) Priority Bitmap Scheduling Class
+ */
+
+#include <lume/sched.h>
+#include <lume/sched_class.h>
+#include <lume/task.h>
+#include <lume/bitops.h>
+#include <kernel/sched_internal.h>
+
+static void o1_init_rq(RunQueue* rq)
+{
+    rq->o1().init();
+}
+
+static void o1_enqueue_task(RunQueue* rq, TaskControlBlock* task)
+{
+    rq->o1().enqueue(task);
+    rq->inc_running();
+}
+
+static void o1_dequeue_task(RunQueue* rq, TaskControlBlock* task)
+{
+    rq->o1().dequeue(task);
+    rq->dec_running();
+}
+
+static TaskControlBlock* o1_pick_next_task(RunQueue* rq)
+{
+    return rq->o1().pick_next();
+}
+
+static void o1_put_prev_task(RunQueue* /*rq*/, TaskControlBlock* /*task*/) {}
+
+int32 priority_to_timeslice(uint8 priority)
+{
+    int32 ts = 20 - static_cast<int32>(priority) * 19 / 31;
+    if (ts < 1) ts = 1;
+    return ts;
+}
+
+static void o1_task_tick(RunQueue* rq, TaskControlBlock* curr)
+{
+    if (curr == rq->idle())
+        return;
+
+    curr->time_slice--;
+    if (curr->time_slice <= 0) {
+        curr->time_slice = priority_to_timeslice(curr->priority);
+        rq->set_need_resched(true);
+    }
+}
+
+static void o1_yield_task(RunQueue* rq, TaskControlBlock* curr)
+{
+    if (curr == rq->idle())
+        return;
+
+    curr->time_slice = 0;
+    rq->set_need_resched(true);
+}
+
+const SchedClass g_sched_class_o1 = {
+    .name = "o1",
+    .next = &g_sched_class_idle,
+    .init_rq = o1_init_rq,
+    .enqueue_task = o1_enqueue_task,
+    .dequeue_task = o1_dequeue_task,
+    .pick_next_task = o1_pick_next_task,
+    .put_prev_task = o1_put_prev_task,
+    .task_tick = o1_task_tick,
+    .yield_task = o1_yield_task,
+};
diff --git a/na/kernel/sched_core.cc b/na/kernel/sched_core.cc
new file mode 100644
index 0000000..85248ef
--- /dev/null
+++ b/na/kernel/sched_core.cc
@@ -0,0 +1,397 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later */
+/* Copyright (C) 2026 Richard Qin */
+
+/*
+ * Scheduler Core
+ *
+ * The core walks sched_class chain to pick the next task. Each class
+ * owns its own runqueue data and policy. This file owns the per-CPU
+ * runqueue array and the scheduling state machine.
+ */
+
+#include <lume/sched.h>
+#include <lume/sched_class.h>
+#include <lume/task.h>
+#include <lume/pmm.h>
+#include <lume/addr.h>
+#include <lume/new.h>
+#include <lume/kprintf.h>
+#include <lume/klog.h>
+#include <arch/cpu.h>
+#include <arch/context.h>
+#include <kernel/sched_internal.h>
+
+/* Assembly trampoline symbol (MD layer) */
+extern "C" void kthread_entry_trampoline();
+
+/* Per-CPU runqueue array */
+static RunQueue g_rq[kMaxCpus];
+
+/* sched_class chain head (highest priority first) */
+static const SchedClass* g_sched_class_head = &g_sched_class_o1;
+
+/* ========================================================================
+ * Idle thread
+ * ======================================================================== */
+
+static void idle_thread_func(void* /*arg*/)
+{
+    while (true) {
+        arch::cpu::intr_on();
+        arch::cpu::halt_until_interrupt();
+    }
+}
+
+/* ========================================================================
+ * Idle class (always returns idle task)
+ * ======================================================================== */
+
+static void idle_init_rq(RunQueue* /*rq*/) {}
+static void idle_enqueue(RunQueue* /*rq*/, TaskControlBlock* /*task*/) {}
+static void idle_dequeue(RunQueue* /*rq*/, TaskControlBlock* /*task*/) {}
+static TaskControlBlock* idle_pick_next(RunQueue* rq) { return rq->idle(); }
+static void idle_put_prev(RunQueue* /*rq*/, TaskControlBlock* /*task*/) {}
+static void idle_task_tick(RunQueue* /*rq*/, TaskControlBlock* /*curr*/) {}
+static void idle_yield(RunQueue* /*rq*/, TaskControlBlock* /*curr*/) {}
+
+const SchedClass g_sched_class_idle = {
+    .name = "idle",
+    .next = nullptr,
+    .init_rq = idle_init_rq,
+    .enqueue_task = idle_enqueue,
+    .dequeue_task = idle_dequeue,
+    .pick_next_task = idle_pick_next,
+    .put_prev_task = idle_put_prev,
+    .task_tick = idle_task_tick,
+    .yield_task = idle_yield,
+};
+
+/* ========================================================================
+ * Internal helpers
+ * ======================================================================== */
+
+static RunQueue* rq_for_cpu(uint32 cpu_id)
+{
+    return &g_rq[cpu_id];
+}
+
+static RunQueue* this_rq()
+{
+    return rq_for_cpu(arch::cpu::id());
+}
+
+static TaskControlBlock* pick_next_task(RunQueue* rq)
+{
+    for (const SchedClass* c = g_sched_class_head; c; c = c->next) {
+        TaskControlBlock* t = c->pick_next_task(rq);
+        if (t) {
+            return t;
+        }
+    }
+    return rq->idle();
+}
+
+static const char* idle_name_for_cpu(uint32 cpu)
+{
+    static char names[kMaxCpus][16];
+    if (cpu >= kMaxCpus) {
+        kernel_panic("sched: cpu_id out of range");
+    }
+
+    if (names[cpu][0] == '\0') {
+        ksnprintf(names[cpu], sizeof(names[cpu]), "idle/%u", cpu);
+    }
+    return names[cpu];
+}
+
+/* ========================================================================
+ * Kernel thread entry — MI-layer function called by MD trampoline
+ * ======================================================================== */
+
+extern "C" void kthread_entry(void (*entry)(void*), void* arg)
+{
+    RunQueue* rq = this_rq();
+    rq->lock().release();
+
+    arch::cpu::intr_on();
+
+    entry(arg);
+
+    kprintf("[sched] kthread exited\n");
+    while (true) {
+        arch::cpu::halt_until_interrupt();
+    }
+}
+
+/* ========================================================================
+ * fork_child_entry — called from ret_from_fork assembly
+ * ======================================================================== */
+extern "C" void userret(arch::TrapFrame* tf);
+
+extern "C" void fork_child_entry(arch::TrapFrame* tf)
+{
+    RunQueue* rq = this_rq();
+    rq->lock().release();
+
+    /* Do NOT enable interrupts here! userret expects them disabled. */
+
+    /* Jump to user space */
+    userret(tf);
+    
+    __builtin_unreachable();
+}
+
+/* ========================================================================
+ * Helper: create a TCB with a kernel stack
+ * ======================================================================== */
+
+static TaskControlBlock* create_tcb(const char* name, uint8 priority,
+                                    uint8 flags, uint8 cpu)
+{
+    auto* task = new TaskControlBlock;
+    if (!task) return nullptr;
+
+    task->context = {};
+
+    Page* stack_page = pmm_alloc_pages(1);
+    if (!stack_page) {
+        delete task;
+        return nullptr;
+    }
+    task->kstack_pa = page_to_pa(stack_page).raw;
+    task->kstack_va = pa_to_va(phys_addr(task->kstack_pa)).raw;
+
+    *reinterpret_cast<uint64*>(task->kstack_va) = kStackCanary;
+
+    task->name = name;
+    task->priority = priority;
+    task->time_slice = priority_to_timeslice(priority);
+    task->state = TaskState::Runnable;
+    task->cpu_id = cpu;
+    task->flags = flags;
+    task->exit_code = 0;
+    task->parent = nullptr;
+    task->mm = nullptr;
+    task->trapframe = nullptr;
+
+    task->run_link.init();
+    task->all_link.init();
+    task->children.init();
+    task->sibling.init();
+    task->wait_chld.init("wait_chld");
+
+    task->sched_class = &g_sched_class_o1;
+    task->sched_entity = &task->run_link;
+
+    return task;
+}
+
+/* ========================================================================
+ * Public API: current_task()
+ * ======================================================================== */
+
+TaskControlBlock* current_task()
+{
+    return this_rq()->current();
+}
+
+/* ========================================================================
+ * Public API: sched_enqueue / sched_dequeue
+ * ======================================================================== */
+
+void sched_enqueue(TaskControlBlock* task)
+{
+    RunQueue* rq = rq_for_cpu(task->cpu_id);
+    LockGuard guard(rq->lock());
+
+    if (!task->sched_class) {
+        task->sched_class = &g_sched_class_o1;
+    }
+    if (!task->sched_entity) {
+        task->sched_entity = &task->run_link;
+    }
+
+    task->sched_class->enqueue_task(rq, task);
+}
+
+void sched_dequeue(TaskControlBlock* task)
+{
+    RunQueue* rq = rq_for_cpu(task->cpu_id);
+    LockGuard guard(rq->lock());
+
+    if (!task->sched_class) {
+        task->sched_class = &g_sched_class_o1;
+    }
+    task->sched_class->dequeue_task(rq, task);
+}
+
+/* ========================================================================
+ * schedule() — main scheduling entry point
+ * ======================================================================== */
+
+void schedule()
+{
+    RunQueue* rq = this_rq();
+    rq->lock().acquire();
+
+    TaskControlBlock* prev = rq->current();
+
+    if (prev && prev->state == TaskState::Running) {
+        prev->state = TaskState::Runnable;
+        prev->sched_class->enqueue_task(rq, prev);
+    } else if (prev) {
+        prev->sched_class->put_prev_task(rq, prev);
+    }
+
+    TaskControlBlock* next = pick_next_task(rq);
+
+    if (prev == next) {
+        /* We picked the same task. It must be Running again. */
+        next->state = TaskState::Running;
+        next->sched_class->dequeue_task(rq, next);
+        rq->clear_need_resched();
+        rq->lock().release();
+        return;
+    }
+
+    if (next != rq->idle()) {
+        next->sched_class->dequeue_task(rq, next);
+    }
+
+    next->state = TaskState::Running;
+    rq->set_current(next);
+    rq->clear_need_resched();
+
+    arch::context_switch(&prev->context, &next->context);
+
+    rq->lock().release();
+}
+
+/* ========================================================================
+ * sched_tick() — called from timer interrupt (MI trap layer)
+ * ======================================================================== */
+
+void sched_tick()
+{
+    RunQueue* rq = this_rq();
+    LockGuard guard(rq->lock());
+
+    TaskControlBlock* curr = rq->current();
+    if (!curr) return;
+
+    curr->sched_class->task_tick(rq, curr);
+}
+
+/* ========================================================================
+ * sched_check_preempt() — called before returning to user mode
+ * ======================================================================== */
+
+void sched_check_preempt()
+{
+    RunQueue* rq = this_rq();
+    if (rq->need_resched()) {
+        schedule();
+    }
+}
+
+/* ========================================================================
+ * sched_yield() — voluntary yield
+ * ======================================================================== */
+
+void sched_yield()
+{
+    RunQueue* rq = this_rq();
+    {
+        LockGuard guard(rq->lock());
+        TaskControlBlock* curr = rq->current();
+        if (curr && curr != rq->idle()) {
+            curr->sched_class->yield_task(rq, curr);
+            rq->set_need_resched(true);
+        }
+    }
+    schedule();
+}
+
+/* ========================================================================
+ * sched_init() — BSP initialization
+ * ======================================================================== */
+
+void sched_init()
+{
+    for (int i = 0; i < kMaxCpus; i++) {
+        RunQueue* rq = &g_rq[i];
+        rq->init(static_cast<uint32>(i));
+
+        g_sched_class_o1.init_rq(rq);
+        g_sched_class_idle.init_rq(rq);
+    }
+
+    g_sched_class_head = &g_sched_class_o1;
+
+    TaskControlBlock* idle0 = create_tcb(idle_name_for_cpu(0), kIdlePriority,
+                                         TF_KTHREAD | TF_IDLE, 0);
+    if (!idle0) kernel_panic("sched_init: failed to create idle/0");
+
+    idle0->pid = 0;
+    idle0->tgid = 0;
+    idle0->sched_class = &g_sched_class_idle;
+    idle0->sched_entity = nullptr;
+
+    uint64 stack_top = idle0->kstack_va + kKernelStackSize;
+    idle0->context = {};
+    idle0->context.ra = reinterpret_cast<uint64>(kthread_entry_trampoline);
+    idle0->context.sp = stack_top;
+    idle0->context.s[0] = reinterpret_cast<uint64>(idle_thread_func);
+    idle0->context.s[1] = 0;
+
+    g_rq[0].set_idle(idle0);
+
+    TaskControlBlock* bsp_task = create_tcb("kernel_main", 0, TF_KTHREAD, 0);
+    if (!bsp_task) kernel_panic("sched_init: failed to create bsp_task");
+
+    bsp_task->pid = -1;
+    bsp_task->tgid = -1;
+    bsp_task->state = TaskState::Running;
+
+    g_rq[0].set_current(bsp_task);
+}
+
+/* ========================================================================
+ * sched_init_ap() — AP per-CPU initialization
+ * ======================================================================== */
+
+void sched_init_ap()
+{
+    uint64 cpu = arch::cpu::id();
+
+    const char* name = idle_name_for_cpu(static_cast<uint32>(cpu));
+
+    TaskControlBlock* idle = create_tcb(name, kIdlePriority,
+                                        TF_KTHREAD | TF_IDLE,
+                                        static_cast<uint8>(cpu));
+    if (!idle) kernel_panic("sched_init_ap: failed to create idle");
+
+    idle->pid = 0;
+    idle->tgid = 0;
+    idle->sched_class = &g_sched_class_idle;
+    idle->sched_entity = nullptr;
+
+    uint64 stack_top = idle->kstack_va + kKernelStackSize;
+    idle->context = {};
+    idle->context.ra = reinterpret_cast<uint64>(kthread_entry_trampoline);
+    idle->context.sp = stack_top;
+    idle->context.s[0] = reinterpret_cast<uint64>(idle_thread_func);
+    idle->context.s[1] = 0;
+
+    g_rq[cpu].set_idle(idle);
+
+    TaskControlBlock* ap_task = create_tcb("ap_boot", 0, TF_KTHREAD,
+                                          static_cast<uint8>(cpu));
+    if (!ap_task) kernel_panic("sched_init_ap: failed to create ap_task");
+
+    ap_task->pid = -1;
+    ap_task->tgid = -1;
+    ap_task->state = TaskState::Running;
+
+    g_rq[cpu].set_current(ap_task);
+}