prometrics/
atomic.rs

1// so the API is "complete" even if not all functions are used
2#![allow(dead_code)]
3
4use std::sync::atomic::{self, Ordering::Relaxed};
5
6#[derive(Debug)]
7pub struct AtomicU64(atomic::AtomicU64);
8
9impl AtomicU64 {
10    pub fn new(v: u64) -> Self {
11        AtomicU64(v.into())
12    }
13
14    pub fn get(&self) -> u64 {
15        self.0.load(Relaxed)
16    }
17
18    pub fn inc(&self) {
19        self.add(1);
20    }
21
22    pub fn add(&self, v: u64) {
23        self.0.fetch_add(v, Relaxed);
24    }
25
26    pub fn update<F>(&self, f: F)
27    where
28        F: Fn(u64) -> u64,
29    {
30        let mut old = self.0.load(Relaxed);
31        loop {
32            let new = f(old);
33            match self.0.compare_exchange_weak(old, new, Relaxed, Relaxed) {
34                Ok(_) => break,
35                Err(v) => old = v, // try again
36            }
37        }
38    }
39
40    pub fn set(&self, v: u64) {
41        self.0.store(v, Relaxed);
42    }
43}
44
45#[derive(Debug)]
46pub struct AtomicI64(atomic::AtomicI64);
47
48impl AtomicI64 {
49    pub fn new(v: i64) -> Self {
50        AtomicI64(v.into())
51    }
52
53    pub fn get(&self) -> i64 {
54        self.0.load(Relaxed)
55    }
56
57    pub fn inc(&self) {
58        self.add(1);
59    }
60
61    pub fn add(&self, v: i64) {
62        self.0.fetch_add(v, Relaxed);
63    }
64
65    pub fn update<F>(&self, f: F)
66    where
67        F: Fn(i64) -> i64,
68    {
69        let mut old = self.0.load(Relaxed);
70        loop {
71            let new = f(old);
72            match self.0.compare_exchange_weak(old, new, Relaxed, Relaxed) {
73                Ok(_) => break,
74                Err(v) => old = v, // try again
75            }
76        }
77    }
78
79    pub fn set(&self, v: i64) {
80        self.0.store(v, Relaxed);
81    }
82}
83
84/// Add (and inc) is not a dedicated atomic instruction, use busy-loop
85#[derive(Debug)]
86pub struct AtomicF64(atomic::AtomicU64);
87
88impl AtomicF64 {
89    pub fn new(v: f64) -> Self {
90        AtomicF64(v.to_bits().into())
91    }
92
93    pub fn get(&self) -> f64 {
94        f64::from_bits(self.0.load(Relaxed))
95    }
96
97    pub fn inc(&self) {
98        self.add(1.0);
99    }
100
101    pub fn add(&self, v: f64) {
102        self.update(|old| old + v);
103    }
104
105    pub fn update<F>(&self, f: F)
106    where
107        F: Fn(f64) -> f64,
108    {
109        let mut old = self.0.load(Relaxed);
110        loop {
111            let new = f(f64::from_bits(old)).to_bits();
112            match self.0.compare_exchange_weak(old, new, Relaxed, Relaxed) {
113                Ok(_) => break,
114                Err(v) => old = v, // try again
115            }
116        }
117    }
118
119    pub fn set(&self, v: f64) {
120        self.0.store(v.to_bits(), Relaxed);
121    }
122}
123
124#[cfg(test)]
125mod test {
126    use super::*;
127
128    #[test]
129    fn atomic_f64_works() {
130        let value = AtomicF64::new(0.0);
131        assert_eq!(value.get(), 0.0);
132
133        value.set(123456789.0);
134        assert_eq!(value.get(), 123456789.0);
135
136        value.update(|v| v + 1.0);
137        assert_eq!(value.get(), 123456790.0);
138    }
139}