From f869ed508f2c19bc07907c46cd0efc07835a6cd6 Mon Sep 17 00:00:00 2001
From: Yang Zhichao <yzc2004.12@qq.com>
Date: Fri, 19 Sep 2025 21:43:15 +0800
Subject: [PATCH] Move `PerCpuCounter` to `aster-util` to broaden its scope of
 use, and add/modify some methods to improve readability.

---
 kernel/libs/aster-util/Cargo.toml             |  1 +
 kernel/libs/aster-util/src/lib.rs             |  1 +
 .../aster-util/src}/per_cpu_counter.rs        | 21 +++++++++++++++++--
 kernel/src/util/mod.rs                        |  1 -
 kernel/src/vm/vmar/mod.rs                     |  6 +++---
 5 files changed, 24 insertions(+), 6 deletions(-)
 rename kernel/{src/util => libs/aster-util/src}/per_cpu_counter.rs (75%)

diff --git a/kernel/libs/aster-util/Cargo.toml b/kernel/libs/aster-util/Cargo.toml
index 96623c540..e78f1479c 100644
--- a/kernel/libs/aster-util/Cargo.toml
+++ b/kernel/libs/aster-util/Cargo.toml
@@ -11,6 +11,7 @@ typeflags-util = { path = "../typeflags-util" }
 aster-rights-proc = { path = "../aster-rights-proc" }
 aster-rights = { path = "../aster-rights" }
 inherit-methods-macro = { git = "https://github.com/asterinas/inherit-methods-macro", rev = "98f7e3e" }
+osdk-heap-allocator = { path = "../../../osdk/deps/heap-allocator" }
 
 [lints]
 workspace = true
diff --git a/kernel/libs/aster-util/src/lib.rs b/kernel/libs/aster-util/src/lib.rs
index c937ada16..1aa0104a5 100644
--- a/kernel/libs/aster-util/src/lib.rs
+++ b/kernel/libs/aster-util/src/lib.rs
@@ -10,6 +10,7 @@ extern crate alloc;
 pub mod coeff;
 pub mod dup;
 pub mod mem_obj_slice;
+pub mod per_cpu_counter;
 pub mod printer;
 pub mod safe_ptr;
 pub mod slot_vec;
diff --git a/kernel/src/util/per_cpu_counter.rs b/kernel/libs/aster-util/src/per_cpu_counter.rs
similarity index 75%
rename from kernel/src/util/per_cpu_counter.rs
rename to kernel/libs/aster-util/src/per_cpu_counter.rs
index af264b29c..34c813c06 100644
--- a/kernel/src/util/per_cpu_counter.rs
+++ b/kernel/libs/aster-util/src/per_cpu_counter.rs
@@ -27,7 +27,7 @@ impl PerCpuCounter {
     }
 
     /// Adds `increment` to the counter on the given CPU.
-    pub fn add(&self, on_cpu: CpuId, increment: isize) {
+    pub fn add_on_cpu(&self, on_cpu: CpuId, increment: isize) {
         self.per_cpu_counter
             .get_on_cpu(on_cpu)
             .fetch_add(increment, Ordering::Relaxed);
@@ -37,7 +37,7 @@ impl PerCpuCounter {
     ///
     /// This function may be inaccurate since other CPUs may be
     /// updating the counter.
-    pub fn get(&self) -> usize {
+    pub fn sum_all_cpus(&self) -> usize {
         let mut total: isize = 0;
         for cpu in all_cpus() {
             total =
@@ -51,4 +51,21 @@ impl PerCpuCounter {
             total as usize
         }
     }
+
+    /// Gets the counter value on a specific CPU.
+    pub fn get_on_cpu(&self, cpu: CpuId) -> usize {
+        let val = self.per_cpu_counter.get_on_cpu(cpu).load(Ordering::Relaxed);
+        if val < 0 {
+            // See explanation in `sum_all_cpus`.
+            0
+        } else {
+            val as usize
+        }
+    }
+}
+
+impl Default for PerCpuCounter {
+    fn default() -> Self {
+        Self::new()
+    }
 }
diff --git a/kernel/src/util/mod.rs b/kernel/src/util/mod.rs
index b8b71451b..c994fba09 100644
--- a/kernel/src/util/mod.rs
+++ b/kernel/src/util/mod.rs
@@ -3,7 +3,6 @@
 mod iovec;
 pub mod net;
 mod padded;
-pub mod per_cpu_counter;
 pub mod random;
 mod read_cstring;
 pub mod ring_buffer;
diff --git a/kernel/src/vm/vmar/mod.rs b/kernel/src/vm/vmar/mod.rs
index 85d591386..6c38341ae 100644
--- a/kernel/src/vm/vmar/mod.rs
+++ b/kernel/src/vm/vmar/mod.rs
@@ -11,6 +11,7 @@ use core::{array, num::NonZeroUsize, ops::Range};
 
 use align_ext::AlignExt;
 use aster_rights::Rights;
+use aster_util::per_cpu_counter::PerCpuCounter;
 use ostd::{
     cpu::CpuId,
     mm::{
@@ -32,7 +33,6 @@ use crate::{
     prelude::*,
     process::{Process, ResourceType},
     thread::exception::PageFaultInfo,
-    util::per_cpu_counter::PerCpuCounter,
     vm::{
         perms::VmPerms,
         vmo::{Vmo, VmoRightsOp},
@@ -735,13 +735,13 @@ impl Vmar_ {
     }
 
     pub fn get_rss_counter(&self, rss_type: RssType) -> usize {
-        self.rss_counters[rss_type as usize].get()
+        self.rss_counters[rss_type as usize].sum_all_cpus()
     }
 
     fn add_rss_counter(&self, rss_type: RssType, val: isize) {
         // There are races but updating a remote counter won't cause any problems.
         let cpu_id = CpuId::current_racy();
-        self.rss_counters[rss_type as usize].add(cpu_id, val);
+        self.rss_counters[rss_type as usize].add_on_cpu(cpu_id, val);
     }
 }