6.s081 Lab4 Copy-on-Write Fork for Xv6

课上谈了个 COW 优化策略，这里就要具体实现了。

关于 COW(Copy-On-Write, 写时拷贝)

The Problem

The fork() system call in xv6 copies all of the parent process's user-space memory into the child. If the parent is large, copying can take a long time. Worse, the work is often largely wasted; for example, a fork() followed by exec() in the child will cause the child to discard the copied memory, probably without ever using most of it. On the other hand, if both parent and child use a page, and one or both writes it, a copy is truly needed.

在 xv6 中，系统调用 fork() 会将父进程页表中所有指向的数据页拷贝到子进程中。有些页不一定用到，用到也不一定写，那就先浅拷贝，等到要写的时候再深拷贝。

The Solution

The goal of copy-on-write (COW) fork() is to defer allocating and copying physical memory pages for the child until the copies are actually needed, if ever.

COW fork() creates just a pagetable for the child, with PTEs for user memory pointing to the parent's physical pages. COW fork() marks all the user PTEs in both parent and child as not writable. When either process tries to write one of these COW pages, the CPU will force a page fault. The kernel page-fault handler detects this case, allocates a page of physical memory for the faulting process, copies the original page into the new page, and modifies the relevant PTE in the faulting process to refer to the new page, this time with the PTE marked writeable. When the page fault handler returns, the user process will be able to write its copy of the page.

COW fork() makes freeing of the physical pages that implement user memory a little trickier. A given physical page may be referred to by multiple processes' page tables, and should be freed only when the last reference disappears.

具体要怎么做呢？在原来的版本，fork() 会调用 uvmcopy() 函数，对于父进程内的每一个虚拟地址 va，其对应一个页表项 pte 及其指向的数据页 pa，调用 kalloc() 在内存中分配一个新的数据页 npa，将 pa 的内容完整拷贝到 npa 中，最后将 pte 的 flag 和 npa 组合成一个新的 npte，插入子进程页表中 va 对应的位置。

引入 COW 优化后，就不需要在这里进行 kalloc() 了。我们可以直接拷贝页表项，毕竟虚拟地址和相应的物理地址都是一样的。但注意，要将 PTE_W 位清除，因为这里有多个进程同时指向同一片内存，如果都允许写的话可能会产生冲突，所以这里将写的权限取消，但保留读权限。此外，还要添加一个 PTE_COW 位，用于标识这一页是享受到 COW 优化了的。

那么一旦尝试对这个虚拟地址进行写入操作，通过查页表发现没有对应物理地址的写权限，就会产生异常，进入 usertrap()。在这个函数中，我们就要对这种特定异常进行处理。通过查看 RISC-V 文档可知，出现该异常时，SCAUSE 寄存器中的值为 15，同时 STVAL 中的值为引起异常的虚拟地址。一旦遇到该异常，说明我们尝试写入一个 PTE_W=0 的页，在 COW 策略下，我们就需要在这里进行拷贝了，那第一步就是判断该虚拟地址对应的物理页是否「有资格享用 COW」。方法很简单，只需要检查页表项的 PTE_V，PTE_U 和 PTE_COW 位是否均为 1 即可。

如果通过检查，那就分配新的数据页 npa，并将原来的页拷贝到新页中，还有别忘了更新页表项，令其指向新的页。哦哦哦，还得设置一下标志位，毕竟指向新页后整个优化策略就与它无关了，就需要去掉 PTE_COW，并且加上 PTE_W。

如果这里分配页失败，那就杀掉进程，处理方式为 p->killed = 1。

这是在用户态写入的情况。还有一种情况是在内核态向用户空间写入数据，那就是 copyout() 函数。同样的，如果遇到一个虚拟地址对应页表项的 PTE_COW 被置 1，说明这一页要在写入时进行拷贝，那流程跟上面基本是一样的，只不过遇到没有空闲内存而导致分配页失败后直接 return -1 即可。

到这里就够了吗？还不够。因为还有一个关键问题我们没有解决，那就是父子进程中的如果都使用 COW 分配新页并修改页表项后，原来的那个页面将没有任何页表项指向——内存泄漏产生了。本质原因在于现有的机制没法意识到一个页面什么时候该被释放，一个合理的措施是维护所有页面的「引用计数」，每当一个页面被一个页表项指向时，调用 pin() 将其计数值加一，同样的，每当一个页表项取消指向时，调用 unpin() 将其计数值减一，而减到 0 后就调用 kfree 进行释放。这样一来就解决了内存泄漏的问题。这样就需要把原本对 kfree 的调用统一替换为 unpin()。

分配一个页面时，其引用计数自动设为 1。

引用计数能做的事情很多，比如当某一进程写入 COW 页却发现这一页的引用计数只有 1 时，它会意识到这一页被自己独占，那就不需要进行新页的分配了，直接修改页表项的标志位即可。

lab 手册提示我们可以将引用计数实现为一个数组，定义在 kernel/kalloc.c 中。那么数组有多大呢？毕竟我们是通过页号进行索引的，那么数组大小就是可用页数，通过查看 kernel/riscv.h 我们发现了两个宏，KERNBASE 与 PHYSTOP，这两值之差便是整个可用内存空间大小，再除以 PGSIZE 那就是可用页数了。

数组大小确定了，那么索引方式也很快能想到，直接 (pa - KERNBASE) >> 12 即可。虽然前面可能有部分页会被浪费，但胜在简单可靠。

思路捋清了，那代码就很好写了。

Preparation

切换到对应分支

$ git fetch
$ git checkout cow
$ make clean

Task: Implement copy-on write

引用计数

首先实现引用计数。

kernel/kalloc.c
#define NPAGES (PHYSTOP-KERNBASE)/PGSIZE

static uint8 ref_count[NPAGES];

// arg start_pa must be PGROUNDed
uint64
page_idx(void* start_pa)
{
  uint64 idx = ((uint64)start_pa-KERNBASE) >> 12;
  if (idx >= NPAGES) {
    return -1;
  }
  return idx;
}

// plus the reference count of page at pa
void
pin(void* pa)
{
  void* start_pa = (void*)PGROUNDDOWN((uint64)pa);
  uint64 idx = page_idx(start_pa);
  if (idx == -1) {
    panic("invalid page\n");
  }

  ref_count[idx]++;
}

// minus the reference count of page at pa
// if the count is 0 after unpin, free it
void
unpin(void* pa)
{
  void* start_pa = (void*)PGROUNDDOWN((uint64)pa);
  uint64 idx = page_idx(start_pa);
  if (idx == -1) {
    panic("invalid page\n");
  }

  if (--ref_count[idx] == 0) {
    kfree(start_pa);
  }
}

uint8
getcount(void* pa)
{
  void* start_pa = (void*)PGROUNDDOWN((uint64)pa);
  uint64 idx = page_idx(start_pa);
  if (idx == -1) {
    panic("invalid page\n");
  }
  return ref_count[idx];
}

int
pinned(void* pa)
{
  return getcount(pa) > 0;
}

int
exown(void* pa)
{
  return getcount(pa) == 1;
}

void
kinit()
{
  initlock(&kmem.lock, "kmem");
  memset(ref_count, 0, sizeof(ref_count));   // (!new)
  freerange(end, (void*)PHYSTOP);
}

// Allocate one 4096-byte page of physical memory.
// Returns a pointer that the kernel can use.
// Returns 0 if the memory cannot be allocated.
void *
kalloc(void)
{
  struct run *r;

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

  if(r){
    memset((char*)r, 5, PGSIZE);
    pin((void*)r);   // (!new)
  }
  return (void*)r;
}

为了在其他文件中调用，部分函数需要在 def.h 中添加声明。

uvmcopy()

kernel/vm.c
int
uvmcopy(pagetable_t old, pagetable_t new, uint64 sz)
{
  pte_t *pte;
  uint64 pa, i;
  uint flags;

  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);
    *pte &= ~PTE_W;   // clear PTE_W
    *pte |= PTE_COW;  // add PTE_COW

    flags = PTE_FLAGS(*pte);
    if (mappages(new, i, PGSIZE, pa, flags) != 0) {
      uvmunmap(new, 0, i / PGSIZE, 0);
      return -1;
    }

    pin((void*)pa); // !!!很重要
  }
  return 0;
}

usertrap()

kernel/trap.c
void
usertrap(void)
{
  ...
  if(r_scause() == 8){
    ...
  } else if (r_scause() == 15) {
    if(p->killed)
      exit(-1);

    uint64 start_va = PGROUNDDOWN(r_stval());
    if (start_va >= MAXVA) // walk 前检查一下，避免 panic 导致测试卡住
      exit(-1);

    pte_t* pte = walk(p->pagetable, start_va, 0);
    if (pte == 0 || (*pte & PTE_V) == 0 || (*pte & PTE_U) == 0 || (*pte & PTE_COW) == 0)
      panic("invalid pte");

    uint64 pa = PTE2PA(*pte);
    if (exown((void*)pa)) { // 如果独占，只需修改标志位，无需拷贝，省时省力
      *pte &= ~PTE_COW;
      *pte |= PTE_W;
    } else {
      char* npa = (char*)kalloc();
      if(npa == 0){
        // 如果没有可用内存，杀掉进程
        acquire(&p->lock);
        p->killed = 1;
        release(&p->lock);
      } else {
        memmove(npa, (char*)pa, PGSIZE);

        *pte &= ~PTE_COW;
        *pte |= PTE_W;
        *pte = PTE_FLAGS(*pte) | PA2PTE(npa);

        unpin((void*)pa);
      }
    }
  }
  ...
}

copyout()

kernel/vm.c
int
copyout(pagetable_t pagetable, uint64 dstva, char *src, uint64 len)
{
  pte_t *pte;
  uint64 n, va0, pa0;
  while(len > 0){
    va0 = PGROUNDDOWN(dstva);
    if (va0 >= MAXVA) // walk 前检查一下，避免 panic 导致测试卡住
      return -1;

    pte = walk(pagetable, va0, 0);
    if (pte == 0 || (*pte & PTE_V) == 0 || (*pte & PTE_U) == 0)
      return -1;

    pa0 = PTE2PA(*pte);
    if (*pte & PTE_COW) {
      if (exown((void*)pa0)) {
        *pte &= ~PTE_COW;
        *pte |= PTE_W;
      } else {
        char* npa = (char*)kalloc();
        if(npa){
          memmove(npa, (char*)pa0, PGSIZE);

          *pte &= ~PTE_COW;
          *pte |= PTE_W;
          *pte = PTE_FLAGS(*pte) | PA2PTE(npa);

          unpin((void*)pa0);
        }
        pa0 = (uint64)npa;  // 如果 npa 为 0，则最终会 return -1
                            // 反之，npa 成为新的 pa0，即 dst pa
      }
    }

    if(pa0 == 0)
      return -1;
    ...
  }
  return 0;
}

测试结果

$ make grade
...
== Test running cowtest ==
$ make qemu-gdb
(6.8s)
== Test   simple ==
  simple: OK
== Test   three ==
  three: OK
== Test   file ==
  file: OK
== Test usertests ==
$ make qemu-gdb
(116.5s)
== Test   usertests: copyin ==
  usertests: copyin: OK
== Test   usertests: copyout ==
  usertests: copyout: OK
== Test   usertests: all tests ==
  usertests: all tests: OK
== Test time ==
time: OK
Score: 110/110

最后的工作

1. git commit -am "" 将所有修改提交到本地;
2. 执行 make handin。由于 lab0 保存了 APIKey，故直接成功提交；

可选的挑战再说吧，没有什么想做的欲望。