feat: implement task management for kernel threads

1 files changed, 399 insertions, 0 deletions

diff --git a/na/kernel/task.cc b/na/kernel/task.cc
new file mode 100644
index 0000000..9f09d62
--- /dev/null
+++ b/na/kernel/task.cc
@@ -0,0 +1,399 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later */
+/* Copyright (C) 2026 Richard Qin */
+
+/*
+ * Task Management — create, exit, destroy, kthread
+ *
+ * Implements task lifecycle management: kernel thread creation,
+ * task exit (zombie transition), and resource reclamation.
+ *
+ * Reference: docs/specs/task.md §4-§6
+ */
+
+#include <lume/task.h>
+#include <lume/sched.h>
+#include <lume/pmm.h>
+#include <lume/addr.h>
+#include <lume/new.h>
+#include <lume/kprintf.h>
+#include <lume/klog.h>
+#include <lume/vmm.h>
+#include <lume/errno.h>
+#include <lume/pid.h>
+#include <arch/cpu.h>
+#include <arch/context.h>
+#include <kernel/sync/spinlock.h>
+
+/* ======================================================================
+ * Forward declarations
+ * ====================================================================== */
+
+extern "C" void kthread_entry_trampoline();
+extern "C" void ret_from_fork();
+
+/* ======================================================================
+ * Global task list
+ * ====================================================================== */
+
+list_node g_all_tasks;
+Spinlock g_task_list_lock;
+int32 g_nr_tasks = 0;
+
+/* Special tasks */
+TaskControlBlock* g_kernel_task = nullptr;  /* PID 0 container */
+TaskControlBlock* g_init_task = nullptr;    /* PID 1 (future) */
+
+/* ======================================================================
+ * task_init() — called once from kernel_main before sched_init
+ * ====================================================================== */
+
+void task_init()
+{
+    g_all_tasks.init();
+    g_task_list_lock.init("task_list");
+    g_pid_alloc.init();
+
+    printk(KERN_INFO "init: task_init done (pid_alloc ready)\n");
+}
+
+/* ======================================================================
+ * Internal: allocate TCB + kernel stack
+ * ====================================================================== */
+
+static TaskControlBlock* alloc_task(const char* name, uint8 priority,
+                                    uint8 flags)
+{
+    auto* task = new TaskControlBlock;
+    if (!task) return nullptr;
+
+    task->context = {};
+
+    /* Allocate 8KB kernel stack (order-1 = 2 pages) */
+    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;
+
+    /* Write stack canary at bottom */
+    *reinterpret_cast<uint64*>(task->kstack_va) = kStackCanary;
+
+    /* Basic fields */
+    task->name = name;
+    task->priority = priority;
+    task->time_slice = priority_to_timeslice(priority);
+    task->state = TaskState::Runnable;
+    task->cpu_id = static_cast<uint8>(arch::cpu::id());
+    task->flags = flags;
+    task->exit_code = 0;
+    task->parent = nullptr;
+    task->mm = nullptr;
+    task->trapframe = nullptr;
+
+    /* Init all list nodes */
+    task->run_link.init();
+    task->all_link.init();
+    task->children.init();
+    task->sibling.init();
+    task->wait_chld.init("wait_chld");
+
+    /* Default sched class */
+    extern const SchedClass g_sched_class_o1;
+    task->sched_class = &g_sched_class_o1;
+    task->sched_entity = &task->run_link;
+
+    return task;
+}
+
+/* Add task to global task list */
+static void task_list_add(TaskControlBlock* task)
+{
+    LockGuard guard(g_task_list_lock);
+
+    list_node* head = &g_all_tasks;
+    task->all_link.prev = head->prev;
+    task->all_link.next = head;
+    head->prev->next = &task->all_link;
+    head->prev = &task->all_link;
+
+    g_nr_tasks++;
+}
+
+/* Remove task from global task list */
+static void task_list_remove(TaskControlBlock* task)
+{
+    LockGuard guard(g_task_list_lock);
+
+    task->all_link.prev->next = task->all_link.next;
+    task->all_link.next->prev = task->all_link.prev;
+    task->all_link.init();
+
+    if (g_nr_tasks > 0)
+        g_nr_tasks--;
+}
+
+/* ======================================================================
+ * kthread_create — create a kernel thread
+ * ====================================================================== */
+
+TaskControlBlock* kthread_create(const char* name,
+                                 void (*entry)(void*),
+                                 void* arg,
+                                 uint8 priority)
+{
+    TaskControlBlock* task = alloc_task(name, priority, TF_KTHREAD);
+    if (!task) return nullptr;
+
+    /* Allocate PID */
+    task->pid = g_pid_alloc.alloc();
+    if (task->pid < 0) {
+        /* PID exhausted — clean up */
+            pmm_free_pages(pa_to_page(phys_addr(task->kstack_pa)), 1);
+        delete task;
+        return nullptr;
+    }
+    task->tgid = task->pid;
+
+    /* Set up context for kthread_entry_trampoline */
+    uint64 stack_top = task->kstack_va + kKernelStackSize;
+    task->context = {};
+    task->context.ra = reinterpret_cast<uint64>(kthread_entry_trampoline);
+    task->context.sp = stack_top;
+    task->context.s[0] = reinterpret_cast<uint64>(entry);
+    task->context.s[1] = reinterpret_cast<uint64>(arg);
+
+    /* Parent is kernel_task (PID 0) */
+    task->parent = g_kernel_task;
+    if (g_kernel_task) {
+        LockGuard guard(g_task_list_lock);
+        list_node* head = &g_kernel_task->children;
+        task->sibling.prev = head->prev;
+        task->sibling.next = head;
+        head->prev->next = &task->sibling;
+        head->prev = &task->sibling;
+    }
+
+    /* Add to global list and enqueue */
+    task_list_add(task);
+    sched_enqueue(task);
+
+    return task;
+}
+
+/* ======================================================================
+ * task_exit — transition current task to Zombie
+ * ====================================================================== */
+
+void task_exit(int code)
+{
+    TaskControlBlock* curr = current_task();
+    if (!curr) kernel_panic("task_exit: no current task");
+
+    curr->exit_code = code;
+
+    /* Release user address space if we're the last reference */
+    if (curr->mm) {
+        curr->mm->destroy();
+        curr->mm = nullptr;
+    }
+
+    /* Reparent children to init (or kernel_task if no init yet) */
+    TaskControlBlock* foster = g_init_task ? g_init_task : g_kernel_task;
+    if (foster && !curr->children.is_empty()) {
+        LockGuard guard(g_task_list_lock);
+
+        list_node* child_node = curr->children.next;
+        while (child_node != &curr->children) {
+            list_node* next_sib = child_node->next;
+            auto* child = list_entry<TaskControlBlock,
+                                     &TaskControlBlock::sibling>(child_node);
+            child->parent = foster;
+
+            /* Move sibling link to foster's children */
+            child_node->prev->next = child_node->next;
+            child_node->next->prev = child_node->prev;
+
+            list_node* fhead = &foster->children;
+            child_node->prev = fhead->prev;
+            child_node->next = fhead;
+            fhead->prev->next = child_node;
+            fhead->prev = child_node;
+
+            child_node = next_sib;
+        }
+        curr->children.init();
+    }
+
+    /* Transition to Zombie — TCB + kstack preserved for parent wait() */
+    curr->state = TaskState::Zombie;
+
+    /* Wake parent->wait_chld queue */
+    if (curr->parent) {
+        curr->parent->wait_chld.wake_all();
+    }
+
+    /* Yield CPU — never returns */
+    schedule();
+
+    __builtin_unreachable();
+}
+
+void task_kill_current(int sig)
+{
+    int code = (sig & 0x7F) | kExitCodeSignalFlag;
+    task_exit(code);
+}
+
+/* ======================================================================
+ * task_destroy — fully reclaim a Zombie task (called by waitpid)
+ * ====================================================================== */
+
+void task_destroy(TaskControlBlock* task)
+{
+    /* Remove from global list */
+    task_list_remove(task);
+
+    /* Remove from parent's children */
+    task->sibling.prev->next = task->sibling.next;
+    task->sibling.next->prev = task->sibling.prev;
+    task->sibling.init();
+
+    /* Free PID */
+    g_pid_alloc.free(task->pid);
+
+    /* Free kernel stack */
+    if (task->kstack_pa != 0) {
+           pmm_free_pages(pa_to_page(phys_addr(task->kstack_pa)), 1);
+    }
+
+    /* Free TCB */
+    delete task;
+}
+
+/* ======================================================================
+ * task_wait — sleep until a child exits, then reclaim it
+ * Returns PID of exited child, or -ECHILD if no children exist.
+ * ====================================================================== */
+
+int32 task_wait(int32* exit_code)
+{
+    TaskControlBlock* curr = current_task();
+    if (!curr) return -EINVAL;
+
+    while (true) {
+        TaskControlBlock* zombie = nullptr;
+
+        /* Look for a Zombie child */
+        g_task_list_lock.acquire();
+        if (curr->children.is_empty()) {
+            g_task_list_lock.release();
+            return -ECHILD;
+        }
+
+        list_node* node = curr->children.next;
+        while (node != &curr->children) {
+            auto* child = list_entry<TaskControlBlock, &TaskControlBlock::sibling>(node);
+            if (child->state == TaskState::Zombie) {
+                zombie = child;
+                break;
+            }
+            node = node->next;
+        }
+        g_task_list_lock.release();
+
+        if (zombie) {
+            int32 pid = zombie->pid;
+            if (exit_code) {
+                /* In a real kernel, we copy_to_user */
+                *exit_code = zombie->exit_code;
+            }
+            task_destroy(zombie);
+            return pid;
+        }
+
+        /* Sleep if no child has exited yet */
+        curr->wait_chld.sleep(&g_task_list_lock);
+    }
+}
+
+/* ======================================================================
+ * task_fork — clone current task
+ * ====================================================================== */
+
+int32 task_fork(arch::TrapFrame* tf)
+{
+    TaskControlBlock* curr = current_task();
+    if (!curr) return -EINVAL;
+
+    TaskControlBlock* child = alloc_task("forked", curr->priority, 0);
+    if (!child) return -ENOMEM;
+
+    /* Allocate PID */
+    child->pid = g_pid_alloc.alloc();
+    if (child->pid < 0) {
+           pmm_free_pages(pa_to_page(phys_addr(child->kstack_pa)), 1);
+        delete child;
+        return -ENOMEM;
+    }
+    child->tgid = child->pid;
+
+    /* Copy VmSpace */
+    if (curr->mm) {
+        child->mm = vmm_clone(curr->mm);
+        if (!child->mm) {
+            g_pid_alloc.free(child->pid);
+            pmm_free_pages(pa_to_page(phys_addr(child->kstack_pa)), 1);
+            delete child;
+            return -ENOMEM;
+        }
+    }
+
+    /* Copy TrapFrame to the child's kernel stack */
+    uint64 tf_addr = child->kstack_va + kKernelStackSize - sizeof(arch::TrapFrame);
+    child->trapframe = reinterpret_cast<arch::TrapFrame*>(tf_addr);
+    *child->trapframe = *tf;
+
+    /* Child returns 0 from fork */
+    child->trapframe->a[0] = 0;
+
+    /* Set up context for ret_from_fork */
+    uint64 stack_top = tf_addr; // Context goes below TrapFrame conceptually, but we just set SP
+    child->context = {};
+    child->context.ra = reinterpret_cast<uint64>(ret_from_fork);
+    child->context.sp = stack_top;
+    child->context.s[0] = tf_addr; /* Passed to a0 in ret_from_fork */
+
+    /* Add to task tree */
+    child->parent = curr;
+    {
+        LockGuard guard(g_task_list_lock);
+        list_node* head = &curr->children;
+        child->sibling.prev = head->prev;
+        child->sibling.next = head;
+        head->prev->next = &child->sibling;
+        head->prev = &child->sibling;
+    }
+
+    /* Add to global list and enqueue */
+    task_list_add(child);
+    sched_enqueue(child);
+
+    return child->pid;
+}
+
+
+/* ======================================================================
+ * task_check_stack_canary
+ * ====================================================================== */
+
+void task_check_stack_canary(TaskControlBlock* task)
+{
+    if (!task || task->kstack_va == 0) return;
+
+    auto* canary = reinterpret_cast<uint64*>(task->kstack_va);
+    if (*canary != kStackCanary) {
+        kernel_panic("kernel stack overflow detected");
+    }
+}