Lab 3

kernel/defs.h

@@ -10,6 +10,15 @@ struct stat;
 struct superblock;
 struct list;
 
+uint64 proc_getsz(struct proc *p);
+
+
+void vmprint(pagetable_t pagetable);
+
+int kref_inc(uint64 pa);
+int kref_dec(uint64 pa);
+int kref_read(uint64 pa);
+
 // bio.c
 void            binit(void);
 struct buf*     bread(uint, uint);

kernel/exec.c

@@ -8,6 +8,7 @@
 #include "elf.h"
 
 static int loadseg(pde_t *, uint64, struct inode *, uint, uint);
+static int vmprint_done = 0;
 
 int flags2perm(int flags)
 {
@@ -128,6 +129,11 @@ exec(char *path, char **argv)
   p->trapframe->sp = sp; // initial stack pointer
   proc_freepagetable(oldpagetable, oldsz);
 
+  if (!vmprint_done) {
+    // vmprint(p->pagetable);
+    vmprint_done = 1;
+  }
+
   return argc; // this ends up in a0, the first argument to main(argc, argv)
 
  bad:

kernel/kalloc.c

@@ -14,6 +14,10 @@ void freerange(void *pa_start, void *pa_end);
 extern char end[]; // first address after kernel.
                    // defined by kernel.ld.
 
+// refcount per physical page.
+// index = (pa - KERNBASE) / PGSIZE
+static int refcount[(PHYSTOP - KERNBASE) / PGSIZE];
+
 struct run {
   struct run *next;
 };
@@ -23,6 +27,18 @@ struct {
   struct run *freelist;
 } kmem;
 
+static inline int
+pa2idx(uint64 pa)
+{
+  return (pa - KERNBASE) / PGSIZE;
+}
+
+static inline int
+npage_total(void)
+{
+  return (PHYSTOP - KERNBASE) / PGSIZE;
+}
+
 void
 kinit()
 {
@@ -35,8 +51,12 @@ freerange(void *pa_start, void *pa_end)
 {
   char *p;
   p = (char*)PGROUNDUP((uint64)pa_start);
-  for(; p + PGSIZE <= (char*)pa_end; p += PGSIZE)
+  for(; p + PGSIZE <= (char*)pa_end; p += PGSIZE){
+    int idx = pa2idx((uint64)p);
+    if (idx >= 0 && idx < npage_total())
+      refcount[idx] = 1; // will be set to 0 in kfree
     kfree(p);
+  }
 }
 
 // Free the page of physical memory pointed at by pa,
@@ -51,12 +71,25 @@ kfree(void *pa)
   if(((uint64)pa % PGSIZE) != 0 || (char*)pa < end || (uint64)pa >= PHYSTOP)
     panic("kfree");
 
+  acquire(&kmem.lock);
+
+  int idx = pa2idx((uint64)pa);
+  int c = refcount[idx];
+  if (c <= 0){
+    release(&kmem.lock);
+    panic("kfree: refcount underflow");
+  }
+  c = --refcount[idx];
+  if (c > 0){
+    release(&kmem.lock);
+    return; // still being used by someone
+  }
+
   // Fill with junk to catch dangling refs.
   memset(pa, 1, PGSIZE);
 
   r = (struct run*)pa;
 
-  acquire(&kmem.lock);
   r->next = kmem.freelist;
   kmem.freelist = r;
   release(&kmem.lock);
@@ -72,11 +105,43 @@ kalloc(void)
 
   acquire(&kmem.lock);
   r = kmem.freelist;
-  if(r)
+  if(r){
     kmem.freelist = r->next;
-  release(&kmem.lock);
+    int idx = pa2idx((uint64) r);
+    refcount[idx] = 1;
+  }
 
+  release(&kmem.lock);
   if(r)
     memset((char*)r, 5, PGSIZE); // fill with junk
   return (void*)r;
 }
+
+
+// --- refcount helpers for COW ---
+int kref_inc(uint64 pa)
+{
+  acquire(&kmem.lock);
+  int idx = pa2idx(pa);
+  int v = ++refcount[idx];
+  release(&kmem.lock);
+  return v;
+}
+
+int kref_dec(uint64 pa)
+{
+  acquire(&kmem.lock);
+  int idx = pa2idx(pa);
+  int v = --refcount[idx];
+  release(&kmem.lock);
+  return v;
+}
+
+int kref_read(uint64 pa)
+{
+  acquire(&kmem.lock);
+  int idx = pa2idx(pa);
+  int v = refcount[idx];
+  release(&kmem.lock);
+  return v;
+}

kernel/proc.c

@@ -26,6 +26,13 @@ extern char trampoline[]; // trampoline.S
 // must be acquired before any p->lock.
 struct spinlock wait_lock;
 
+// helper for calling from vm
+uint64
+proc_getsz(struct proc* p)
+{
+  return p->sz;
+}
+
 // Allocate a page for each process's kernel stack.
 // Map it high in memory, followed by an invalid
 // guard page.
@@ -408,14 +415,18 @@ wait(uint64 addr)
         if(pp->state == ZOMBIE){
           // Found one.
           pid = pp->pid;
+          freeproc(pp);
+          release(&pp->lock);
+
           if(addr != 0 && copyout(p->pagetable, addr, (char *)&pp->xstate,
                                   sizeof(pp->xstate)) < 0) {
-            release(&pp->lock);
-            release(&wait_lock);
-            return -1;
+            // release(&pp->lock);
+            // release(&wait_lock);
+            // return -1;
           }
-          freeproc(pp);
-          release(&pp->lock);
+
+          // freeproc(pp);
+          // release(&pp->lock);
           release(&wait_lock);
           return pid;
         }

kernel/riscv.h

@@ -362,6 +362,7 @@ typedef uint64 *pagetable_t; // 512 PTEs
 #define PTE_W (1L << 2)
 #define PTE_X (1L << 3)
 #define PTE_U (1L << 4) // user can access
+#define PTE_COW (1L << 8) // RSW
 
 // shift a physical address to the right place for a PTE.
 #define PA2PTE(pa) ((((uint64)pa) >> 12) << 10)

kernel/sysproc.c

@@ -40,11 +40,29 @@ sys_sbrk(void)
 {
   uint64 addr;
   int n;
+  struct proc *p = myproc();
 
   argint(0, &n);
-  addr = myproc()->sz;
-  if(growproc(n) < 0)
-    return -1;
+  addr = p->sz;
+  if (n > 0) {
+    uint64 temp = p->sz + n;
+    if (temp >= MAXVA)
+      return -1;
+    p->sz += n;
+
+  } else if (n < 0) {
+    uint64 dec = -n;
+    if (dec > p->sz)
+      return -1;
+    uint64 oldsz = p->sz;
+    uint64 newsz = oldsz - dec;
+    uint64 a = PGROUNDUP(newsz);
+    uint64 end = PGROUNDUP(oldsz);
+    if (end > a) {
+      uvmunmap(p->pagetable, a, (end - a) / PGSIZE, 1);
+    }
+    p->sz = newsz;
+  }
   return addr;
 }
 

kernel/trap.c

@@ -67,12 +67,78 @@ usertrap(void)
     syscall();
   } else if((which_dev = devintr()) != 0){
     // ok
+  } else if (r_scause() == 13 || r_scause() == 15) {
+    uint64 va = r_stval();
+    uint64 a = PGROUNDDOWN(va);
+
+    if (a >= MAXVA) {
+      goto unexpected_page_fault;
+    }
+
+    pte_t *pte = walk(p->pagetable, a, 0);
+
+
+    // lazy alloc logic
+    if (pte == 0 || (*pte & PTE_V) == 0) {
+      if (a < p->sz) {
+        char *mem = kalloc();
+        if (mem == 0) {
+          // printf("lazy alloc: kalloc fault, pid=%d, va=0x%lx\n", p->pid, va);
+          setkilled(p);
+          goto done;
+        }
+        memset(mem, 0, PGSIZE);
+        if (mappages(p->pagetable, a, PGSIZE, (uint64)mem, PTE_R | PTE_W | PTE_U) != 0) {
+          kfree(mem);
+          // printf("lazy alloc: mappages fault, pid=%d, va=0x%lx\n", p->pid, va);
+          setkilled(p);
+          goto done;
+        }
+        sfence_vma();
+        goto done;
+      }
+    }
+
+    // make sure its COW page
+    if (pte && (*pte & PTE_V) && (*pte & PTE_COW) && ((*pte & PTE_W) == 0 )) {
+      uint64 pa = PTE2PA(*pte);
+      uint flags = PTE_FLAGS(*pte);
+      int refs = kref_read(pa);
+
+      if (refs > 1) {
+        char *mem = kalloc();
+        if (mem == 0) {
+          // no memory 
+          // printf("COW: out of memory pid=%d va=0x%lx\n", p->pid, va);
+          setkilled(p);
+        }
+        else{
+          memmove(mem, (void*)pa, PGSIZE);
+          // new pte
+          *pte = PA2PTE((uint64)mem);
+          *pte |= ((flags | PTE_W) & ~PTE_COW); // unset COW, set W
+          kref_dec(pa);
+          sfence_vma(); // update tlb for new va
+        }
+      } else {
+        // page is not used anymore
+        *pte = (*pte | PTE_W) & ~PTE_COW;
+        sfence_vma();
+      }
+    } else {
+      goto unexpected_page_fault;
+    }
   } else {
-    printf("usertrap(): unexpected scause 0x%lx pid=%d\n", r_scause(), p->pid);
-    printf("            sepc=0x%lx stval=0x%lx\n", r_sepc(), r_stval());
-    setkilled(p);
+      goto unexpected_page_fault;
   }
+  goto done;
 
+unexpected_page_fault:
+      // printf("usertrap(): unexpected scause 0x%lx pid=%d\n", r_scause(), p->pid);
+      // printf("            sepc=0x%lx stval=0x%lx\n", r_sepc(), r_stval());
+      setkilled(p);
+
+done:
   if(killed(p))
     exit(-1);