XV6 Lab 5: Copy On Write

xv6 中的 fork() 系统调用将父进程的所有用户空间内存复制到子进程中。如果父进程很大，复制可能需要很长的时间。更糟糕的是，这项工作通常是无用的：fork() 之后通常是子进程的 exec()，它放弃了复制的内存，没有使用复制过来的大部分内存。如果父代和子代都使用了复制的页面，并且其中一个或两个都写了它，那么这个复制才是真正需要的。

所以写时复制(copy on write, cow) 这个技术变得十分的重要。在这个 lab 中，我们需要实现一个写时复制的 fork()。我们需要修改原本的 fork() 程序中做内存申请的模块，同时我们还需要实现 usertrap 来在写时(在实际写的时候碰到 cow flag，抛出分页错误)的时候处理这个 trap。

因为一个程序很可能被 fork() 了很多的分支，所以我们需要一个计数器来确定这个 page 是要被释放还是保留。

一般来说，一个正常运行的程序在写时复制的时候会有下面几个流程:

fork() 出一个子进程 child。 child 拥有父进程的页的引用，并且页的 flag 有 cow 标识。
- 并且要把页引用 ++
child 需要向自己的页中写入新的数据. 此时需要重新分配内存，页引用 –。
当 child 返回的时候，可能有内存需要由 kernel 转换到 user。
当父进程销毁的时候，如果计数器为 0，则销毁，反之，不变。

每一页都需要一个计数器来进行计数，所以我们需要在 kernel 中加入一个数组来记录，查阅 xv6 book，我们发现有如下的图:

我们需要在 kernel data 之后的区域内申请一块内存来作为计数器存储的数组。我们可以在 kalloc.c 中实现。

// ./kernel/kalloc.c
struct {
  struct spinlock lock;
  struct run *freelist;
  // lab 5
  uint* ref_cnt;
  struct spinlock ref_cnt_lock;
  char * pa_start;
} kmem;

在这个 kmem 结构体中加入了一个 ref_cnt_lock 来保证计数器的正确引用。一个 ref_cnt pointer 来指示 ref array 的起始位置，使用 pa_start 来表示实际的 free memory 的起始位置。我们还需要修改初始化程序来正确的申请计数器数组，并且对 free memory 填充上一个初始值。

// ./kernel/kalloc.c
void
kinit()
{
  initlock(&kmem.lock, "kmem");

  // lab 5
  initlock(&kmem.ref_cnt_lock, "kmemrefcnt");
  uint64 pg_size = (((PHYSTOP - (uint64)end)) >> PGSHIFT) + 1; // get the page size
  uint64 ref_array_size = ((pg_size * sizeof(uint)) >> PGSHIFT) + 1; // caculate how many page wes need to store ref cnt.
  kmem.ref_cnt = (uint*)end; // the start of ref cnt array
  kmem.pa_start = end + (ref_array_size << PGSHIFT); // the start of user memory space.

  char *p;
  p = (char*)PGROUNDUP((uint64)kmem.pa_start);
  for (; p + PGSIZE <= (char*)PHYSTOP; p+= PGSIZE) {
    acquire(&kmem.ref_cnt_lock);
    uint64 idx = (uint64)((PGROUNDDOWN((uint64)p) - PGROUNDUP((uint64)kmem.pa_start)) / PGSIZE);
    kmem.ref_cnt[idx] = 1; // set to 1, force freerange(.., ...) function to remove all mem in start->end scope.
    release(&kmem.ref_cnt_lock);
  }

  freerange((void*)kmem.pa_start, (void*)PHYSTOP); // free memory in those scope.
}

PGSHIFT, PHYSTOP 是定义在 riscv.h 中便于使用的 macro。下面是 kfree 的代码，需要在 ref > 1 的时候减少引用，在 ref == 1 的时候释放这块内存，并且把计数器的引用次数置成 0。

void
kfree(void *pa)
{
  struct run *r;

  // lab 5
  // If ref is not 0
  acquire(&kmem.ref_cnt_lock);
  // get the ref cnt location of pa in ref cnt array
  uint64 idx = (uint64)((PGROUNDDOWN((uint64)pa) - PGROUNDUP((uint64)kmem.pa_start)) / PGSIZE);
  if (kmem.ref_cnt[idx] > 1) {
    kmem.ref_cnt[idx] --;
    release(&kmem.ref_cnt_lock);
    return;
  }

  // If ref is 1;
  if(((uint64)pa % PGSIZE) != 0 || (char*)pa < end || (uint64)pa >= PHYSTOP)
    panic("kfree");

  // Fill with junk to catch dangling refs.
  memset(pa, 1, PGSIZE);

  // lab 5
  kmem.ref_cnt[idx] = 0;
  release(&kmem.ref_cnt_lock);

  r = (struct run*)pa;

  acquire(&kmem.lock);
  r->next = kmem.freelist;
  kmem.freelist = r;
  release(&kmem.lock);
}

在申请内存的同时，也要加上计数器:

void *
kalloc(void)
{
  struct run *r;

  acquire(&kmem.lock);
  r = kmem.freelist;
  if(r)
    kmem.freelist = r->next;
  release(&kmem.lock);

  // lab 5
  if (r) {
    acquire(&kmem.ref_cnt_lock);
    uint64 idx = (uint64)((PGROUNDDOWN((uint64)r) - PGROUNDUP((uint64)kmem.pa_start)) / PGSIZE);
    kmem.ref_cnt[idx] = 1;
    release(&kmem.ref_cnt_lock);
  }

  if(r)
    memset((char*)r, 5, PGSIZE); // fill with junk
  return (void*)r;
}

在 vm.c 中我们需要修改 uvmcopy 函数，在这个函数中，我们要注释掉原来的内存申请的代码，然后使用 cow 策略来“申请”内存

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");
  }
  pa = PTE2PA(*pte);
  // clear write flag and add read flag.
  *pte = ((*pte) & (~PTE_W)) | PTE_COW;
  flags = PTE_FLAGS(*pte);
  // Map to child process's pagetable directly. Make both read-only.
  if (mappages(new, i, PGSIZE, pa, flags) != 0) {
    goto err;
  }
  // increase ref of page old.
  kincre_ref(pa);
}
return 0;

在上面的这段代码中，我将 flag 由原来的可写转变成了 cow(PTE_COW 需要在 reiscv.h 中定义)，然后将页共享，如果共享成功，则将 reference ++。

当用户需要写内存的时候，会触发一个 trap，同上一个实验，我们需要在 usertrap 中加入代码来处理这个 trap。

else if (r_scause() == 15) {
     uint64 va = r_stval();
     if (va > p->sz) p->killed = 1;
     else if (kalloc_cow(p->pagetable, va) != 0) p->killed = 1;
   }

int kalloc_cow(pagetable_t pagetable, uint64 va) {
  va = PGROUNDDOWN(va);
  if (va >= MAXVA) return -1;
  pte_t *pte = walk(pagetable, va, 0);
  if (!pte) return -1;
  uint64 pa = PTE2PA(*pte);
  if (!pa) return -1;
  uint64 flags = PTE_FLAGS(*pte);
  if (flags & PTE_COW) {
    uint64 mem = (uint64)kalloc();
    if (!mem) return -1;
    memmove((char*)mem, (char*)pa, PGSIZE);
    uvmunmap(pagetable, va, 1, 1);
    flags = (flags | PTE_W) & (~PTE_COW);
    if (mappages(pagetable, va, PGSIZE, mem, flags) != 0) {
      kfree((void*)mem);
      return -1;
    }
  }
  return 0;
}

当用户所访问的内存需要写时复制的时候，由 trap 处理程序来进行处理，kalloc_cow 程序负责申请出与原来的内存区块相同大小的区间，然后把内存 copy 过去，并且把 flag 改为可写而不是 cow。

当 kernel 需要把内存交换到 user 区时，也需要对 cow 特殊对待。在 copyout 函数中，需要加入下述代码:

if (kalloc_cow(pagetable, va0) != 0) return -1;