Lab: Copy-on-Write Fork for xv6

该实验旨在为内核实现经典的 Copy-on-Write(COW) 机制，COW 的原理非常简单，然而在实现时依旧需要注重每一处细节。正如补充资料中提到的，在实际的内核中实现 COW 可能非常棘手，2021 年对 Linux 内核的 COW 代码有两个主要的更改，不那么令人惊讶的是，这些变化都产生了意想不到的后果，并破坏了一切。

Implement copy-on-write fork

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 -q’ programs successfully.

为记录每个页面的使用情况，需在内核空间中申请一个全局数组，数组中记录每个物理页的引用进程数。在内核中操作多进程共享的资源时，需要为该资源配置自旋锁。这里引出第一个需要思考的问题 – 自旋锁的粒度如何设置？是为每个物理页单独配置一把锁？还是为整个数组配置一把大锁？两种实现均是可行的，然而为每个物理页配置一把锁会消耗大量的内存地址空间用于存储锁的相关信息（详见 struct spinlock 的结构体定义），故这里选择一把大锁保护整个数组，定义如下：

1
2
3
4
struct {
  int reference_count[PHYSTOP / PGSIZE]; // Array of process number using this page
  struct spinlock lock;
} ppages;
1
2
3
4
5
6
7
8
9
10
11
// Initialize the physical pages tracking array
void
initppages()
{
  // Assume all pages are used by the kernel initially, then freed in function freerange
  for (int i = 0; i < sizeof(ppages.reference_count) / sizeof(int); i++) {
    ppages.reference_count[i] = 1;
  }
  // Initialize the ppages lock
  initlock(&ppages.lock, "ppages");
}

这里我们假设每一物理页初始都被内核使用，且 reference_count 均为 1（调用 kfree() 函数时直接释放）。因为在 initppages() 函数后将调用 freerange(end, (void*)PHYSTOP) 函数初始化空闲页链表，对每一物理页调用 kfree() 函数。

随后开始更改 kalloc 及 kfree 函数：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
void *
kalloc(void)
{
    ...
  if(r) {
    memset((char*)r, 5, PGSIZE); // fill with junk 
    ppages_increase((uint64)r); // Increase reference count for the allocated page
  }
    ...
}

void
kfree(void *pa)
{
    ...
  acquire(&ppages.lock);
  int count = get_reference_count((uint64)pa);
  if(count <= 0) {
    panic("kfree: reference count is already zero");
  } else if(count > 1) {
    // Decrement the reference count for the physical page
    ppages.reference_count[(uint64)pa / PGSIZE]--;
    release(&ppages.lock);
    return; // Do not free if still in use
  }
  
  // count is 1 -- the page should be freed normally
  ppages.reference_count[(uint64)pa / PGSIZE] = 0; // Set reference count to zero
  release(&ppages.lock);
    ...
}

函数 kalloc() 只需在成功分配页面时将其 reference_count 加一，函数 kfree() 需根据当前页面的 reference_count 执行不同操作：当 count > 1 时将计数减一并立即返回，不释放该物理页；count == 1 时将计数归零并释放该物理页，因为当前物理页被该进程私有。值得注意的是，对 reference_count 的所有访问都应在自旋锁的保护下进行，否则锁的设置就是无意义的，必然出现临界区。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
void
usertrap(void)
{
    ...
  else if(r_scause() == 15) {
    // Store page fault
    uint64 stval = r_stval();

    intr_on();
    
    if(cow_handler(p->pagetable, stval) < 0) {
      // If copy on write fails, kill the process.
      printf("usertrap(): copy on write failed at %p pid=%d\n", stval, p->pid);
      printf("            sepc=%p stval=%p\n", r_sepc(), stval);
      setkilled(p);
    }
  }
    ...
}

阅读 RISC-V 手册可以知道 Store/AMO page fault 对应的异常码为 15，故在 scause == 15 时调用函数 cow_handler() 实现写时复制，将当前进程页表 p->pagetable 及异常时的虚拟地址 stval 作为函数参数。函数执行失败时返回 -1，此时杀死此进程。

PS：测试样例中有针对非法地址的操作，期望内核仅仅杀死异常进程而不直接 panic。

写时复制 (Copy-on-Write) 的核心函数 cow_handle 实现如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
int cow_handler(pagetable_t pagetable , uint64 va)
{
    uint64 pa, flags;
    void *mem;
    pte_t *pte;
    int count;

    if(va >= MAXVA) {
        return -1; // Invalid virtual address
    }
    
    if((pte = walk(pagetable, va, 0)) == 0) {
        return -1; // Page not found
    }

    if((*pte & PTE_C) == 0 || (*pte & PTE_U) == 0 || (*pte & PTE_V) == 0) {
        return -1; // Not a copy-on-write page
    }

    pa = PTE2PA(*pte);
    flags = PTE_FLAGS(*pte);

    acquire(&ppages.lock);
    if((count = get_reference_count(pa)) <= 0) {
        release(&ppages.lock);
        panic("cow_handler: reference count is zero");
    } else if(count == 1) {
        // the page is already private, no need to copy
        // just modify the flags
        flags = (flags & ~PTE_C) | PTE_W; // Clear the copy-on-write flag and set writable
        *pte |= flags;
        release(&ppages.lock);
    } else {
        // The page is shared, we need to create a private copy

        // Initially, release the ppages lock to avoid deadlock
        // because kalloc() need to acquire ppages.lock !!!!!
        release(&ppages.lock);

        // Allocate a new page of memory
        if((mem = kalloc()) == 0) {
            return -1; // Allocation failed
        }
        memmove(mem, (void*)pa, PGSIZE); // Copy the contents of the original page to the new page
        
        flags = (flags & ~PTE_C) | PTE_W; // Clear the copy-on-write flag and set writable
        *pte = PA2PTE(mem) | flags; // Update the page table entry to point to the new page

        kfree((void*)pa); // Free the original page
    }
    return 0;
}

为了区分写时复制导致的写入异常及单纯的非法写入操作，定义 PTE_C 为 1L << 8，用于表示当前是由于写时复制 COW 机制导致的写入异常。为了避免无意义的内核开销，如果某页是单独进程私有的 (count == 1)，无需再次申请物理页并复制数据，只需修改对应 PTE 中的标志位（抹去 PTE_C 并设置 PTE_W 位）。

PS：这段代码调试了很长时间，主要由于忽略了 kalloc 中涉及使用 ppages.lock，重复申请锁导致内核 panic。故在调用 kalloc 前需释放 ppages.lock，但这又会形成多核背景下的临界区。虽然有临界区的存在，后续处理代码在极端情况下依旧保持正确性，只会略微增大内核的复制开销。

The devil is in the details：实际上在 count == 1 时武断的认为该物理页被某一进程私有是危险的 —— 详见 https://lwn.net/Articles/849638/ 。内核中为了优化性能，常常会带来潜藏的风险，因为可能有其他优化与之冲突。