MIT 6.s081 Lab6: Copy-on-Write Fork for xv6

lab6 是实现 COW fork，在 fork 子进程的时候不直接将父进程的物理内存复制给子进程，而是只复制页表，并且把双方 PTE 设置为只读的。在进程需要进行写操作的适合，会触发 page fault，处理 page fault 的时候再复制一份内存，映射给进程，然后把 PTE 重新映射并且修改为可读写。

COW 能够节约复制的时间，只复制要修改的 page，并且可以很大程度上避免复制完后调用 exec 的浪费。

Implement copy-on write (hard)

Your task is to implement copy-on-write fork in the xv6 kernel. You are done if your modified kernel executes both the cowtest and usertests programs successfully.

首先来实现 reasonable plans 中的第三点，修改 kernel/kalloc.c。因为实现 COW 之后，就会存在多个进程引用同一个 page 的情况，所以当进程退出，要释放内存的时候不能直接释放，而是要看还有没有别的进程在引用该 page。

我们需要用一个计数器来记录每个 page 的引用数量，只有当引用数为 0 的时候才能释放。当调用 kernel/kalloc.c#kalloc 函数的时候，将这个新分配的 page 的引用计数设置为 1。当调用 kernel/kalloc.c#kfree 函数的时候，将需要释放的 page 的引用计数减 1，引用计数为 0 的时候释放物理内存。当执行 fork 系统调用的时候，在复制 page 的映射时，将 page 的引用计数加 1。

添加以下全局变量：

// ... kalloc.c

struct spinlock pincountlock; // 锁住计数器，防止发生竞态，记得在 kinit 中初始化锁

int pincounts[32768]; // 每个页的引用计数器

再实现以下函数，用以操作 page 的引用计数：

// pa 转化成计数器的下标
int
pa2pincount(uint64 pa)
{
  return (pa - KERNBASE) / PGSIZE;
}

// 将对应 page 的引用计数加 1
void
pin(uint64 pa)
{
  pa = PGROUNDUP(pa);
  acquire(&pincountlock);
  if (pincounts[pa2pincount(pa)] > 0){
    pincounts[pa2pincount(pa)]++; 
  }
  release(&pincountlock);
}

// 将对应 page 的引用计数减 1
void
unpin(uint64 pa)
{
  pa = PGROUNDUP(pa);
  acquire(&pincountlock);
  if (pincounts[pa2pincount(pa)] > 0) {
    pincounts[pa2pincount(pa)]--;
  }
  release(&pincountlock);
}

修改 kalloc 和 kfree 函数，加上对引用计数的操作：

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

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

  unpin((uint64)pa);
  acquire(&pincountlock);
  if (pincounts[pa2pincount((uint64)pa)] != 0){
    release(&pincountlock);
    return;
  }
  release(&pincountlock);

  // 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);
}

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

  acquire(&kmem.lock);
  r = kmem.freelist;
  if(r) {
    kmem.freelist = r->next;
    memset((char*)r, 5, PGSIZE); // fill with junk
    uint64 pa = (uint64) r;
    pincounts[pa2pincount(pa)] = 1;
  }
  release(&kmem.lock);
  return (void*)r;
}

然后修改 kernel/vm.c#uvmcopy() 函数，让它不实际复制内存，而是复制一下页表，并且把父子进程的 W 位置为 0。并且我们需要在 PTE 中的 flags 中选择一位用作 COW 标识。我这里选择的是 RSW 的低位。

我在 riscv.h 中定义了一个宏来方便计算：

#define PTE_COW (1L << 8)

我直接将 lab3 写过的复制页表的函数复制过来，稍微做了一些修改，然后在原本复制内存的函数中调用：

int
uvmcopy(pagetable_t old, pagetable_t new, uint64 sz)
{
  return uvmcopypg(old, new, sz, 0);
}

int
uvmcopypg(pagetable_t pagetable, pagetable_t npagetable, uint64 sz, uint64 start)
{
  pte_t *pte;
  uint64 pa, i;
  uint flags;

  for (i = PGROUNDUP(start); i < sz; i+=PGSIZE){
    if((pte = walk(pagetable, i, 0)) == 0)
      panic("uvmcopypg: pte should exist");
    if((*pte & PTE_V) == 0)
      panic("uvmcopypg: page not present");
    pa = PTE2PA(*pte);
    *pte &= ~PTE_W; // 将 W 位置为 0
    *pte |= PTE_COW; // 将 COW 位置为 1
    flags = PTE_FLAGS(*pte);
    if(mappages(npagetable, i, PGSIZE, pa, flags) != 0){
      goto err;
    }
    pin(pa); // 映射成功，将这一页的引用计数加 1
  }

  return 0;

  err:
  uvmunmap(npagetable, PGROUNDUP(start), (i - PGROUNDUP(start)) / PGSIZE, 0);
  return -1;
  
}

接下来实现一个 page fault handler 来处理 COW fork page fault，在这个函数中进行一些判断看能否处理这个 page fault，如果不行就返回 -1，否则返回 1，用于在 usertrap 中杀死进程。

int
cowhandler(uint64 va)
{
  struct proc *p = myproc();
  pte_t *pte;
  uint flags; 
  va = PGROUNDDOWN(va);
  if (va >= MAXVA) { // 测试 copyout 的时候 walk 中爆了 panic，va 存在大于 MAXVA 的情况
    return -1;
  }
  pte = walk(p->pagetable, va, 0);
  // va 必须在 heap 中，并且对应的 pte 必须存在且有效，同时 COW 位应该是 1
  if (va >= p->sz || ((pte = walk(p->pagetable, va, 0)) == 0) || (*pte & PTE_COW) == 0 || (*pte & PTE_V) == 0) {
    return -1;
  }
  uint64 pa = PTE2PA(*pte); // 将旧 page 的地址取出用于复制
  uint64 ka = (uint64) kalloc();   // 分配新的 page
  if (ka == 0) { // hints 中说如果没有内存可以分配了就杀死进程
    return -1;
  }
  memmove((void *)ka, (char*)pa, PGSIZE);
  // 将标识设置为可写，并去掉 COW 标识
  flags = PTE_FLAGS(*pte);
  flags &= ~PTE_COW;
  flags |= PTE_W;
  uvmunmap(p->pagetable, va, 1, 1); // 在取消映射时会调用 kfree，这里不用调用 unpin
  if (mappages(p->pagetable, va, PGSIZE, ka, flags)) {
    kfree((void *) ka);
    return -1;
  }
  return 0;
}

跟 lazy allocation 一样，需要在 usertrap 中调用 page fault handler 来处理 page fault。

  // ... usertrap
  intr_on();

  syscall();
} else if((which_dev = devintr()) != 0){
// ok
} else if(r_scause() == 13 || r_scause() == 15) {
  uint64 va = r_stval();
  if (cowhandler(va) == -1) { // 调用 cowhandler 返回值为 -1 就杀死进程
    p->killed = 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;
}

最后只需要在 copyout 中调用 cowhandler 函数来处理 dstva 指向的是 on-demand page 的情况。

int
copyout(pagetable_t pagetable, uint64 dstva, char *src, uint64 len)
{
  uint64 n, va0, pa0;
  while(len > 0){
    cowhandler(dstva);
    va0 = PGROUNDDOWN(dstva);
    pa0 = walkaddr(pagetable, va0);
    if(pa0 == 0)
      return -1;
    n = PGSIZE - (dstva - va0);
    if(n > len)
      n = len;
    memmove((void *)(pa0 + (dstva - va0)), src, n);

    len -= n;
    src += n;
    dstva = va0 + PGSIZE;
  }
  return 0;
}

总结

lab6 只有一个部分，任务指导书也讲得非常清楚，实现起来不是很难。我在做得时候由于一开始没有对 pincount 上锁，导致发生竞态卡了不少时间。

到这个 lab 应该就是虚拟内存系列 lab 的最后一个了，感觉还是 lab3 是最为困难的，在写完 lab3 之后对整个 xv6 就会有一定的理解了，对操作虚拟内存也会更加熟悉，做后面这几个就没有那么困难了。

对于 page table、traps、page fault 有了较为深入的理解，收获很大。