Count RSS when dropping RssDelta

kernel/src/vm/vmar/mod.rs

@@ -385,9 +385,8 @@ impl Vmar_ {
         if let Some(vm_mapping) = inner.vm_mappings.find_one(&address) {
             debug_assert!(vm_mapping.range().contains(&address));
 
-            let rss_increment = vm_mapping.handle_page_fault(&self.vm_space, page_fault_info)?;
+            let mut rss_delta = RssDelta::new(self);
-            self.add_rss_counter(vm_mapping.rss_type(), rss_increment as isize);
+            return vm_mapping.handle_page_fault(&self.vm_space, page_fault_info, &mut rss_delta);
-            return Ok(());
         }
 
         return_errno_with_message!(Errno::EACCES, "page fault addr is not in current vmar");
@@ -413,14 +412,13 @@ impl Vmar_ {
 
     pub fn remove_mapping(&self, range: Range<usize>) -> Result<()> {
         let mut inner = self.inner.write();
-        let mut rss_delta = RssDelta::new();
+        let mut rss_delta = RssDelta::new(self);
         inner.alloc_free_region_exact_truncate(
             &self.vm_space,
             range.start,
             range.len(),
             &mut rss_delta,
         )?;
-        self.add_rss_delta(rss_delta);
         Ok(())
     }
 
@@ -490,7 +488,7 @@ impl Vmar_ {
             let mut new_cursor = new_vmspace.cursor_mut(&preempt_guard, &range).unwrap();
             let cur_vmspace = self.vm_space();
             let mut cur_cursor = cur_vmspace.cursor_mut(&preempt_guard, &range).unwrap();
-            let mut rss_delta = RssDelta::new();
+            let mut rss_delta = RssDelta::new(&new_vmar_);
 
             for vm_mapping in inner.vm_mappings.iter() {
                 let base = vm_mapping.map_to_addr();
@@ -508,7 +506,6 @@ impl Vmar_ {
 
                 rss_delta.add(vm_mapping.rss_type(), num_copied as isize);
             }
-            new_vmar_.add_rss_delta(rss_delta);
 
             cur_cursor.flusher().issue_tlb_flush(TlbFlushOp::All);
             cur_cursor.flusher().dispatch_tlb_flush();
@@ -527,14 +524,6 @@ impl Vmar_ {
         let cpu_id = CpuId::current_racy();
         self.rss_counters[rss_type as usize].add(cpu_id, val);
     }
-
-    fn add_rss_delta(&self, rss_delta: RssDelta) {
-        for i in 0..NUM_RSS_COUNTERS {
-            let rss_type = RssType::try_from(i).unwrap();
-            let delta = rss_delta.get(rss_type);
-            self.add_rss_counter(rss_type, delta);
-        }
-    }
 }
 
 /// Sets mappings in the source page table as read-only to trigger COW, and
@@ -816,14 +805,13 @@ where
                 Errno::EINVAL,
                 "offset cannot be None since can overwrite is set",
             ))?;
-            let mut rss_delta = RssDelta::new();
+            let mut rss_delta = RssDelta::new(&parent.0);
             inner.alloc_free_region_exact_truncate(
                 parent.vm_space(),
                 offset,
                 map_size,
                 &mut rss_delta,
             )?;
-            parent.0.add_rss_delta(rss_delta);
             offset
         } else if let Some(offset) = offset {
             inner.alloc_free_region_exact(offset, map_size)?;
@@ -918,7 +906,7 @@ pub fn get_intersected_range(range1: &Range<usize>, range2: &Range<usize>) -> Ra
 /// The type representing categories of Resident Set Size (RSS).
 ///
 /// See <https://github.com/torvalds/linux/blob/fac04efc5c793dccbd07e2d59af9f90b7fc0dca4/include/linux/mm_types_task.h#L26..L32>
-#[repr(usize)]
+#[repr(u32)]
 #[expect(non_camel_case_types)]
 #[derive(Debug, Clone, Copy, TryFromInt)]
 pub enum RssType {
@@ -928,19 +916,35 @@ pub enum RssType {
 
 const NUM_RSS_COUNTERS: usize = 2;
 
-struct RssDelta([isize; NUM_RSS_COUNTERS]);
+pub(super) struct RssDelta<'a> {
+    delta: [isize; NUM_RSS_COUNTERS],
+    operated_vmar: &'a Vmar_,
+}
 
-impl RssDelta {
+impl<'a> RssDelta<'a> {
-    pub(self) fn new() -> Self {
+    pub(self) fn new(operated_vmar: &'a Vmar_) -> Self {
-        Self([0; NUM_RSS_COUNTERS])
+        Self {
+            delta: [0; NUM_RSS_COUNTERS],
+            operated_vmar,
+        }
     }
 
     pub(self) fn add(&mut self, rss_type: RssType, increment: isize) {
-        self.0[rss_type as usize] += increment;
+        self.delta[rss_type as usize] += increment;
     }
 
-    pub(self) fn get(&self, rss_type: RssType) -> isize {
+    fn get(&self, rss_type: RssType) -> isize {
-        self.0[rss_type as usize]
+        self.delta[rss_type as usize]
+    }
+}
+
+impl Drop for RssDelta<'_> {
+    fn drop(&mut self) {
+        for i in 0..NUM_RSS_COUNTERS {
+            let rss_type = RssType::try_from(i as u32).unwrap();
+            let delta = self.get(rss_type);
+            self.operated_vmar.add_rss_counter(rss_type, delta);
+        }
     }
 }
 

kernel/src/vm/vmar/vm_mapping.rs

@@ -14,7 +14,7 @@ use ostd::{
     task::disable_preempt,
 };
 
-use super::{interval_set::Interval, RssType};
+use super::{interval_set::Interval, RssDelta, RssType};
 use crate::{
     fs::utils::Inode,
     prelude::*,
@@ -151,12 +151,13 @@ impl VmMapping {
 /****************************** Page faults **********************************/
 
 impl VmMapping {
-    /// Handles a page fault and returns the number of pages mapped.
+    /// Handles a page fault.
-    pub fn handle_page_fault(
+    pub(super) fn handle_page_fault(
         &self,
         vm_space: &VmSpace,
         page_fault_info: &PageFaultInfo,
-    ) -> Result<usize> {
+        rss_delta: &mut RssDelta,
+    ) -> Result<()> {
         if !self.perms.contains(page_fault_info.required_perms) {
             trace!(
                 "self.perms {:?}, page_fault_info.required_perms {:?}, self.range {:?}",
@@ -173,7 +174,7 @@ impl VmMapping {
         let is_write = page_fault_info.required_perms.contains(VmPerms::WRITE);
 
         if !is_write && self.vmo.is_some() && self.handle_page_faults_around {
-            let (rss_increment, res) = self.handle_page_faults_around(vm_space, address);
+            let res = self.handle_page_faults_around(vm_space, address, rss_delta);
 
             // Errors caused by the "around" pages should be ignored, so here we
             // only return the error if the faulting page is still not mapped.
@@ -184,14 +185,13 @@ impl VmMapping {
                     &(page_aligned_addr..page_aligned_addr + PAGE_SIZE),
                 )?;
                 if let (_, Some((_, _))) = cursor.query().unwrap() {
-                    return Ok(rss_increment);
+                    return Ok(());
                 }
             }
 
-            return res.map(|_| rss_increment);
+            return res;
         }
 
-        let mut rss_increment: usize = 0;
         'retry: loop {
             let preempt_guard = disable_preempt();
             let mut cursor = vm_space.cursor_mut(
@@ -206,14 +206,14 @@ impl VmMapping {
                         // The page fault is already handled maybe by other threads.
                         // Just flush the TLB and return.
                         TlbFlushOp::Range(va).perform_on_current();
-                        return Ok(0);
+                        return Ok(());
                     }
                     assert!(is_write);
                     // Perform COW if it is a write access to a shared mapping.
 
                     // Skip if the page fault is already handled.
                     if prop.flags.contains(PageFlags::W) {
-                        return Ok(0);
+                        return Ok(());
                     }
 
                     // If the forked child or parent immediately unmaps the page after
@@ -233,7 +233,7 @@ impl VmMapping {
                         let new_frame = duplicate_frame(&frame)?;
                         prop.flags |= new_flags;
                         cursor.map(new_frame.into(), prop);
-                        rss_increment += 1;
+                        rss_delta.add(self.rss_type(), 1);
                     }
                     cursor.flusher().sync_tlb_flush();
                 }
@@ -270,13 +270,13 @@ impl VmMapping {
                     let map_prop = PageProperty::new_user(page_flags, CachePolicy::Writeback);
 
                     cursor.map(frame, map_prop);
-                    rss_increment += 1;
+                    rss_delta.add(self.rss_type(), 1);
                 }
             }
             break 'retry;
         }
 
-        Ok(rss_increment)
+        Ok(())
     }
 
     fn prepare_page(
@@ -310,14 +310,12 @@ impl VmMapping {
     }
 
     /// Handles a page fault and maps additional surrounding pages.
-    ///
-    /// Returns a tuple `(mapped_pages, result)`, where `mapped_pages` is the number
-    /// of pages mapped successfully, even if the `result` is some error.
     fn handle_page_faults_around(
         &self,
         vm_space: &VmSpace,
         page_fault_addr: Vaddr,
-    ) -> (usize, Result<()>) {
+        mut rss_delta: &mut RssDelta,
+    ) -> Result<()> {
         const SURROUNDING_PAGE_NUM: usize = 16;
         const SURROUNDING_PAGE_ADDR_MASK: usize = !(SURROUNDING_PAGE_NUM * PAGE_SIZE - 1);
 
@@ -332,19 +330,12 @@ impl VmMapping {
         );
 
         let vm_perms = self.perms - VmPerms::WRITE;
-        let mut rss_increment: usize = 0;
 
         'retry: loop {
             let preempt_guard = disable_preempt();
+            let mut cursor = vm_space.cursor_mut(&preempt_guard, &(start_addr..end_addr))?;
 
-            let mut cursor = match vm_space.cursor_mut(&preempt_guard, &(start_addr..end_addr)) {
+            let rss_delta_ref = &mut rss_delta;
-                Ok(cursor) => cursor,
-                Err(e) => {
-                    return (rss_increment, Err(e.into()));
-                }
-            };
-
-            let rss_increment_ref = &mut rss_increment;
             let operate =
                 move |commit_fn: &mut dyn FnMut()
                     -> core::result::Result<UFrame, VmoCommitError>| {
@@ -356,7 +347,7 @@ impl VmMapping {
                         let page_prop = PageProperty::new_user(page_flags, CachePolicy::Writeback);
                         let frame = commit_fn()?;
                         cursor.map(frame, page_prop);
-                        *rss_increment_ref += 1;
+                        rss_delta_ref.add(self.rss_type(), 1);
                     } else {
                         let next_addr = cursor.virt_addr() + PAGE_SIZE;
                         if next_addr < end_addr {
@@ -369,16 +360,14 @@ impl VmMapping {
             let start_offset = start_addr - self.map_to_addr;
             let end_offset = end_addr - self.map_to_addr;
             match vmo.try_operate_on_range(&(start_offset..end_offset), operate) {
-                Ok(_) => return (rss_increment, Ok(())),
+                Ok(_) => return Ok(()),
                 Err(VmoCommitError::NeedIo(index)) => {
                     drop(preempt_guard);
-                    if let Err(e) = vmo.commit_on(index, CommitFlags::empty()) {
+                    vmo.commit_on(index, CommitFlags::empty())?;
-                        return (rss_increment, Err(e));
-                    }
                     start_addr = index * PAGE_SIZE + self.map_to_addr;
                     continue 'retry;
                 }
-                Err(VmoCommitError::Err(e)) => return (rss_increment, Err(e)),
+                Err(VmoCommitError::Err(e)) => return Err(e),
             }
         }
     }