feat: support COW

modules/axdma/src/dma.rs

@@ -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/axhal/src/arch/x86_64/trap.rs

@@ -78,6 +78,7 @@ fn vec_to_str(vec: u64) -> &'static str {
 fn err_code_to_flags(err_code: u64) -> Result<MappingFlags, u64> {
     let code = PageFaultErrorCode::from_bits_truncate(err_code);
     let reserved_bits = (PageFaultErrorCode::CAUSED_BY_WRITE
+        | PageFaultErrorCode::PROTECTION_VIOLATION
         | PageFaultErrorCode::USER_MODE
         | PageFaultErrorCode::INSTRUCTION_FETCH)
         .complement();

modules/axmm/Cargo.toml

@@ -9,11 +9,16 @@ homepage.workspace = true
 repository = "https://github.com/arceos-org/arceos/tree/main/modules/axmm"
 documentation = "https://arceos-org.github.io/arceos/axmm/index.html"
 
+[features]
+default = []
+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/src/backend/alloc.rs

@@ -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/backend/mod.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/frameinfo.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/lib.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/backend/page_iter_wrapper.rs → modules/axmm/src/page_iter_wrapper.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,
         }
     }
 }