Improve cpu stat accuracy

src/context/arch/x86_64.rs

@@ -239,11 +239,6 @@ pub unsafe fn empty_cr3() -> rmm::PhysicalAddress {
 /// Switch to the next context by restoring its stack and registers
 pub unsafe fn switch_to(prev: &mut super::Context, next: &mut super::Context) {
     unsafe {
-        // Update contexts' timestamps
-        let switch_time = crate::time::monotonic();
-        prev.cpu_time += switch_time.saturating_sub(prev.switch_time);
-        next.switch_time = switch_time;
-
         let pcr = crate::gdt::pcr();
 
         if let Some(ref stack) = next.kstack {

src/context/switch.rs

@@ -139,9 +139,6 @@ pub enum SwitchResult {
 pub fn switch(token: &mut CleanLockToken) -> SwitchResult {
     let percpu = PercpuBlock::current();
     cpu_stats::add_context_switch();
-    percpu
-        .stats
-        .add_time(percpu.switch_internals.pit_ticks.get());
 
     //set PIT Interrupt counter to 0, giving each process same amount of PIT ticks
     percpu.switch_internals.pit_ticks.set(0);
@@ -218,6 +215,13 @@ pub fn switch(token: &mut CleanLockToken) -> SwitchResult {
         }
     };
 
+    // Update per-cpu times
+    let switch_time = crate::time::monotonic();
+    let percpu_nanos = switch_time.saturating_sub(percpu.switch_internals.switch_time.get()) as u64;
+    let percpu_ms = percpu_nanos / 1_000_000;
+    percpu.stats.add_time(percpu_ms);
+    percpu.switch_internals.switch_time.set(switch_time);
+
     // Switch process states, TSS stack pointer, and store new context ID
     match switch_context_opt {
         Some((mut prev_context_guard, mut next_context_guard)) => {
@@ -233,7 +237,14 @@ pub fn switch(token: &mut CleanLockToken) -> SwitchResult {
             // Set the CPU ID for the next context
             next_context.cpu_id = Some(cpu_id);
 
-            let percpu = PercpuBlock::current();
+            // Update times
+            prev_context.cpu_time += switch_time.saturating_sub(prev_context.switch_time);
+            next_context.switch_time = switch_time;
+            if next_context.userspace {
+                percpu.stats.set_state(cpu_stats::CpuState::User);
+            } else {
+                percpu.stats.set_state(cpu_stats::CpuState::Kernel);
+            }
             unsafe {
                 percpu.switch_internals.set_current_context(Arc::clone(
                     ArcContextLockWriteGuard::rwlock(&next_context_guard),
@@ -294,12 +305,6 @@ pub fn switch(token: &mut CleanLockToken) -> SwitchResult {
             // need to use the `switch_finish_hook` to be able to release the locks. Newly created
             // contexts will return directly to the function pointer passed to context::spawn, and not
             // reach this code until the next context switch back.
-            if next_context.userspace {
-                percpu.stats.set_state(cpu_stats::CpuState::User);
-            } else {
-                percpu.stats.set_state(cpu_stats::CpuState::Kernel);
-            }
-
             SwitchResult::Switched
         }
         _ => {
@@ -319,6 +324,7 @@ pub fn switch(token: &mut CleanLockToken) -> SwitchResult {
 /// as well as fields required for managing ptrace sessions and signals.
 pub struct ContextSwitchPercpu {
     switch_result: Cell<Option<SwitchResultInner>>,
+    switch_time: Cell<u128>,
     pit_ticks: Cell<usize>,
 
     current_ctxt: RefCell<Option<Arc<ContextLock>>>,
@@ -333,6 +339,7 @@ impl ContextSwitchPercpu {
     pub const fn default() -> Self {
         Self {
             switch_result: Cell::new(None),
+            switch_time: Cell::new(0),
             pit_ticks: Cell::new(0),
             current_ctxt: RefCell::new(None),
             idle_ctxt: RefCell::new(None),

src/cpu_stats.rs

@@ -1,5 +1,5 @@
 use alloc::{string::String, vec::Vec};
-use core::sync::atomic::{AtomicU8, AtomicUsize, Ordering};
+use core::sync::atomic::{AtomicU8, AtomicUsize, AtomicU64, Ordering};
 
 use crate::cpu_set::LogicalCpuId;
 
@@ -28,15 +28,15 @@ pub enum CpuState {
 #[derive(Debug, Default)]
 pub struct CpuStats {
     /// Number of ticks spent on userspace contexts
-    user: AtomicUsize,
+    user: AtomicU64,
     /// Number of ticks spent on Niced userspace contexts
-    nice: AtomicUsize,
+    nice: AtomicU64,
     /// Number of ticks spent on kernel contexts
-    kernel: AtomicUsize,
+    kernel: AtomicU64,
     /// Number of ticks spent idle
-    idle: AtomicUsize,
+    idle: AtomicU64,
     /// Number of times the CPU handled an interrupt
-    irq: AtomicUsize,
+    irq: AtomicU64,
     /// Current state of the CPU
     state: AtomicU8,
 }
@@ -44,11 +44,11 @@ pub struct CpuStats {
 impl CpuStats {
     pub const fn default() -> Self {
         Self {
-            user: AtomicUsize::new(0),
-            nice: AtomicUsize::new(0),
-            kernel: AtomicUsize::new(0),
-            idle: AtomicUsize::new(0),
-            irq: AtomicUsize::new(0),
+            user: AtomicU64::new(0),
+            nice: AtomicU64::new(0),
+            kernel: AtomicU64::new(0),
+            idle: AtomicU64::new(0),
+            irq: AtomicU64::new(0),
             state: AtomicU8::new(0),
         }
     }
@@ -56,15 +56,15 @@ impl CpuStats {
 
 pub struct CpuStatsData {
     /// Number of ticks spent on userspace contexts
-    pub user: usize,
+    pub user: u64,
     /// Number of ticks spent on Niced userspace contexts
-    pub nice: usize,
+    pub nice: u64,
     /// Number of ticks spent on kernel contexts
-    pub kernel: usize,
+    pub kernel: u64,
     /// Number of ticks spent idle
-    pub idle: usize,
+    pub idle: u64,
     /// Number of times the CPU handled an interrupt
-    pub irq: usize,
+    pub irq: u64,
 }
 
 impl CpuStats {
@@ -82,13 +82,13 @@ impl CpuStats {
     /// Which statistic is incremented depends on the [`State`] of the CPU.
     ///
     /// # Parameters
-    /// * `ticks` - NUmber of ticks to add.
+    /// * `nanos` - Number of nanoseconds to add.
     #[inline]
-    pub fn add_time(&self, ticks: usize) {
+    pub fn add_time(&self, nanos: u64) {
         match self.state.load(Ordering::Relaxed) {
-            val if val == CpuState::Idle as u8 => self.idle.fetch_add(ticks, Ordering::Relaxed),
-            val if val == CpuState::User as u8 => self.user.fetch_add(ticks, Ordering::Relaxed),
-            val if val == CpuState::Kernel as u8 => self.kernel.fetch_add(ticks, Ordering::Relaxed),
+            val if val == CpuState::Idle as u8 => self.idle.fetch_add(nanos, Ordering::Relaxed),
+            val if val == CpuState::User as u8 => self.user.fetch_add(nanos, Ordering::Relaxed),
+            val if val == CpuState::Kernel as u8 => self.kernel.fetch_add(nanos, Ordering::Relaxed),
             _ => unreachable!("all possible values are covered"),
         };
     }