From bec2c97637e18d916a14c17a8f35978daa42fa55 Mon Sep 17 00:00:00 2001
From: Zhang Junyang <junyang@stu.pku.edu.cn>
Date: Sun, 1 Sep 2024 12:36:43 +0800
Subject: [PATCH] Don't copy on write if this is the only reference

---
 kernel/src/vm/vmar/vm_mapping.rs |  9 ++++++++-
 ostd/src/mm/frame/mod.rs         | 15 +++++++++++++++
 ostd/src/mm/page/mod.rs          | 15 +++++++++++++++
 3 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/kernel/src/vm/vmar/vm_mapping.rs b/kernel/src/vm/vmar/vm_mapping.rs
index 0c04b514f..1859f6947 100644
--- a/kernel/src/vm/vmar/vm_mapping.rs
+++ b/kernel/src/vm/vmar/vm_mapping.rs
@@ -241,7 +241,14 @@ impl VmMapping {
                 return Ok(());
             }
 
-            if self.is_shared {
+            // If the forked child or parent immediately unmaps the page after
+            // the fork without accessing it, we are the only reference to the
+            // frame. We can directly map the frame as writable without
+            // copying. In this case, the reference count of the frame is 2 (
+            // one for the mapping and one for the frame handle itself).
+            let only_reference = frame.reference_count() == 2;
+
+            if self.is_shared || only_reference {
                 cursor.protect(PAGE_SIZE, |p| p.flags |= PageFlags::W);
             } else {
                 let new_frame = duplicate_frame(&frame)?;
diff --git a/ostd/src/mm/frame/mod.rs b/ostd/src/mm/frame/mod.rs
index bc41dc1ff..dc0123686 100644
--- a/ostd/src/mm/frame/mod.rs
+++ b/ostd/src/mm/frame/mod.rs
@@ -80,6 +80,21 @@ impl Frame {
             core::ptr::copy_nonoverlapping(src.as_ptr(), self.as_mut_ptr(), self.size());
         }
     }
+
+    /// Get the reference count of the frame.
+    ///
+    /// It returns the number of all references to the page, including all the
+    /// existing page handles ([`Frame`]) and all the mappings in the page
+    /// table that points to the page.
+    ///
+    /// # Safety
+    ///
+    /// The function is safe to call, but using it requires extra care. The
+    /// reference count can be changed by other threads at any time including
+    /// potentially between calling this method and acting on the result.
+    pub fn reference_count(&self) -> u32 {
+        self.page.reference_count()
+    }
 }
 
 impl From<Page<FrameMeta>> for Frame {
diff --git a/ostd/src/mm/page/mod.rs b/ostd/src/mm/page/mod.rs
index c74c087c1..9c4004df9 100644
--- a/ostd/src/mm/page/mod.rs
+++ b/ostd/src/mm/page/mod.rs
@@ -164,6 +164,21 @@ impl<M: PageMeta> Page<M> {
         unsafe { &*(self.ptr as *const M) }
     }
 
+    /// Get the reference count of the page.
+    ///
+    /// It returns the number of all references to the page, including all the
+    /// existing page handles ([`Page`], [`DynPage`]), and all the mappings in the
+    /// page table that points to the page.
+    ///
+    /// # Safety
+    ///
+    /// The function is safe to call, but using it requires extra care. The
+    /// reference count can be changed by other threads at any time including
+    /// potentially between calling this method and acting on the result.
+    pub fn reference_count(&self) -> u32 {
+        self.ref_count().load(Ordering::Relaxed)
+    }
+
     fn ref_count(&self) -> &AtomicU32 {
         unsafe { &(*self.ptr).ref_count }
     }