72 lines
2.1 KiB
Rust
72 lines
2.1 KiB
Rust
// SPDX-License-Identifier: MPL-2.0
|
|
|
|
//! A fast and scalable per-CPU counter.
|
|
|
|
use core::sync::atomic::{AtomicIsize, Ordering};
|
|
|
|
use osdk_heap_allocator::{alloc_cpu_local, CpuLocalBox};
|
|
use ostd::cpu::{all_cpus, CpuId};
|
|
|
|
/// A fast, SMP-friendly, dynamically allocated, per-CPU counter.
|
|
///
|
|
/// Updating it is fast and scalable, but reading is slow and inaccurate.
|
|
///
|
|
// TODO: Reuse the code from [`osdk_frame_allocator::fast_smp_counter`],
|
|
// which may need to extract that code into a separate crate that needs
|
|
// to be published. Do that after we somehow stabilize the per-CPU counter.
|
|
pub struct PerCpuCounter {
|
|
per_cpu_counter: CpuLocalBox<AtomicIsize>,
|
|
}
|
|
|
|
impl PerCpuCounter {
|
|
/// Creates a new, zero-valued per-CPU counter.
|
|
pub fn new() -> Self {
|
|
Self {
|
|
per_cpu_counter: alloc_cpu_local(|_| AtomicIsize::new(0)).unwrap(),
|
|
}
|
|
}
|
|
|
|
/// Adds `increment` to the counter on the given CPU.
|
|
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);
|
|
}
|
|
|
|
/// Gets the total counter value.
|
|
///
|
|
/// This function may be inaccurate since other CPUs may be
|
|
/// updating the counter.
|
|
pub fn sum_all_cpus(&self) -> usize {
|
|
let mut total: isize = 0;
|
|
for cpu in all_cpus() {
|
|
total =
|
|
total.wrapping_add(self.per_cpu_counter.get_on_cpu(cpu).load(Ordering::Relaxed));
|
|
}
|
|
if total < 0 {
|
|
// The counter is unsigned. But an observer may see a negative
|
|
// value due to race conditions. We return zero if it happens.
|
|
0
|
|
} else {
|
|
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()
|
|
}
|
|
}
|