feat: support CopyOn-Write mechanism (#50)

* Add basc COW

- support x86_64

refactor: improve page fault handling and page management

- Fix whitespace in PageFaultErrorCode flags

- Simplify page fault COW logic

- Rename page ref count methods for clarity

- Move page copy logic to Page struct

- Add better documentation for page management

refactor: improve copy_with_cow

    - use `insert_area`

    - add `MemorySet` patch

fmt: fmt

chore: update test workflow dependencies

- Add memory_set crate dependency

- Add memory_addr crate dependency

refactor: copy_with_cow and add support alloc contiguous pages

- Remove myself git dependencies for memory_set and memory_addr

- Update Cargo.toml and Cargo.lock

- Adjust page allocation to support contiguous pages

- Improve COW handling with page size parameter

* Squash merge cow-backend into cow

* refactor: replace Arc-based frame tracking with frame info table

- Remove Arc and FrameTracker in favor of atomic ref counting

- Add frameinfo module for physical frame management

- Simplify COW fault handling with direct frame operations

- Update backend allocator to use new frame tracking system

- only support with x86

* refactor: simplify address space handling logic

- Remove redundant backend matching in fork

- Reorganize COW fault handling comments

- Clean up page table update flow

* refactor: improve memory management safety

- Replace ok() with expect() for protect/map operations

- Simplify page copy with copy_nonoverlapping

- Add physical address validation in phys_to_pfn

* chore: clean up

* refactor: Rename `populate_area` to `ensure_region_mapped` with COW support

- Renamed `populate_area` to `ensure_region_mapped`
- Refactored `handle_cow_fault` to take a mutable page table reference

* chore: Modify according to comment

* refactor: replace PageIter4K with PageIterWrapper in AddrSpace

- Use align from backend for page iteration

- Remove TODO and FIXME comments

- Simplify page iteration logic

* refactor: simplify frame reference counting

- Replace direct frame info access with FrameRefTable API

- Remove redundant align parameter from protect()

- Rename ensure_region_mapped to populate_area

- Simplify COW handling logic

- Consolidate frame info operations into FrameRefTable struct

* refactor: (populate_area) use a iterator to walk through the area covered by the specified access interval

* feat: add COW feature for memory management

- Add conditional compilation for COW feature

- Refactor clone_or_err into try_clone with COW support

- Update frame allocation/deallocation logic

- Move frameinfo module behind feature flag

- Add new cow feature to Cargo.toml

* refactor: simplify COW handling and frame table initialization

- Consolidate COW fault handling logic
- Replace LazyInit with lazy_static for frame table
- Remove redundant frame table initialization
- Clean up memory area copying logic

* refactor: reorganize memory management modules

- Move PageIterWrapper to root module

- Clean up unused imports

- Simplify cow feature logic

- Fix docstring alignment

- Remove redundant page size handling

* doc: fix doc link

Author: mingzi47

Cargo.lock

@@ -2,10 +2,7 @@ use core::{alloc::Layout, ptr::NonNull};
 
 use allocator::{AllocError, AllocResult, BaseAllocator, ByteAllocator};
 use axalloc::{DefaultByteAllocator, global_allocator};
-use axhal::{
-    mem::virt_to_phys,
-    paging::{MappingFlags, PageSize},
-};
+use axhal::{mem::virt_to_phys, paging::MappingFlags};
 use kspin::SpinNoIrq;
 use log::{debug, error};
 use memory_addr::{PAGE_SIZE_4K, VirtAddr, va};
@@ -97,7 +94,7 @@ impl DmaAllocator {
         let expand_size = num_pages * PAGE_SIZE_4K;
         axmm::kernel_aspace()
             .lock()
-            .protect(vaddr, expand_size, flags, PageSize::Size4K)
+            .protect(vaddr, expand_size, flags)
             .map_err(|e| {
                 error!("change table flag fail: {e:?}");
                 AllocError::NoMemory

modules/axdma/src/dma.rs

@@ -12,12 +12,14 @@ documentation = "https://arceos-org.github.io/arceos/axmm/index.html"
 [features]
 default = []
 copy-from = ["page_table_multiarch/copy-from"]
+cow = []
 
 [dependencies]
 axhal = { workspace = true, features = ["paging"] }
 axalloc = { workspace = true }
 axconfig = { workspace = true }
 
+lazy_static = { version = "1.5", features = ["spin_no_std"] }
 log = "=0.4.21"
 axerrno = "0.1"
 lazyinit = "0.2"

modules/axmm/Cargo.toml

@@ -1,9 +1,12 @@
-use crate::backend::page_iter_wrapper::PageIterWrapper;
+use crate::page_iter_wrapper::PageIterWrapper;
 use axalloc::global_allocator;
 use axhal::mem::{phys_to_virt, virt_to_phys};
 use axhal::paging::{MappingFlags, PageSize, PageTable};
 use memory_addr::{PAGE_SIZE_4K, PhysAddr, VirtAddr};
 
+#[cfg(feature = "cow")]
+use crate::frameinfo::frame_table;
+
 use super::Backend;
 
 /// Allocates a physical frame, with an option to zero it out.
@@ -25,14 +28,18 @@ use super::Backend;
 /// - If `zeroed` is `true`, the function uses `unsafe` operations to zero out the memory.
 /// - The allocated memory must be accessed via its physical address, which requires
 ///   conversion using `virt_to_phys`.
-fn alloc_frame(zeroed: bool, align: PageSize) -> Option<PhysAddr> {
+pub fn alloc_frame(zeroed: bool, align: PageSize) -> Option<PhysAddr> {
     let page_size: usize = align.into();
     let num_pages = page_size / PAGE_SIZE_4K;
     let vaddr = VirtAddr::from(global_allocator().alloc_pages(num_pages, page_size).ok()?);
     if zeroed {
         unsafe { core::ptr::write_bytes(vaddr.as_mut_ptr(), 0, page_size) };
     }
     let paddr = virt_to_phys(vaddr);
+
+    #[cfg(feature = "cow")]
+    frame_table().inc_ref(paddr);
+
     Some(paddr)
 }
 
@@ -43,6 +50,9 @@ fn alloc_frame(zeroed: bool, align: PageSize) -> Option<PhysAddr> {
 /// The size of the memory to be freed is determined by the `align` parameter,
 /// which must be a multiple of 4KiB.
 ///
+/// If `cow` feature is enabled, this function decreases the reference count associated with the frame.
+/// When the reference count reaches 1, it actually frees the frame memory.
+///
 /// # Parameters
 /// - `frame`: The physical address of the memory to be freed.
 /// - `align`: The alignment requirement for the memory, must be a multiple of 4KiB.
@@ -52,10 +62,15 @@ fn alloc_frame(zeroed: bool, align: PageSize) -> Option<PhysAddr> {
 ///   otherwise undefined behavior may occur.
 /// - If the deallocation fails, the function will call `panic!`. Details about
 ///   the failure can be obtained from the global memory allocator’s error messages.
-fn dealloc_frame(frame: PhysAddr, align: PageSize) {
+pub fn dealloc_frame(frame: PhysAddr, align: PageSize) {
+    #[cfg(feature = "cow")]
+    if frame_table().dec_ref(frame) > 1 {
+        return;
+    }
+
+    let vaddr = phys_to_virt(frame);
     let page_size: usize = align.into();
     let num_pages = page_size / PAGE_SIZE_4K;
-    let vaddr = phys_to_virt(frame);
     global_allocator().dealloc_pages(vaddr.as_usize(), num_pages);
 }
 

modules/axmm/src/aspace.rs

@@ -3,12 +3,12 @@
 use axhal::paging::{MappingFlags, PageTable};
 use memory_addr::VirtAddr;
 use memory_set::MappingBackend;
-pub use page_iter_wrapper::PageIterWrapper;
 use page_table_multiarch::PageSize;
 
 mod alloc;
 mod linear;
-mod page_iter_wrapper;
+
+pub use alloc::{alloc_frame, dealloc_frame};
 
 /// A unified enum type for different memory mapping backends.
 ///

modules/axmm/src/backend/alloc.rs

@@ -0,0 +1,85 @@
+//! FrameInfo
+//!
+//! A simple physical FrameInfo manager is provided to track and manage
+//! the reference count for every 4KB memory page frame in the system.
+//!
+//! There is a [`FrameInfo`] struct for each physical page frame
+//! that keeps track of its reference count.
+//! NOTE: If the page is huge page, its [`FrameInfo`] is placed at the
+//! starting physical address.
+use core::{
+    array,
+    sync::atomic::{AtomicUsize, Ordering},
+};
+
+use alloc::boxed::Box;
+use lazy_static::lazy_static;
+use memory_addr::PhysAddr;
+// 4 kb page
+const FRAME_SHIFT: usize = 12;
+
+pub const MAX_FRAME_NUM: usize = axconfig::plat::PHYS_MEMORY_SIZE >> FRAME_SHIFT;
+
+lazy_static! {
+    static ref FRAME_INFO_TABLE: FrameRefTable = FrameRefTable::default();
+}
+
+pub(crate) fn frame_table() -> &'static FrameRefTable {
+    &FRAME_INFO_TABLE
+}
+
+pub(crate) struct FrameRefTable {
+    data: Box<[FrameInfo; MAX_FRAME_NUM]>,
+}
+
+impl Default for FrameRefTable {
+    fn default() -> Self {
+        FrameRefTable {
+            data: Box::new(array::from_fn(|_| FrameInfo::default())),
+        }
+    }
+}
+
+impl FrameRefTable {
+    fn info(&self, paddr: PhysAddr) -> &FrameInfo {
+        let index = (paddr.as_usize() - axconfig::plat::PHYS_MEMORY_BASE) >> FRAME_SHIFT;
+        &self.data[index]
+    }
+
+    /// Increases the reference count of the frame associated with a physical address.
+    ///
+    /// # Parameters
+    /// - `paddr`: It must be an aligned physical address; if it's a huge page,
+    ///   it must be the starting physical address.
+    pub fn inc_ref(&self, paddr: PhysAddr) {
+        self.info(paddr).ref_count.fetch_add(1, Ordering::SeqCst);
+    }
+
+    /// Decreases the reference count of the frame associated with a physical address.
+    ///
+    /// - `paddr`: It must be an aligned physical address; if it's a huge page,
+    ///   it must be the starting physical address.
+    ///
+    /// # Returns
+    /// The updated reference count after decrementing.
+    pub fn dec_ref(&self, paddr: PhysAddr) -> usize {
+        self.info(paddr).ref_count.fetch_sub(1, Ordering::SeqCst)
+    }
+
+    /// Returns the `FrameInfo` structure associated with a given physical address.
+    ///
+    /// # Parameters
+    /// - `paddr`: It must be an aligned physical address; if it's a huge page,
+    ///   it must be the starting physical address.
+    ///
+    /// # Returns
+    /// A reference to the `FrameInfo` associated with the given physical address.
+    pub fn ref_count(&self, paddr: PhysAddr) -> usize {
+        self.info(paddr).ref_count.load(Ordering::SeqCst)
+    }
+}
+
+#[derive(Default)]
+pub(crate) struct FrameInfo {
+    ref_count: AtomicUsize,
+}

modules/axmm/src/backend/mod.rs

@@ -8,6 +8,9 @@ extern crate alloc;
 
 mod aspace;
 mod backend;
+#[cfg(feature = "cow")]
+mod frameinfo;
+mod page_iter_wrapper;
 
 pub use self::aspace::AddrSpace;
 pub use self::backend::Backend;

modules/axmm/src/frameinfo.rs

@@ -56,7 +56,6 @@ impl PageIterWrapper {
             PageSize::Size4K => PageIter4K::<VirtAddr>::new(start, end).map(Self::Size4K),
             PageSize::Size2M => PageIter2M::<VirtAddr>::new(start, end).map(Self::Size2M),
             PageSize::Size1G => PageIter1G::<VirtAddr>::new(start, end).map(Self::Size1G),
-            _ => None,
         }
     }
 }