Inject post_schedule/user_page_fault handler to OSTD

kernel/src/lib.rs

@@ -92,6 +92,7 @@ pub fn main() {
 }
 
 pub fn init() {
+    thread::init();
     util::random::init();
     driver::init();
     time::init();

kernel/src/thread/exception.rs

@@ -3,9 +3,10 @@
 #![expect(unused_variables)]
 
 use aster_rights::Full;
-use ostd::{cpu::*, mm::VmSpace, task::Task};
+use ostd::cpu::*;
 
 use crate::{
+    current_userspace,
     prelude::*,
     process::signal::signals::fault::FaultSignal,
     vm::{page_fault_handler::PageFaultHandler, perms::VmPerms, vmar::Vmar},
@@ -41,26 +42,8 @@ pub fn handle_exception(ctx: &Context, context: &UserContext) {
     generate_fault_signal(trap_info, ctx);
 }
 
-/// Handles the page fault occurs in the input `VmSpace`.
-pub(crate) fn handle_page_fault_from_vm_space(
-    vm_space: &VmSpace,
-    page_fault_info: &PageFaultInfo,
-) -> core::result::Result<(), ()> {
-    let task = Task::current().unwrap();
-    let current_root_vmar = task.as_thread_local().unwrap().root_vmar().borrow();
-    let root_vmar = current_root_vmar.as_ref().unwrap();
-
-    // If page is not present or due to write access, we should ask the vmar try to commit this page
-    debug_assert_eq!(
-        Arc::as_ptr(root_vmar.vm_space()),
-        vm_space as *const VmSpace
-    );
-
-    handle_page_fault_from_vmar(root_vmar, page_fault_info)
-}
-
 /// Handles the page fault occurs in the input `Vmar`.
-pub(crate) fn handle_page_fault_from_vmar(
+fn handle_page_fault_from_vmar(
     root_vmar: &Vmar<Full>,
     page_fault_info: &PageFaultInfo,
 ) -> core::result::Result<(), ()> {
@@ -92,3 +75,7 @@ fn log_trap_info(trap_info: &CpuExceptionInfo) {
         trace!("[Trap][{exception:?}][err = {}]", trap_info.error_code)
     }
 }
+
+pub(super) fn page_fault_handler(info: &CpuExceptionInfo) -> core::result::Result<(), ()> {
+    handle_page_fault_from_vmar(current_userspace!().root_vmar(), &info.try_into().unwrap())
+}

kernel/src/thread/mod.rs

@@ -24,6 +24,23 @@ pub mod work_queue;
 
 pub type Tid = u32;
 
+fn post_schedule_handler() {
+    let task = Task::current().unwrap();
+    let Some(thread_local) = task.as_thread_local() else {
+        return;
+    };
+
+    let root_vmar = thread_local.root_vmar().borrow();
+    if let Some(vmar) = root_vmar.as_ref() {
+        vmar.vm_space().activate()
+    }
+}
+
+pub(super) fn init() {
+    ostd::task::inject_post_schedule_handler(post_schedule_handler);
+    ostd::arch::trap::inject_user_page_fault_handler(exception::page_fault_handler);
+}
+
 /// A thread is a wrapper on top of task.
 #[derive(Debug)]
 pub struct Thread {

kernel/src/vm/vmar/mod.rs

@@ -11,10 +11,7 @@ use core::{num::NonZeroUsize, ops::Range};
 
 use align_ext::AlignExt;
 use aster_rights::Rights;
-use ostd::{
+use ostd::mm::{tlb::TlbFlushOp, PageFlags, PageProperty, VmSpace, MAX_USERSPACE_VADDR};
-    cpu::CpuExceptionInfo,
-    mm::{tlb::TlbFlushOp, PageFlags, PageProperty, VmSpace, MAX_USERSPACE_VADDR},
-};
 
 use self::{
     interval_set::{Interval, IntervalSet},
@@ -24,7 +21,7 @@ use super::page_fault_handler::PageFaultHandler;
 use crate::{
     prelude::*,
     process::{Process, ResourceType},
-    thread::exception::{handle_page_fault_from_vm_space, PageFaultInfo},
+    thread::exception::PageFaultInfo,
     vm::{
         perms::VmPerms,
         vmo::{Vmo, VmoRightsOp},
@@ -305,7 +302,7 @@ impl Vmar_ {
 
     fn new_root() -> Arc<Self> {
         let vmar_inner = VmarInner::new();
-        let mut vm_space = VmSpace::new();
+        let vm_space = VmSpace::new();
         Vmar_::new(vmar_inner, Arc::new(vm_space), 0, ROOT_VMAR_CAP_ADDR)
     }
 
@@ -429,7 +426,7 @@ impl Vmar_ {
     pub(super) fn new_fork_root(self: &Arc<Self>) -> Result<Arc<Self>> {
         let new_vmar_ = {
             let vmar_inner = VmarInner::new();
-            let mut new_space = VmSpace::new();
+            let new_space = VmSpace::new();
             Vmar_::new(vmar_inner, Arc::new(new_space), self.base, self.size)
         };
 
@@ -466,14 +463,6 @@ impl Vmar_ {
     }
 }
 
-/// This is for fallible user space write handling.
-fn handle_page_fault_wrapper(
-    vm_space: &VmSpace,
-    trap_info: &CpuExceptionInfo,
-) -> core::result::Result<(), ()> {
-    handle_page_fault_from_vm_space(vm_space, &trap_info.try_into().unwrap())
-}
-
 impl<R> Vmar<R> {
     /// The base address, i.e., the offset relative to the root VMAR.
     ///