From 149e169b2c456a9f20a2b728a1ce914ee2882f95 Mon Sep 17 00:00:00 2001
From: Jianfeng Jiang <jiangjianfeng.jjf@antgroup.com>
Date: Tue, 4 Jul 2023 10:44:07 +0800
Subject: [PATCH] Drop parent once all inherited pages is committed by child

---
 Cargo.lock                                |  1 +
 services/libs/jinux-std/Cargo.toml        |  5 +++++
 services/libs/jinux-std/src/vm/vmo/mod.rs | 22 +++++++++++++++++++++-
 3 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/Cargo.lock b/Cargo.lock
index 172c61fc3..a8e6c2135 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -573,6 +573,7 @@ dependencies = [
  "align_ext",
  "ascii",
  "bitflags",
+ "bitvec",
  "controlled",
  "core2",
  "cpio-decoder",
diff --git a/services/libs/jinux-std/Cargo.toml b/services/libs/jinux-std/Cargo.toml
index 684cc5453..efbf6568a 100644
--- a/services/libs/jinux-std/Cargo.toml
+++ b/services/libs/jinux-std/Cargo.toml
@@ -47,3 +47,8 @@ getrandom = { version = "0.2.10", default-features = false, features = ["rdrand"
 [dependencies.lazy_static]
 version = "1.0"
 features = ["spin_no_std"]
+
+[dependencies.bitvec]
+version = "1"
+default-features = false
+features = ["atomic", "alloc"]
diff --git a/services/libs/jinux-std/src/vm/vmo/mod.rs b/services/libs/jinux-std/src/vm/vmo/mod.rs
index 1d37a6117..535def9e1 100644
--- a/services/libs/jinux-std/src/vm/vmo/mod.rs
+++ b/services/libs/jinux-std/src/vm/vmo/mod.rs
@@ -3,6 +3,8 @@
 use core::ops::Range;
 
 use align_ext::AlignExt;
+use bitvec::bitvec;
+use bitvec::vec::BitVec;
 use jinux_frame::vm::{VmAllocOptions, VmFrameVec, VmIo};
 use jinux_rights::Rights;
 
@@ -167,6 +169,8 @@ struct InheritedPages {
     /// page with index `idx + parent_page_idx_offset` in parent vmo
     parent_page_idx_offset: usize,
     is_copy_on_write: bool,
+    /// The pages already committed by child
+    committed_pages: BitVec,
 }
 
 impl InheritedPages {
@@ -176,11 +180,13 @@ impl InheritedPages {
         parent_page_idx_offset: usize,
         is_copy_on_write: bool,
     ) -> Self {
+        let committed_pages = bitvec![0; page_range.len()];
         Self {
             parent: Some(parent),
             page_range,
             parent_page_idx_offset,
             is_copy_on_write,
+            committed_pages,
         }
     }
 
@@ -236,6 +242,19 @@ impl InheritedPages {
     fn should_child_commit_page(&self, write_page: bool) -> bool {
         !self.is_copy_on_write || (self.is_copy_on_write && write_page)
     }
+
+    fn mark_page_as_commited(&mut self, page_idx: usize) {
+        if !self.contains_page(page_idx) {
+            return;
+        }
+        let idx = page_idx - self.page_range.start;
+        debug_assert_eq!(self.committed_pages[idx], false);
+        debug_assert!(self.parent.is_some());
+        self.committed_pages.set(idx, true);
+        if self.committed_pages.count_ones() == self.page_range.len() {
+            self.parent.take();
+        }
+    }
 }
 
 impl VmoInner {
@@ -292,10 +311,11 @@ impl VmoInner {
             return Ok(frames);
         }
 
-        let inherited_pages = self.inherited_pages.as_ref().unwrap();
+        let inherited_pages = self.inherited_pages.as_mut().unwrap();
         let frame = inherited_pages.get_frame_from_parent(page_idx, write_page, commit_if_none)?;
 
         if inherited_pages.should_child_commit_page(write_page) {
+            inherited_pages.mark_page_as_commited(page_idx);
             self.insert_frame(page_idx, frame.clone());
         }