1 files changed, 42 insertions, 1 deletions

diff --git a/kernel/kalloc.c b/kernel/kalloc.c
index c2fdb86..581f0f6 100644
--- a/kernel/kalloc.c
+++ b/kernel/kalloc.c
@@ -23,10 +23,41 @@ struct {
   struct run *freelist;
 } kmem;
 
+int phypg_refcnt[PHYSTOP/PGSIZE];
+
+// Increase the refcnt
+int
+refcnt_inc(uint64 pa)
+{
+  acquire(&kmem.lock);
+  int *prefcnt = &phypg_refcnt[pa/PGSIZE];
+  if(pa > PHYSTOP || *prefcnt < 1)
+    panic("increase refcnt");
+  (*prefcnt)++;
+  release(&kmem.lock);
+  return *prefcnt;
+}
+
+// Decrease the refcnt
+int
+refcnt_dec(uint64 pa)
+{
+  acquire(&kmem.lock);
+  int *prefcnt = &phypg_refcnt[pa/PGSIZE];
+  if(pa > PHYSTOP || *prefcnt < 1)
+    panic("decrease refcnt");
+  (*prefcnt)--;
+  release(&kmem.lock);
+  return *prefcnt;
+}
+
 void
 kinit()
 {
   initlock(&kmem.lock, "kmem");
+  // init all refcnt to 1, which would later be freed to 0 by kfree()
+  for(uint64 p = PGROUNDUP((uint64)end); p + PGSIZE <= PHYSTOP; p += PGSIZE)
+    phypg_refcnt[p/PGSIZE] = 1;
   freerange(end, (void*)PHYSTOP);
 }
 
@@ -51,6 +82,12 @@ kfree(void *pa)
   if(((uint64)pa % PGSIZE) != 0 || (char*)pa < end || (uint64)pa >= PHYSTOP)
     panic("kfree");
 
+  refcnt_dec((uint64)pa);
+
+  if(phypg_refcnt[(uint64)pa/PGSIZE] > 0)
+    // We still have refs to this phy page, do not actually free it
+    return;
+
   // Fill with junk to catch dangling refs.
   memset(pa, 1, PGSIZE);
 
@@ -72,8 +109,12 @@ kalloc(void)
 
   acquire(&kmem.lock);
   r = kmem.freelist;
-  if(r)
+  if(r){
+    if(phypg_refcnt[(uint64)r/PGSIZE])
+      panic("kalloc: invalid refcnt");
+    phypg_refcnt[(uint64)r/PGSIZE] = 1;
     kmem.freelist = r->next;
+  }
   release(&kmem.lock);
 
   if(r)