Add copy-on-write feature

kernel/defs.h

@@ -170,6 +170,7 @@ uint64          walkaddr(pagetable_t, uint64);
 int             copyout(pagetable_t, uint64, char *, uint64);
 int             copyin(pagetable_t, char *, uint64, uint64);
 int             copyinstr(pagetable_t, char *, uint64, uint64);
+pte_t *         walk(pagetable_t pagetable, uint64 va, int alloc);
 
 // plic.c
 void            plicinit(void);

kernel/kalloc.c

@@ -23,10 +23,13 @@ struct {
   struct run *freelist;
 } kmem;
 
+int refcnt[(PHYSTOP - KERNBASE) / PGSIZE];
+
 void
 kinit()
 {
   initlock(&kmem.lock, "kmem");
+  memset(refcnt, 0, sizeof(refcnt));
   freerange(end, (void*)PHYSTOP);
 }
 
@@ -51,6 +54,12 @@ kfree(void *pa)
   if(((uint64)pa % PGSIZE) != 0 || (char*)pa < end || (uint64)pa >= PHYSTOP)
     panic("kfree");
 
+  if(refcnt[PA2IDX((uint64)pa)] > 1){
+    refcnt[PA2IDX((uint64)pa)]--;
+    return;
+  }
+  refcnt[PA2IDX((uint64)pa)] = 0;
+
   // Fill with junk to catch dangling refs.
   memset(pa, 1, PGSIZE);
 
@@ -78,5 +87,8 @@ kalloc(void)
 
   if(r)
     memset((char*)r, 5, PGSIZE); // fill with junk
+
+  if(r) refcnt[PA2IDX((uint64)r)] = 1;
+
   return (void*)r;
 }

kernel/proc.c

@@ -20,6 +20,8 @@ static void freeproc(struct proc *p);
 
 extern char trampoline[]; // trampoline.S
 
+extern int refcnt[];
+
 // helps ensure that wakeups of wait()ing
 // parents are not lost. helps obey the
 // memory model when using p->parent.
@@ -150,8 +152,10 @@ allocproc(void)
 static void
 freeproc(struct proc *p)
 {
-  if(p->trapframe)
+  if(p->trapframe) {
+    refcnt[PA2IDX((uint64)p->trapframe)] = 0;
     kfree((void*)p->trapframe);
+  }
   p->trapframe = 0;
   if(p->pagetable)
     proc_freepagetable(p->pagetable, p->sz);

kernel/riscv.h

@@ -337,12 +337,14 @@ sfence_vma()
 
 #define PGROUNDUP(sz)  (((sz)+PGSIZE-1) & ~(PGSIZE-1))
 #define PGROUNDDOWN(a) (((a)) & ~(PGSIZE-1))
-
-#define PTE_V (1L << 0) // valid
-#define PTE_R (1L << 1)
-#define PTE_W (1L << 2)
-#define PTE_X (1L << 3)
-#define PTE_U (1L << 4) // 1 -> user can access
+#define PA2IDX(pa) ((PGROUNDDOWN(pa) - KERNBASE) / PGSIZE)
+
+#define PTE_V   (1L << 0) // valid
+#define PTE_R   (1L << 1)
+#define PTE_W   (1L << 2)
+#define PTE_X   (1L << 3)
+#define PTE_U   (1L << 4) // 1 -> user can access
+#define PTE_RSW (1L << 8) // True if the PTE is COW mapping
 
 // shift a physical address to the right place for a PTE.
 #define PA2PTE(pa) ((((uint64)pa) >> 12) << 10)

kernel/trap.c

@@ -16,6 +16,8 @@ void kernelvec();
 
 extern int devintr();
 
+extern int refcnt[];
+
 void
 trapinit(void)
 {
@@ -37,6 +39,10 @@ void
 usertrap(void)
 {
   int which_dev = 0;
+  pte_t *pte;
+  uint64 va, pa;
+  uint flags;
+  char *mem;
 
   if((r_sstatus() & SSTATUS_SPP) != 0)
     panic("usertrap: not from user mode");
@@ -67,10 +73,39 @@ usertrap(void)
     syscall();
   } else if((which_dev = devintr()) != 0){
     // ok
+  }
+  else if(r_scause() == 15) { // Page fault
+    va = r_stval(); // stval register contains the addr couldn't be translated.
+
+    if(va >= MAXVA) exit(-1);
+
+    pte = walk(p->pagetable, va, 0);
+
+    if(((*pte & PTE_W) == 0) && (*pte & PTE_RSW) && (*pte & PTE_V) && (*pte & PTE_U)) {
+      pa = PTE2PA(*pte);
+
+      if((mem = kalloc()) == 0) exit(-1);
+
+      memmove(mem, (char*)pa, PGSIZE);
+
+      *pte |= PTE_W;
+      *pte &= ~PTE_RSW;
+      *pte &= ~PTE_V;
+      flags = PTE_FLAGS(*pte);
+
+      if(mappages(p->pagetable, PGROUNDDOWN(va), PGSIZE, (uint64)mem, flags) != 0){
+        refcnt[PA2IDX((uint64)mem)] = 0;
+        kfree(mem);
+        exit(-1);
+      }
+
+      kfree((void *)pa);
+    }
+    else exit(-1);
   } else {
-    printf("usertrap(): unexpected scause %p pid=%d\n", r_scause(), p->pid);
-    printf("            sepc=%p stval=%p\n", r_sepc(), r_stval());
-    p->killed = 1;
+      printf("usertrap(): unexpected scause %p pid=%d\n", r_scause(), p->pid);
+      printf("            sepc=%p stval=%p\n", r_sepc(), r_stval());
+      p->killed = 1;
   }
 
   if(p->killed)

kernel/vm.c

@@ -15,6 +15,8 @@ extern char etext[];  // kernel.ld sets this to end of kernel code.
 
 extern char trampoline[]; // trampoline.S
 
+extern int refcnt[];
+
 // Make a direct-map page table for the kernel.
 pagetable_t
 kvmmake(void)
@@ -45,7 +47,7 @@ kvmmake(void)
 
   // map kernel stacks
   proc_mapstacks(kpgtbl);
-  
+
   return kpgtbl;
 }
 
@@ -178,10 +180,9 @@ uvmunmap(pagetable_t pagetable, uint64 va, uint64 npages, int do_free)
       panic("uvmunmap: not mapped");
     if(PTE_FLAGS(*pte) == PTE_V)
       panic("uvmunmap: not a leaf");
-    if(do_free){
-      uint64 pa = PTE2PA(*pte);
-      kfree((void*)pa);
-    }
+    uint64 pa = PTE2PA(*pte);
+    if(do_free) kfree((void*)pa);
+    else refcnt[PA2IDX(pa)]--;
     *pte = 0;
   }
 }
@@ -291,10 +292,9 @@ uvmfree(pagetable_t pagetable, uint64 sz)
   freewalk(pagetable);
 }
 
-// Given a parent process's page table, copy
-// its memory into a child's page table.
-// Copies both the page table and the
-// physical memory.
+// Given a parent process's page table,
+// map the parent's physical page into child.
+// Copies the page table only.
 // returns 0 on success, -1 on failure.
 // frees any allocated pages on failure.
 int
@@ -303,22 +303,21 @@ uvmcopy(pagetable_t old, pagetable_t new, uint64 sz)
   pte_t *pte;
   uint64 pa, i;
   uint flags;
-  char *mem;
-
   for(i = 0; i < sz; i += PGSIZE){
     if((pte = walk(old, i, 0)) == 0)
       panic("uvmcopy: pte should exist");
     if((*pte & PTE_V) == 0)
       panic("uvmcopy: page not present");
+    if(*pte & PTE_W) {
+      *pte &= ~PTE_W; // Clear W bit
+      *pte |= PTE_RSW; // Indicate that this PTE is COW mapping
+    }
     pa = PTE2PA(*pte);
     flags = PTE_FLAGS(*pte);
-    if((mem = kalloc()) == 0)
-      goto err;
-    memmove(mem, (char*)pa, PGSIZE);
-    if(mappages(new, i, PGSIZE, (uint64)mem, flags) != 0){
-      kfree(mem);
+    if(mappages(new, i, PGSIZE, (uint64)pa, flags) != 0){
       goto err;
     }
+    refcnt[PA2IDX((uint64)pa)]++;
   }
   return 0;
 
@@ -333,7 +332,7 @@ void
 uvmclear(pagetable_t pagetable, uint64 va)
 {
   pte_t *pte;
-  
+
   pte = walk(pagetable, va, 0);
   if(pte == 0)
     panic("uvmclear");
@@ -347,12 +346,39 @@ int
 copyout(pagetable_t pagetable, uint64 dstva, char *src, uint64 len)
 {
   uint64 n, va0, pa0;
+  uint flags;
+  pte_t *pte;
+  char *mem;
 
   while(len > 0){
+    if(dstva >= MAXVA) return -1;
     va0 = PGROUNDDOWN(dstva);
     pa0 = walkaddr(pagetable, va0);
+    pte = walk(pagetable, va0, 0);
     if(pa0 == 0)
       return -1;
+
+    if(((*pte & PTE_W) == 0) && (*pte & PTE_RSW)) { // write on COW page
+      if((mem = kalloc()) == 0) {
+        return -1;
+      }
+
+      memmove(mem, (char *)pa0, PGSIZE);
+
+      *pte |= PTE_W;
+      *pte &= ~PTE_RSW;
+      *pte &= ~PTE_V;
+      flags = PTE_FLAGS(*pte);
+
+      if(mappages(pagetable, va0, PGSIZE, (uint64)mem, flags) != 0) {
+        kfree(mem);
+        return -1;
+      }
+
+      kfree((void *)pa0);
+      pa0 = (uint64)mem;
+    }
+
     n = PGSIZE - (dstva - va0);
     if(n > len)
       n = len;