still buggy, maybe concurrency

kernel/defs.h

@@ -63,6 +63,8 @@ void            ramdiskrw(struct buf*);
 void*           kalloc(void);
 void            kfree(void *);
 void            kinit(void);
+void            refcount_add(uint64 pa, int val);
+int             refcount_get(uint64 pa);
 
 // log.c
 void            initlog(int, struct superblock*);

kernel/kalloc.c

@@ -13,6 +13,7 @@ void freerange(void *pa_start, void *pa_end);
 
 extern char end[]; // first address after kernel.
                    // defined by kernel.ld.
+int refcount[(PHYSTOP - KERNBASE) >> PGSHIFT];
 
 struct run {
   struct run *next;
@@ -27,16 +28,37 @@ void
 kinit()
 {
   initlock(&kmem.lock, "kmem");
+  memset(refcount, 0, sizeof(refcount));
   freerange(end, (void*)PHYSTOP);
 }
+int total_free;
+int alloc_count;
+int free_count;
+int count_ref() {
+  int count = 0;
+  for (int i = 0; i < NELEM(refcount); ++ i) {
+    count += (refcount[i] > 0);
+  }
+  return count;
+}
+
+void show_page_count()
+{
+  int ref_count = count_ref();
+  printf("count ref: %d / %d; alloc=%d, free=%d\n", ref_count, total_free, alloc_count, free_count);
+}
 
 void
 freerange(void *pa_start, void *pa_end)
 {
   char *p;
+  total_free = alloc_count = free_count = 0;
   p = (char*)PGROUNDUP((uint64)pa_start);
-  for(; p + PGSIZE <= (char*)pa_end; p += PGSIZE)
+  for(; p + PGSIZE <= (char*)pa_end; p += PGSIZE, total_free ++){
+    refcount_add((uint64)p, 1);
     kfree(p);
+  }
+  show_page_count();
 }
 
 // Free the page of physical memory pointed at by pa,
@@ -51,12 +73,29 @@ kfree(void *pa)
   if(((uint64)pa % PGSIZE) != 0 || (char*)pa < end || (uint64)pa >= PHYSTOP)
     panic("kfree");
 
+//* [COW] if still has refcount after dec, keep it
+  refcount_add((uint64)pa, -1);
+  free_count++;
+  acquire(&kmem.lock);
+  if (refcount_get((uint64)pa) > 0) {
+    release(&kmem.lock);
+    return;
+  } else if (refcount_get((uint64)pa) < 0) {
+    panic("kfree: invalid refcount");
+  }
+  refcount[((uint64)pa - KERNBASE) >> PGSHIFT] = 0;
+  release(&kmem.lock);
+
   // Fill with junk to catch dangling refs.
   memset(pa, 1, PGSIZE);
 
   r = (struct run*)pa;
 
   acquire(&kmem.lock);
+  if (alloc_count) {
+    if (r == 0) panic("kfree r");
+    if (kmem.freelist == 0) panic("kfree");
+  }
   r->next = kmem.freelist;
   kmem.freelist = r;
   release(&kmem.lock);
@@ -72,11 +111,36 @@ kalloc(void)
 
   acquire(&kmem.lock);
   r = kmem.freelist;
-  if(r)
+  if(r) {
     kmem.freelist = r->next;
+    int idx = ((uint64)r - KERNBASE) >> PGSHIFT;
+    if (idx < 0 || idx >= NELEM(refcount)) panic("!!!");
+    refcount[idx] = 1;
+    alloc_count ++;
+  } else {
+    show_page_count();
+  }
   release(&kmem.lock);
 
   if(r)
     memset((char*)r, 5, PGSIZE); // fill with junk
   return (void*)r;
 }
+
+void
+refcount_add(uint64 pa, int val)
+{
+  acquire(&kmem.lock);
+  int idx = (pa - KERNBASE) >> PGSHIFT;
+  if (idx < 0 || idx >= NELEM(refcount)) panic("!!!");
+  refcount[idx] += val;
+  release(&kmem.lock);
+}
+
+int
+refcount_get(uint64 pa)
+{
+  int idx = (pa - KERNBASE) >> PGSHIFT;
+  if (idx < 0 || idx >= NELEM(refcount)) panic("!!!");
+  return refcount[idx];
+}

kernel/proc.c

@@ -321,7 +321,8 @@ fork(void)
   acquire(&np->lock);
   np->state = RUNNABLE;
   release(&np->lock);
-
+extern void show_page_count();
+  show_page_count();
   return pid;
 }
 

kernel/riscv.h

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

kernel/trap.c

@@ -65,6 +65,37 @@ usertrap(void)
     intr_on();
 
     syscall();
+  } else if(r_scause() == 15){
+    // Store/AMO page fault
+    uint64 stval = r_stval();
+    pte_t* pte = walk(p->pagetable, stval, 0);
+    uint64 pa0 = PTE2PA((*pte));
+    if ((*pte & PTE_COW) && !(*pte & PTE_W)) {
+      // printf("COW: sepc=%p, stval=%p, pte=%p\n", r_sepc(), r_stval(), *pte);
+      if (refcount_get(pa0) > 1) {
+        uint64 newpage = (uint64)kalloc();
+        if (newpage == 0) {
+          printf("killed due to no mem\n");
+          setkilled(p);
+        } else {
+          pte_t save_pte = * pte;
+          memmove((void*)newpage, (void*)pa0, PGSIZE);
+          uvmunmap(p->pagetable, PGROUNDDOWN(stval), 1, 0);
+          refcount_add(pa0, -1);
+          mappages(p->pagetable, PGROUNDDOWN(stval), PGSIZE, newpage,
+                  (PTE_FLAGS(save_pte) | PTE_W) & (~PTE_COW));
+        }
+      } else {
+        //* [COW] this pa is no more shared
+        *pte |= PTE_W;
+        *pte &= ~PTE_COW;
+      }
+    } else {
+      printf("Store/AMO page fault\n");
+      printf("    pid=%d PTE=%p\n", p->pid, *pte);
+      printf("    sepc=%p stval=%p\n", r_sepc(), r_stval());
+      setkilled(p);
+    }
   } else if((which_dev = devintr()) != 0){
     // ok
   } else {

kernel/vm.c

@@ -308,7 +308,7 @@ uvmcopy(pagetable_t old, pagetable_t new, uint64 sz)
   pte_t *pte;
   uint64 pa, i;
   uint flags;
-  char *mem;
+  // char *mem;
 
   for(i = 0; i < sz; i += PGSIZE){
     if((pte = walk(old, i, 0)) == 0)
@@ -316,14 +316,23 @@ uvmcopy(pagetable_t old, pagetable_t new, uint64 sz)
     if((*pte & PTE_V) == 0)
       panic("uvmcopy: page not present");
     pa = PTE2PA(*pte);
+    if (*pte & PTE_W) {
+      //* [COW] let go those read-only pages
+      *pte &= ~PTE_W;
+      *pte |= PTE_COW;
+    }
+    //* [COW] clear parent's PTE_W, and 'll be carried to child's pgtbl
     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);
+    // *[COW] no alloc and copy here
+    // if((mem = kalloc()) == 0)
+    //   goto err;
+    // memmove(mem, (char*)pa, PGSIZE);
+    // if(mappages(new, i, PGSIZE, (uint64)mem, flags) != 0){
+    if(mappages(new, i, PGSIZE, pa, flags) != 0){
+      // kfree(mem);
       goto err;
     }
+    refcount_add(pa, 1);
   }
   return 0;
 
@@ -361,6 +370,28 @@ copyout(pagetable_t pagetable, uint64 dstva, char *src, uint64 len)
     n = PGSIZE - (dstva - va0);
     if(n > len)
       n = len;
+    pte_t* pte = walk(pagetable, va0, 0);
+    if ((*pte & PTE_COW) && !(*pte & PTE_W)) {
+      if (refcount_get(PTE2PA(*pte)) > 1) {
+        uint64 newpage = (uint64)kalloc();
+        if (newpage == 0) {
+          printf("killed due to no mem in copyout\n");
+          setkilled(myproc());
+          exit(-1);
+        }
+        pte_t save_pte = *pte;
+        memmove((void*)newpage, (void*)pa0, PGSIZE);
+        uvmunmap(pagetable, PGROUNDDOWN(va0), 1, 0);
+        refcount_add(pa0, -1);
+        mappages(pagetable, PGROUNDDOWN(va0), PGSIZE, newpage,
+                (PTE_FLAGS(save_pte) | PTE_W) & (~PTE_COW));
+        pa0 = newpage;
+      } else {
+        //* [COW] this pa is no more shared
+        *pte |= PTE_W;
+        *pte &= ~PTE_COW;
+      }
+    }
     memmove((void *)(pa0 + (dstva - va0)), src, n);
 
     len -= n;