From dbba25aa24f86c1f8c3a0589c82985840adc1b6b Mon Sep 17 00:00:00 2001
From: Richard Qin <richard.qin.x@gmail.com>
Date: Wed, 31 Dec 2025 18:00:55 +0800
Subject: dev:Refactor Memory Management and Process Management

---
 src/kernel/proc/proc.cc | 434 +++++++++++++++++++++++++++++-------------------
 1 file changed, 266 insertions(+), 168 deletions(-)

(limited to 'src/kernel/proc')

diff --git a/src/kernel/proc/proc.cc b/src/kernel/proc/proc.cc
index 77a078e..e99b83d 100644
--- a/src/kernel/proc/proc.cc
+++ b/src/kernel/proc/proc.cc
@@ -1,120 +1,198 @@
 #include "kernel/proc.h"
+#include "kernel/proc_def.h"
+#include "kernel/runqueue.h"
+#include "kernel/cpu.h"
+#include "kernel/spinlock.h"
 #include "kernel/pmm.h"
 #include "kernel/mm.h"
-#include "drivers/uart.h"
 #include "kernel/trap.h"
-#include "common/types.h"
 #include "kernel/riscv.h"
-#include "kernel/config.h"
-#include "kernel/spinlock.h"
+#include "drivers/uart.h"
 #include "lib/string.h"
 
-#define NPROC 10
+#define NPROC 64
 #define NCPU 8
 
 extern "C" void swtch(struct Context *old_ctx, struct Context *new_ctx);
-extern char uservec[];
+extern char uservec[], trampoline[];
+extern "C" void usertrap();
 extern "C" void userret(uint64, uint64);
-extern char trampoline[];
+extern char initcode_start[], initcode_end[];
 
 struct Proc procs[NPROC];
-struct Context main_context; // The main_context here should actually be one per CPU, but it's temporarily kept for compatibility with single-core logic
-struct Proc *current_proc = nullptr;
-struct Spinlock proc_lock;
-
+struct Spinlock proc_mem_lock;
 struct cpu cpus[NCPU];
+struct RunQueue runqueues[NCPU];
+
+
+static int nextpid = 1;
+
+
+struct cpu *mycpu()
+{
+    int id = r_tp();
+    return &cpus[id];
+}
+
 
-int nextpid = 1;
+void RunQueue::init(const char *name)
+{
+    lock.init(name);
+    count = 0;
+    for (int i = 0; i < N_PRIO; i++)
+    {
+        heads[i] = nullptr;
+        tails[i] = nullptr;
+    }
+}
+
+void RunQueue::enqueue(struct Proc *p)
+{
+    if (p->priority < 0 || p->priority >= N_PRIO)
+        p->priority = PRIO_MIN;
+    p->next = nullptr;
+    if (heads[p->priority] == nullptr)
+    {
+        heads[p->priority] = p;
+        tails[p->priority] = p;
+    }
+    else
+    {
+        tails[p->priority]->next = p;
+        tails[p->priority] = p;
+    }
+    count++;
+}
 
-extern char initcode_start[];
-extern char initcode_end[];
+struct Proc *RunQueue::dequeue()
+{
+    for (int i = 0; i < N_PRIO; i++)
+    {
+        if (heads[i])
+        {
+            struct Proc *p = heads[i];
+            heads[i] = p->next;
+            if (!heads[i])
+                tails[i] = nullptr;
+            p->next = nullptr;
+            count--;
+            return p;
+        }
+    }
+    return nullptr;
+}
 
-struct Proc *myproc()
+struct Proc *RunQueue::steal()
 {
-    // Keep a simple implementation for now, should be changed later to:
-    // push_off();
-    // struct cpu *c = mycpu();
-    // struct Proc *p = c->proc;
-    // pop_off();
-    // return p;
-    return current_proc;
+    for (int i = N_PRIO - 1; i >= 0; i--)
+    {
+        if (heads[i])
+        {
+            struct Proc *p = heads[i];
+            heads[i] = p->next;
+            if (!heads[i])
+                tails[i] = nullptr;
+            p->next = nullptr;
+            count--;
+            return p;
+        }
+    }
+    return nullptr;
 }
 
+//  Trap Return Functions
+
 void usertrapret()
 {
     struct Proc *p = myproc();
 
     intr_off();
 
+    // Set stvec to point to uservec in the trampoline
+    // When an interrupt occurs in user mode, the CPU will jump here
     uint64 trampoline_uservec = TRAMPOLINE + (uservec - trampoline);
     w_stvec(trampoline_uservec);
 
+    // Fill in the Trapframe, used when uservec saves the kernel state
     p->tf->kernel_satp = r_satp();
     p->tf->kernel_sp = (uint64)p->kstack + PGSIZE;
     p->tf->kernel_trap = (uint64)usertrap;
     p->tf->kernel_hartid = r_tp();
 
+    // Set SSTATUS
+    // SPP = 0 (return to User Mode)
+    // SPIE = 1 (enable interrupts)
     unsigned long x = r_sstatus();
     x &= ~SSTATUS_SPP;
     x |= SSTATUS_SPIE;
     w_sstatus(x);
 
+    // Set SEPC to the user program counter
     w_sepc(p->tf->epc);
 
+    // Prepare user page table SATP
     uint64 satp = MAKE_SATP(p->pagetable);
 
+    // Calculate the virtual address of userret in the trampoline
     uint64 fn = TRAMPOLINE + ((uint64)userret - (uint64)trampoline);
+
+    // Jump to userret, with parameters being the TRAPFRAME address and the user page table SATP
     ((void (*)(uint64, uint64))fn)(TRAPFRAME, satp);
 }
 
+
 void forkret()
 {
-    //....
-    proc_lock.release();
+
+    struct cpu *c = mycpu();
+    c->runqueue->lock.release();
 
     Drivers::uart_puts("[Proc] forkret: jumping to user.\n");
+
+    // Enter user mode
     usertrapret();
 }
 
+//  ProcManager Implementation
+
 static Proc *allocproc()
 {
-    struct Proc* p;
-    proc_lock.acquire();
+    struct Proc *p;
+    proc_mem_lock.acquire();
     for (p = procs; p < &procs[NPROC]; p++)
     {
         if (p->state == UNUSED)
-        {
             goto found;
-        }
     }
-
-    proc_lock.release();
+    proc_mem_lock.release();
     return 0;
 
 found:
     p->state = USED;
     p->pid = nextpid++;
+    p->priority = 0;
+    p->cpu = mycpu();
 
     if ((p->tf = (struct Trapframe *)PMM::alloc_page()) == 0)
     {
         p->state = UNUSED;
-        proc_lock.release();
+        proc_mem_lock.release();
         return 0;
     }
     if ((p->kstack = PMM::alloc_page()) == 0)
     {
         PMM::free_page(p->tf);
         p->state = UNUSED;
-        proc_lock.release();
+        proc_mem_lock.release();
         return 0;
     }
     memset(p->kstack, 0, PGSIZE);
 
     memset(&p->context, 0, sizeof(p->context));
+
+    // Set the context so that the next time this process is scheduled, it starts executing directly from forkret
     p->context.ra = (uint64)forkret;
     p->context.sp = (uint64)p->kstack + PGSIZE;
-
-    proc_lock.release();
     return p;
 }
 
@@ -123,15 +201,12 @@ static void freeproc(struct Proc *p)
     if (p->tf)
         PMM::free_page(p->tf);
     p->tf = 0;
-
     if (p->kstack)
         PMM::free_page(p->kstack);
     p->kstack = 0;
-
     if (p->pagetable)
         VM::uvmfree(p->pagetable, p->sz);
     p->pagetable = 0;
-
     p->sz = 0;
     p->pid = 0;
     p->parent = 0;
@@ -142,204 +217,225 @@ static void freeproc(struct Proc *p)
 
 int fork()
 {
-    int pid;
     struct Proc *np;
     struct Proc *p = myproc();
 
     if ((np = allocproc()) == 0)
-    {
         return -1;
-    }
 
     np->pagetable = VM::uvmcreate();
-
     if (VM::uvmcopy(p->pagetable, np->pagetable, p->sz) < 0)
     {
         freeproc(np);
-        proc_lock.release(); // Maybe a bug
+        proc_mem_lock.release();
         return -1;
     }
     np->sz = p->sz;
 
-    if (VM::mappages(np->pagetable, TRAPFRAME, PGSIZE,
-                     (uint64)np->tf, PTE_R | PTE_W) < 0)
+    if (VM::mappages(np->pagetable, TRAPFRAME, PGSIZE, (uint64)np->tf, PTE_R | PTE_W) < 0)
     {
         freeproc(np);
+        proc_mem_lock.release();
         return -1;
     }
 
     *(np->tf) = *(p->tf);
-
     np->tf->a0 = 0;
-
     strcpy(np->name, p->name);
     np->parent = p;
+    np->priority = p->priority;
 
-    proc_lock.acquire();
+    proc_mem_lock.release();
+
+    struct cpu *c = mycpu();
+    c->runqueue->lock.acquire();
     np->state = RUNNABLE;
-    pid = np->pid;
-    proc_lock.release();
+    c->runqueue->enqueue(np);
+    c->runqueue->lock.release();
 
-    return pid;
+    return np->pid;
 }
 
 namespace ProcManager
 {
+
     void init()
     {
-        proc_lock.init("proc");
-
+        proc_mem_lock.init("proc_mem");
         for (int i = 0; i < NPROC; i++)
-        {
             procs[i].state = UNUSED;
-            procs[i].kstack = 0;
+
+        for (int i = 0; i < NCPU; i++)
+        {
+            cpus[i].runqueue = &runqueues[i];
+            cpus[i].runqueue->init("runqueue");
         }
     }
 
-    void create_kernel_thread(void (*func)(), const char *name)
+    void scheduler()
     {
-        struct Proc *p = nullptr;
-        for(int i = 0; i < NPROC; i++)
-        {
-            if(procs[i].state == UNUSED)
-            {
-                p = &procs[i];
-                p->pid = i + 1;
-                break;
-            }
-        }
+        struct Proc *p;
+        struct cpu *c = mycpu();
+        c->proc = 0;
+        Drivers::uart_puts("[Proc] Scheduler started (Per-CPU).\n");
 
-        if (!p)
+        for (;;)
         {
-            Drivers::uart_puts("Error: No free process slots!\n");
-            return;
-        }
+            intr_on();
 
-        p->kstack = PMM::alloc_page();
-        if (!p->kstack)
-        {
-            Drivers::uart_puts("Error: Alloc kstack failed!\n");
-            return;
-        }
+            // Get from the local queue
+            c->runqueue->lock.acquire();
+            p = c->runqueue->dequeue();
+            c->runqueue->lock.release();
 
-        p->context.ra = (uint64)func;
+            // Work Stealing
+            if (!p)
+            {
+                for (int i = 0; i < NCPU; i++)
+                {
+                    if ((uint64)i == r_tp())
+                        continue;
+
+                    struct cpu *victim = &cpus[i];
+                    victim->runqueue->lock.acquire();
+                    if (victim->runqueue->count > 0)
+                        p = victim->runqueue->steal();
+                    victim->runqueue->lock.release();
+                    if (p)
+                        break;
+                }
+            }
 
-        p->context.sp = (uint64)p->kstack + 4096;
+            if (p)
+            {
+                p->state = RUNNING;
+                c->proc = p;
 
-        p->state = RUNNABLE;
-        strcpy(p->name, name);
+                c->runqueue->lock.acquire();
+                swtch(&c->context, &p->context);
+                c->runqueue->lock.release();
 
-        Drivers::uart_puts("[Proc] Created thread: ");
-        Drivers::uart_puts(name);
-        Drivers::uart_puts("\n");
+                c->proc = 0;
+            }
+        }
     }
 
     void yield()
     {
         struct Proc *p = myproc();
-        proc_lock.acquire();
+        struct cpu *c = mycpu();
+        c->runqueue->lock.acquire();
         if (p && p->state == RUNNING)
         {
             p->state = RUNNABLE;
-            swtch(&p->context, &main_context);
+            if (p->priority < PRIO_MIN)
+                p->priority++;
+            c->runqueue->enqueue(p);
+            swtch(&p->context, &c->context);
         }
-        proc_lock.release();
+        c->runqueue->lock.release();
     }
 
-    void scheduler()
+    void user_init()
     {
-        struct Proc *p = nullptr;
-        Drivers::uart_puts("[Proc] Scheduler started.\n");
+        struct Proc *p = allocproc();
+        if (!p)
+            return;
 
-        for (;;)
+        p->pagetable = VM::uvmcreate();
+
+        uint64 initcode_sz = initcode_end - initcode_start;
+        Drivers::uart_puts("[Proc] initcode size: ");
+        Drivers::uart_put_int(initcode_sz);
+        Drivers::uart_puts(" bytes\n");
+
+        uint64 code_pages = PGROUNDUP(initcode_sz) / PGSIZE;
+
+        for (uint64 i = 0; i < code_pages; i++)
         {
-            intr_on();
+            void *mem = PMM::alloc_page();
+            if (!mem)
+            {
+                Drivers::uart_puts("user_init: out of memory for code\n");
+                return;
+            }
+            memset(mem, 0, PGSIZE);
 
-            proc_lock.acquire();
+            uint64 src_offset = i * PGSIZE;
+            uint64 copy_len = PGSIZE;
+            if (src_offset + copy_len > initcode_sz)
+            {
+                copy_len = initcode_sz - src_offset;
+            }
+            memmove(mem, (void *)(initcode_start + src_offset), copy_len);
 
-            for (p = procs; p < &procs[NPROC]; p++)
+            if (VM::mappages(p->pagetable, i * PGSIZE, PGSIZE, (uint64)mem, PTE_R | PTE_W | PTE_X | PTE_U) < 0)
             {
-                if (p->state == RUNNABLE)
-                {
-                    p->state = RUNNING;
-                    current_proc = p;
-                    // Switch context: from scheduler to process
-                    swtch(&main_context, &p->context);
-                    // The process returns here after yielding the CPU
-                    current_proc = nullptr;
-                }
+                Drivers::uart_puts("user_init: mappages code failed\n");
+                return;
             }
-            proc_lock.release();
         }
-    }
 
-    void user_init()
-    {
-        struct Proc *p;
+        uint64 stack_va = (code_pages + 1) * PGSIZE;
 
-        p = allocproc();
-        if (!p)
-        {
-            Drivers::uart_puts("user_init: failed\n");
+        void *stack = PMM::alloc_page();
+        if (!stack)
             return;
-        }
-
-        p->pagetable = VM::uvmcreate();
+        memset(stack, 0, PGSIZE);
 
-        uint64 initcode_sz = initcode_end - initcode_start;
-        VM::uvminit(p->pagetable, (uchar *)initcode_start, initcode_sz);
-        p->sz = PGSIZE;
+        if (VM::mappages(p->pagetable, stack_va, PGSIZE, (uint64)stack, PTE_R | PTE_W | PTE_U) < 0)
+            return;
 
-        if (VM::mappages(p->pagetable, TRAPFRAME, PGSIZE,
-                         (uint64)p->tf, PTE_R | PTE_W) < 0)
-        {
-            Drivers::uart_puts("user_init: map trapframe failed\n");
+        if (VM::mappages(p->pagetable, TRAPFRAME, PGSIZE, (uint64)p->tf, PTE_R | PTE_W) < 0)
             return;
-        }
 
-        
         p->tf->epc = 0;
-        p->tf->sp = PGSIZE;
-
+        p->tf->sp = stack_va + PGSIZE;
+        p->sz = stack_va + PGSIZE;
         strcpy(p->name, "initcode");
 
-        
-        proc_lock.acquire();
+        proc_mem_lock.release();
+
+        struct cpu *c = &cpus[0];
+        c->runqueue->lock.acquire();
         p->state = RUNNABLE;
-        proc_lock.release();
+        c->runqueue->enqueue(p);
+        c->runqueue->lock.release();
 
         Drivers::uart_puts("[Proc] user_init success. PID: 1\n");
+        asm volatile("fence.i");
     }
 
-    void sleep(void *chan, struct Spinlock *lk)
+    void sleep(void *chan, Spinlock *lk)
     {
         struct Proc *p = myproc();
-        if(lk != &proc_lock)
-        {
-            proc_lock.acquire();
+        struct cpu *c = mycpu();
+
+        if (lk && lk != &proc_mem_lock)
             lk->release();
-        }
+
+        c->runqueue->lock.acquire();
         p->chan = chan;
         p->state = SLEEPING;
-
-        swtch(&p->context, &main_context);
-
+        swtch(&p->context, &c->context);
         p->chan = 0;
+        c->runqueue->lock.release();
 
-        if (lk != &proc_lock)
-        {
-            proc_lock.release();
+        if (lk && lk != &proc_mem_lock)
             lk->acquire();
-        }
     }
 
     void wakeup(void *chan)
     {
-        for(struct Proc *p = procs; p < &procs[NPROC]; p++)
+        for (struct Proc *p = procs; p < &procs[NPROC]; p++)
         {
             if (p->state == SLEEPING && p->chan == chan)
             {
+                struct cpu *c = mycpu();
+                c->runqueue->lock.acquire();
                 p->state = RUNNABLE;
+                c->runqueue->enqueue(p);
+                c->runqueue->lock.release();
             }
         }
     }
@@ -348,15 +444,10 @@ namespace ProcManager
     {
         struct Proc *p = myproc();
         struct Proc *initproc = &procs[0];
-
         if (p == initproc)
-        {
-            Drivers::uart_puts("init exiting!\n");
             while (1)
                 ;
-        }
 
-        proc_lock.acquire();
         wakeup(p->parent);
 
         for (struct Proc *pp = procs; pp < &procs[NPROC]; pp++)
@@ -369,24 +460,27 @@ namespace ProcManager
         }
 
         p->xstate = status;
-        p->state = ZOMBIE;
-
-        swtch(&p->context, &main_context);
 
-        while(1);
+        struct cpu *c = mycpu();
+        c->runqueue->lock.acquire();
+        p->state = ZOMBIE;
+        swtch(&p->context, &c->context);
+        while (1)
+            ;
     }
 
     int wait(uint64 addr)
     {
-        struct Proc* np;
+        (void)addr;
+
+        struct Proc *np;
         int havekids, pid;
         struct Proc *p = myproc();
-
-        proc_lock.acquire();
-
-        for(;;)
+        for (;;)
         {
             havekids = 0;
+            proc_mem_lock.acquire();
+
             for (np = procs; np < &procs[NPROC]; np++)
             {
                 if (np->parent == p)
@@ -394,30 +488,34 @@ namespace ProcManager
                     havekids = 1;
                     if (np->state == ZOMBIE)
                     {
-                        // Found one
-                        pid = np->pid;
-                        if (addr != 0)
+
+                        struct cpu *child_cpu = np->cpu;
+                        if (child_cpu)
                         {
-                            // copyout(p->pagetable, addr, (char*)&np->xstate, sizeof(np->xstate));
-                            // 暂时不做 copyout
+                            proc_mem_lock.release();
+                            child_cpu->runqueue->lock.acquire();
+                            child_cpu->runqueue->lock.release();
+                            proc_mem_lock.acquire();
                         }
+
+                        pid = np->pid;
                         freeproc(np);
-                        proc_lock.release();
+                        proc_mem_lock.release();
                         return pid;
                     }
                 }
             }
+            proc_mem_lock.release();
+
             if (!havekids || p->state == ZOMBIE)
-            {
-                proc_lock.release();
                 return -1;
-            }
-            // sleep will atomically release proc_lock and sleep
-            // When returning, it will reacquire proc_lock
-            sleep(p, &proc_lock);
+            sleep(p, nullptr);
         }
     }
 
-
-
+    void create_kernel_thread(void (*func)(), const char *name)
+    {
+        (void)func;
+        (void)name;
+    }
 }
\ No newline at end of file
-- 
cgit v1.2.3