From 52f1787d3524f3ae1fa18b404e9bfc9e5c5a2fd3 Mon Sep 17 00:00:00 2001
From: Zhang Junyang <junyang@stu.pku.edu.cn>
Date: Tue, 24 Sep 2024 10:26:49 +0800
Subject: [PATCH] Remove syncronizations for the `VmSpace` PF handler

---
 kernel/src/vm/vmar/mod.rs |  4 ++--
 ostd/src/mm/vm_space.rs   | 15 +++++----------
 2 files changed, 7 insertions(+), 12 deletions(-)

diff --git a/kernel/src/vm/vmar/mod.rs b/kernel/src/vm/vmar/mod.rs
index e1a53c0bf..c6d98e063 100644
--- a/kernel/src/vm/vmar/mod.rs
+++ b/kernel/src/vm/vmar/mod.rs
@@ -229,7 +229,7 @@ impl Vmar_ {
             vm_mappings: BTreeMap::new(),
             free_regions,
         };
-        let vm_space = VmSpace::new();
+        let mut vm_space = VmSpace::new();
         vm_space.register_page_fault_handler(handle_page_fault_wrapper);
         Vmar_::new(vmar_inner, Arc::new(vm_space), 0, ROOT_VMAR_CAP_ADDR, None)
     }
@@ -661,7 +661,7 @@ impl Vmar_ {
             let vm_space = if let Some(parent) = parent {
                 parent.vm_space().clone()
             } else {
-                let new_space = VmSpace::new();
+                let mut new_space = VmSpace::new();
                 new_space.register_page_fault_handler(handle_page_fault_wrapper);
                 Arc::new(new_space)
             };
diff --git a/ostd/src/mm/vm_space.rs b/ostd/src/mm/vm_space.rs
index ea6729372..68c7ccb10 100644
--- a/ostd/src/mm/vm_space.rs
+++ b/ostd/src/mm/vm_space.rs
@@ -15,8 +15,6 @@ use core::{
     sync::atomic::{AtomicPtr, Ordering},
 };
 
-use spin::Once;
-
 use super::{
     io::Fallible,
     kspace::KERNEL_PAGE_TABLE,
@@ -56,7 +54,7 @@ use crate::{
 #[derive(Debug)]
 pub struct VmSpace {
     pt: PageTable<UserMode>,
-    page_fault_handler: Once<fn(&VmSpace, &CpuExceptionInfo) -> core::result::Result<(), ()>>,
+    page_fault_handler: Option<fn(&VmSpace, &CpuExceptionInfo) -> core::result::Result<(), ()>>,
     /// A CPU can only activate a `VmSpace` when no mutable cursors are alive.
     /// Cursors hold read locks and activation require a write lock.
     activation_lock: RwLock<()>,
@@ -67,7 +65,7 @@ impl VmSpace {
     pub fn new() -> Self {
         Self {
             pt: KERNEL_PAGE_TABLE.get().unwrap().create_user_page_table(),
-            page_fault_handler: Once::new(),
+            page_fault_handler: None,
             activation_lock: RwLock::new(()),
         }
     }
@@ -156,21 +154,18 @@ impl VmSpace {
         &self,
         info: &CpuExceptionInfo,
     ) -> core::result::Result<(), ()> {
-        if let Some(func) = self.page_fault_handler.get() {
+        if let Some(func) = self.page_fault_handler {
             return func(self, info);
         }
         Err(())
     }
 
     /// Registers the page fault handler in this `VmSpace`.
-    ///
-    /// The page fault handler of a `VmSpace` can only be initialized once.
-    /// If it has been initialized before, calling this method will have no effect.
     pub fn register_page_fault_handler(
-        &self,
+        &mut self,
         func: fn(&VmSpace, &CpuExceptionInfo) -> core::result::Result<(), ()>,
     ) {
-        self.page_fault_handler.call_once(|| func);
+        self.page_fault_handler = Some(func);
     }
 
     /// Creates a reader to read data from the user space of the current task.