dev(spinlock): refactoring the implementation of locks

2 files changed, 196 insertions, 101 deletions

diff --git a/src/kernel/proc/proc.cc b/src/kernel/proc/proc.cc
index b0ab370..77a078e 100644
--- a/src/kernel/proc/proc.cc
+++ b/src/kernel/proc/proc.cc
@@ -10,6 +10,7 @@
 #include "lib/string.h"
 
 #define NPROC 10
+#define NCPU 8
 
 extern "C" void swtch(struct Context *old_ctx, struct Context *new_ctx);
 extern char uservec[];
@@ -17,9 +18,11 @@ extern "C" void userret(uint64, uint64);
 extern char trampoline[];
 
 struct Proc procs[NPROC];
-struct Context main_context;
+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_lock;
+
+struct cpu cpus[NCPU];
 
 int nextpid = 1;
 
@@ -28,7 +31,12 @@ extern char initcode_end[];
 
 struct Proc *myproc()
 {
-    // 简化处理
+    // 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;
 }
 
@@ -62,7 +70,7 @@ void usertrapret()
 void forkret()
 {
     //....
-    Spinlock::release(&proc_lock);
+    proc_lock.release();
 
     Drivers::uart_puts("[Proc] forkret: jumping to user.\n");
     usertrapret();
@@ -71,7 +79,7 @@ void forkret()
 static Proc *allocproc()
 {
     struct Proc* p;
-    Spinlock::acquire(&proc_lock);
+    proc_lock.acquire();
     for (p = procs; p < &procs[NPROC]; p++)
     {
         if (p->state == UNUSED)
@@ -80,7 +88,7 @@ static Proc *allocproc()
         }
     }
 
-    Spinlock::release(&proc_lock);
+    proc_lock.release();
     return 0;
 
 found:
@@ -90,14 +98,14 @@ found:
     if ((p->tf = (struct Trapframe *)PMM::alloc_page()) == 0)
     {
         p->state = UNUSED;
-        Spinlock::release(&proc_lock);
+        proc_lock.release();
         return 0;
     }
     if ((p->kstack = PMM::alloc_page()) == 0)
     {
         PMM::free_page(p->tf);
         p->state = UNUSED;
-        Spinlock::release(&proc_lock);
+        proc_lock.release();
         return 0;
     }
     memset(p->kstack, 0, PGSIZE);
@@ -106,7 +114,7 @@ found:
     p->context.ra = (uint64)forkret;
     p->context.sp = (uint64)p->kstack + PGSIZE;
 
-    Spinlock::release(&proc_lock);
+    proc_lock.release();
     return p;
 }
 
@@ -148,7 +156,7 @@ int fork()
     if (VM::uvmcopy(p->pagetable, np->pagetable, p->sz) < 0)
     {
         freeproc(np);
-        Spinlock::release(&proc_lock);
+        proc_lock.release(); // Maybe a bug
         return -1;
     }
     np->sz = p->sz;
@@ -167,10 +175,10 @@ int fork()
     strcpy(np->name, p->name);
     np->parent = p;
 
-    Spinlock::acquire(&proc_lock);
+    proc_lock.acquire();
     np->state = RUNNABLE;
     pid = np->pid;
-    Spinlock::release(&proc_lock);
+    proc_lock.release();
 
     return pid;
 }
@@ -179,7 +187,7 @@ namespace ProcManager
 {
     void init()
     {
-        Spinlock::init(&proc_lock, "proc");
+        proc_lock.init("proc");
 
         for (int i = 0; i < NPROC; i++)
         {
@@ -229,13 +237,13 @@ namespace ProcManager
     void yield()
     {
         struct Proc *p = myproc();
-        Spinlock::acquire(&proc_lock);
+        proc_lock.acquire();
         if (p && p->state == RUNNING)
         {
             p->state = RUNNABLE;
             swtch(&p->context, &main_context);
         }
-        Spinlock::release(&proc_lock);
+        proc_lock.release();
     }
 
     void scheduler()
@@ -247,7 +255,7 @@ namespace ProcManager
         {
             intr_on();
 
-            Spinlock::acquire(&proc_lock);
+            proc_lock.acquire();
 
             for (p = procs; p < &procs[NPROC]; p++)
             {
@@ -255,13 +263,13 @@ namespace ProcManager
                 {
                     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;
                 }
             }
-            Spinlock::release(&proc_lock);
+            proc_lock.release();
         }
     }
 
@@ -296,21 +304,20 @@ namespace ProcManager
         strcpy(p->name, "initcode");
 
         
-        Spinlock::acquire(&proc_lock);
+        proc_lock.acquire();
         p->state = RUNNABLE;
-        Spinlock::release(&proc_lock);
+        proc_lock.release();
 
         Drivers::uart_puts("[Proc] user_init success. PID: 1\n");
     }
 
-    void sleep(void *chan, struct spinlock *lk)
+    void sleep(void *chan, struct Spinlock *lk)
     {
         struct Proc *p = myproc();
         if(lk != &proc_lock)
         {
-            Drivers::uart_puts("sleep: lock mismatch\n");
-            while (1)
-                ;
+            proc_lock.acquire();
+            lk->release();
         }
         p->chan = chan;
         p->state = SLEEPING;
@@ -318,6 +325,12 @@ namespace ProcManager
         swtch(&p->context, &main_context);
 
         p->chan = 0;
+
+        if (lk != &proc_lock)
+        {
+            proc_lock.release();
+            lk->acquire();
+        }
     }
 
     void wakeup(void *chan)
@@ -343,7 +356,8 @@ namespace ProcManager
                 ;
         }
 
-        Spinlock::acquire(&proc_lock);
+        proc_lock.acquire();
+        wakeup(p->parent);
 
         for (struct Proc *pp = procs; pp < &procs[NPROC]; pp++)
         {
@@ -354,8 +368,6 @@ namespace ProcManager
             }
         }
 
-        wakeup(p->parent);
-
         p->xstate = status;
         p->state = ZOMBIE;
 
@@ -370,7 +382,7 @@ namespace ProcManager
         int havekids, pid;
         struct Proc *p = myproc();
 
-        Spinlock::acquire(&proc_lock);
+        proc_lock.acquire();
 
         for(;;)
         {
@@ -390,16 +402,18 @@ namespace ProcManager
                             // 暂时不做 copyout
                         }
                         freeproc(np);
-                        Spinlock::release(&proc_lock);
+                        proc_lock.release();
                         return pid;
                     }
                 }
             }
             if (!havekids || p->state == ZOMBIE)
             {
-                Spinlock::release(&proc_lock);
+                proc_lock.release();
                 return -1;
             }
+            // sleep will atomically release proc_lock and sleep
+            // When returning, it will reacquire proc_lock
             sleep(p, &proc_lock);
         }
     }
diff --git a/src/kernel/spinlock.cc b/src/kernel/spinlock.cc
index 210e0b9..426f639 100644
--- a/src/kernel/spinlock.cc
+++ b/src/kernel/spinlock.cc
@@ -2,106 +2,187 @@
 #include "kernel/riscv.h"
 #include "kernel/proc.h"
 #include "drivers/uart.h"
-#include "common/types.h"
 
-#define NCPU 8
+extern struct cpu cpus[];
 
-namespace Spinlock
+// Get the pointer of the current CPU
+// Depends on r_tp() in riscv.h to read the tp register
+struct cpu *mycpu()
 {
-    int g_noff[NCPU];
+    int id = r_tp();
+    return &cpus[id];
+}
 
-    int g_intena[NCPU];
+// Get the current process
+// Must be called with interrupts disabled, or read atomically in a safe manner
+struct Proc *myproc()
+{
+    push_off();
+    struct cpu *c = mycpu();
+    struct Proc *p = c->proc;
+    pop_off();
+    return p;
+}
 
-    int cpuid()
-    {
-        return r_tp();
-    }
+// Interrupt Depth Management
+void push_off(void);
+void pop_off(void);
+
+// Spinlock Implementation
+
+void Spinlock::init(const char *lock_name)
+{
+    name = lock_name;
+    locked = 0;
+    cpu = nullptr;
+}
 
-    void init(struct spinlock *lk, char* name)
+void Spinlock::acquire()
+{
+    push_off(); // Disable interrupts: prevent deadlocks (interrupt handler requests a lock that is already held)
+
+    if (holding())
     {
-        lk->name = name;
-        lk->locked = 0;
-        lk->cpu_id = -1;
+        Drivers::uart_puts("panic: recursive acquire ");
+        Drivers::uart_puts(name);
+        Drivers::uart_puts("\n");
+        while (1)
+            ;
     }
 
-    void push_off()
+    // Atomic Exchange
+    // __ATOMIC_ACQUIRE: Ensures that reads and writes after the acquire are not reordered before the acquire
+    while (__atomic_exchange_n(&locked, 1, __ATOMIC_ACQUIRE))
     {
-        int old = intr_get();
-        intr_off();
-
-        int id = cpuid();
-        if (g_noff[id] == 0)
+        // Busy-wait (Spinning)
+        // Optimization: Read-only, no writes, reducing cache coherence traffic
+        while (__atomic_load_n(&locked, __ATOMIC_RELAXED) == 1)
         {
-            g_intena[id] = old;
+            asm volatile("nop"); // RISC-V indicates that the CPU is in a spin state
         }
-        g_noff[id] += 1;
     }
 
-    void pop_off()
+    // Record holder
+    cpu = mycpu();
+}
+
+void Spinlock::release()
+{
+    if (!holding())
     {
-        int id = cpuid();
-        if (intr_get())
-            Drivers::uart_puts("pop_off - interruptible\n");
+        Drivers::uart_puts("panic: release not holding\n");
+        while (1)
+            ;
+    }
 
-        if (g_noff[id] < 1)
-            Drivers::uart_puts("pop_off - count error\n");
+    cpu = nullptr;
 
-        g_noff[id] -= 1;
+    // Atomic Store
+    // __ATOMIC_RELEASE: Ensures all reads and writes before the Release are completed
+    __atomic_store_n(&locked, 0, __ATOMIC_RELEASE);
 
-        if (g_noff[id] == 0 && g_intena[id])
-        {
-            intr_on();
-        }
-    }
+    pop_off(); // 2. 恢复中断状态
+}
+
+// Check if the current CPU holds the lock
+bool Spinlock::holding()
+{
+    // Interrupt checks must be disabled, otherwise a process switch may occur causing CPU fluctuations.
+    int r;
+    push_off();
+    r = (locked && cpu == mycpu());
+    pop_off();
+    return r;
+}
+
+
+// Mutex Implementation (Sleep Locks)
+
+void Mutex::init(const char *lock_name)
+{
+    name = lock_name;
+    locked = 0;
+    pid = 0;
+    lk.init("mutex_internal"); // 初始化内部自旋锁
+}
+
+void Mutex::acquire()
+{
+    lk.acquire(); // Acquire the internal lock to protect the locked field
 
-    int holding(struct spinlock *lk)
+    while (locked)
     {
-        int r;
-        push_off();
-        r = (lk->locked && lk->cpu_id == cpuid());
-        pop_off();
-        return r;
+        // Lock is held, preparing to sleep
+        // Atomicity contract of sleep():
+        // 1. Mark the current process as SLEEPING
+        // 2. Release lk (so others can operate on locked)
+        // 3. Switch CPU (context switch)
+        // 4. After being woken up, reacquire the lk lock and return
+        ProcManager::sleep(this, &lk);
     }
 
-    void acquire(struct spinlock *lk)
-    {
-        push_off();
+    locked = 1;
+    pid = myproc()->pid;
 
-        if(holding(lk))
-        {
-            Drivers::uart_puts("acquire: panic - recursive lock ");
-            Drivers::uart_puts(lk->name);
-            Drivers::uart_puts("\n");
-            while (1)
-                ;
-        }
+    lk.release();
+}
 
-        while (__sync_lock_test_and_set(&lk->locked, 1) != 0)
-        {
-            // 自旋等待
-        }
+void Mutex::release()
+{
+    lk.acquire();
 
-        __sync_synchronize();
+    locked = 0;
+    pid = 0;
 
-        lk->cpu_id = cpuid();
-    }
+    // Wake up all processes sleeping on this lock (chan = this)
+    ProcManager::wakeup(this);
 
-    void release(struct spinlock *lk)
-    {
-        if (!holding(lk))
-        {
-            Drivers::uart_puts("release: panic - not holding\n");
-            while (1)
-                ;
-        }
+    lk.release();
+}
 
-        lk->cpu_id = -1;
+bool Mutex::holding()
+{
+    lk.acquire(); // Must use a lock when reading to prevent race conditions.
+    bool r = locked && (pid == myproc()->pid);
+    lk.release();
+    return r;
+}
+
+// Interrupt Depth Management
+
+void push_off(void)
+{
+    int old = intr_get();
+    intr_off();
+
+    struct cpu *c = mycpu();
 
-        __sync_synchronize();
+    if (c->noff == 0)
+        c->intena = old; // 记录最外层的中断状态
 
-        __sync_lock_release(&lk->locked);
+    c->noff += 1;
+}
 
-        pop_off();
+void pop_off(void)
+{
+    struct cpu *c = mycpu();
+
+    if (intr_get())
+    {
+        Drivers::uart_puts("panic: pop_off - interruptible\n");
+        while (1)
+            ;
     }
+    if (c->noff < 1)
+    {
+        Drivers::uart_puts("panic: pop_off count\n");
+        while (1)
+            ;
+    }
+
+    c->noff -= 1;
 
-} // namespace Spinlock
+    // An interrupt is only truly enabled when the nesting level returns to zero and the outermost layer allows interrupts.
+    if (c->noff == 0 && c->intena)
+        intr_on();
+}
\ No newline at end of file