diff --git a/src/rust/lqos_utils/src/units/atomic_down_up.rs b/src/rust/lqos_utils/src/units/atomic_down_up.rs index e7a37380..7289ec40 100644 --- a/src/rust/lqos_utils/src/units/atomic_down_up.rs +++ b/src/rust/lqos_utils/src/units/atomic_down_up.rs @@ -1,7 +1,17 @@ +//! AtomicDownUp is a struct that contains two atomic u64 values, one for down and one for up. +//! We frequently order things down and then up in kernel maps, keeping the ordering explicit +//! helps reduce directional confusion/bugs. + use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering::Relaxed; use crate::units::DownUpOrder; +/// AtomicDownUp is a struct that contains two atomic u64 values, one for down and one for up. +/// It's typically used for throughput, but can be used for any pairing that needs to keep track +/// of values by direction. +/// +/// Note that unlike the DownUpOrder struct, it is not intended for direct serialization, and +/// is not generic. #[derive(Debug)] pub struct AtomicDownUp { down: AtomicU64, @@ -9,6 +19,7 @@ pub struct AtomicDownUp { } impl AtomicDownUp { + /// Create a new `AtomicDownUp` with both values set to zero. pub const fn zeroed() -> Self { Self { down: AtomicU64::new(0), @@ -16,11 +27,14 @@ impl AtomicDownUp { } } + /// Set both down and up to zero. pub fn set_to_zero(&self) { self.up.store(0, Relaxed); self.down.store(0, Relaxed); } + /// Add a tuple of u64 values to the down and up values. The addition + /// is checked, and will not occur if it would result in an overflow. pub fn checked_add_tuple(&self, n: (u64, u64)) { let n0 = self.down.load(std::sync::atomic::Ordering::Relaxed); if let Some(n) = n0.checked_add(n.0) { @@ -32,7 +46,9 @@ impl AtomicDownUp { self.up.store(n, std::sync::atomic::Ordering::Relaxed); } } - + + /// Add a DownUpOrder to the down and up values. The addition + /// is checked, and will not occur if it would result in an overflow. pub fn checked_add(&self, n: DownUpOrder) { let n0 = self.down.load(std::sync::atomic::Ordering::Relaxed); if let Some(n) = n0.checked_add(n.down) { @@ -45,26 +61,60 @@ impl AtomicDownUp { } } + /// Get the down value. pub fn get_down(&self) -> u64 { self.down.load(Relaxed) } + /// Get the up value. pub fn get_up(&self) -> u64 { self.up.load(Relaxed) } + /// Set the down value. pub fn set_down(&self, n: u64) { self.down.store(n, Relaxed); } + /// Set the up value. pub fn set_up(&self, n: u64) { self.up.store(n, Relaxed); } + /// Transform the AtomicDownUp into a `DownUpOrder`. pub fn as_down_up(&self) -> DownUpOrder { DownUpOrder::new( self.get_down(), self.get_up() ) } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_atomic_down_up() { + let adu = AtomicDownUp::zeroed(); + assert_eq!(adu.get_down(), 0); + assert_eq!(adu.get_up(), 0); + + adu.set_down(1); + adu.set_up(2); + assert_eq!(adu.get_down(), 1); + assert_eq!(adu.get_up(), 2); + + adu.checked_add(DownUpOrder::new(1, 2)); + assert_eq!(adu.get_down(), 2); + assert_eq!(adu.get_up(), 4); + + adu.checked_add_tuple((1, 2)); + assert_eq!(adu.get_down(), 3); + assert_eq!(adu.get_up(), 6); + + adu.set_to_zero(); + assert_eq!(adu.get_down(), 0); + assert_eq!(adu.get_up(), 0); + } } \ No newline at end of file