From 99a22ff124972872838fa5dc7094fa9a5eea557e Mon Sep 17 00:00:00 2001
From: Chen Chengjun <chenchengjun.ccj@antgroup.com>
Date: Thu, 15 Aug 2024 16:53:54 +0800
Subject: [PATCH] Improve efficiency of global TLB flushing

---
 kernel/aster-nix/src/vm/vmar/mod.rs | 17 +-----------
 ostd/src/mm/vm_space.rs             | 43 +++++++++++++++++++++++++----
 2 files changed, 39 insertions(+), 21 deletions(-)

diff --git a/kernel/aster-nix/src/vm/vmar/mod.rs b/kernel/aster-nix/src/vm/vmar/mod.rs
index 01151278d..7d35ec5f1 100644
--- a/kernel/aster-nix/src/vm/vmar/mod.rs
+++ b/kernel/aster-nix/src/vm/vmar/mod.rs
@@ -101,16 +101,6 @@ pub(super) struct Vmar_ {
     parent: Weak<Vmar_>,
 }
 
-impl Drop for Vmar_ {
-    fn drop(&mut self) {
-        let mut cursor = self
-            .vm_space
-            .cursor_mut(&(self.base..self.base + self.size))
-            .unwrap();
-        cursor.unmap(self.size);
-    }
-}
-
 struct VmarInner {
     /// Whether the vmar is destroyed
     is_destroyed: bool,
@@ -295,12 +285,7 @@ impl Vmar_ {
         if !self.is_root_vmar() {
             return_errno_with_message!(Errno::EACCES, "The vmar is not root vmar");
         }
-        let mut cursor = self
-            .vm_space
-            .cursor_mut(&(self.base..self.base + self.size))
-            .unwrap();
-        cursor.unmap(self.size);
-        drop(cursor);
+        self.vm_space.clear();
         let mut inner = self.inner.lock();
         inner.child_vmar_s.clear();
         inner.vm_mappings.clear();
diff --git a/ostd/src/mm/vm_space.rs b/ostd/src/mm/vm_space.rs
index 6cc660949..a00eefbec 100644
--- a/ostd/src/mm/vm_space.rs
+++ b/ostd/src/mm/vm_space.rs
@@ -21,8 +21,8 @@ use super::{
 };
 use crate::{
     arch::mm::{
-        current_page_table_paddr, tlb_flush_addr, tlb_flush_addr_range, PageTableEntry,
-        PagingConsts,
+        current_page_table_paddr, tlb_flush_addr, tlb_flush_addr_range,
+        tlb_flush_all_excluding_global, PageTableEntry, PagingConsts,
     },
     cpu::CpuExceptionInfo,
     mm::{
@@ -102,6 +102,29 @@ impl VmSpace {
         self.pt.activate();
     }
 
+    /// Clears all mappings.
+    pub fn clear(&self) {
+        let mut cursor = self.pt.cursor_mut(&(0..MAX_USERSPACE_VADDR)).unwrap();
+        loop {
+            // SAFETY: It is safe to un-map memory in the userspace.
+            let result = unsafe { cursor.take_next(MAX_USERSPACE_VADDR - cursor.virt_addr()) };
+            match result {
+                PageTableItem::Mapped { page, .. } => {
+                    drop(page);
+                }
+                PageTableItem::NotMapped { .. } => {
+                    break;
+                }
+                PageTableItem::MappedUntracked { .. } => {
+                    panic!("found untracked memory mapped into `VmSpace`");
+                }
+            }
+        }
+        // TODO: currently this method calls x86_64::flush_all(), which rewrite the Cr3 register.
+        // We should replace it with x86_64::flush_pcid(InvPicdCommand::AllExceptGlobal) after enabling PCID.
+        tlb_flush_all_excluding_global();
+    }
+
     pub(crate) fn handle_page_fault(
         &self,
         info: &CpuExceptionInfo,
@@ -136,10 +159,20 @@ impl VmSpace {
             prop.flags -= PageFlags::W;
         };
 
-        // SAFETY: It is safe to protect memory in the userspace.
-        while let Some(range) = unsafe { cursor.protect_next(end - cursor.virt_addr(), &mut op) } {
-            tlb_flush_addr(range.start);
+        loop {
+            // SAFETY: It is safe to protect memory in the userspace.
+            unsafe {
+                if cursor
+                    .protect_next(end - cursor.virt_addr(), &mut op)
+                    .is_none()
+                {
+                    break;
+                }
+            };
         }
+        // TODO: currently this method calls x86_64::flush_all(), which rewrite the Cr3 register.
+        // We should replace it with x86_64::flush_pcid(InvPicdCommand::AllExceptGlobal) after enabling PCID.
+        tlb_flush_all_excluding_global();
 
         let page_fault_handler = {
             let new_handler = Once::new();