1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
// Copyright 2019 Conflux Foundation. All rights reserved.
// Conflux is free software and distributed under GNU General Public License.
// See http://www.gnu.org/licenses/

use crate::metrics::ORDER;
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize};

/// EWMAs continuously calculate an exponentially-weighted moving average based
/// on an outside source of clock ticks.
pub struct EWMA {
    uncounted: AtomicUsize,
    alpha: f64,
    rate: AtomicU64,
    init: AtomicBool,
}

impl EWMA {
    /// Constructs a new EWMA for a n-minutes moving average.
    pub fn new(n: f64) -> Self {
        EWMA::new_with_alpha(1.0 - (-5.0 / 60.0 / n).exp())
    }

    fn new_with_alpha(alpha: f64) -> Self {
        EWMA {
            uncounted: AtomicUsize::new(0),
            alpha,
            rate: AtomicU64::new(0),
            init: AtomicBool::new(false),
        }
    }

    /// Rate returns the moving average rate of events per second.
    pub fn rate(&self) -> f64 { f64::from_bits(self.rate.load(ORDER)) * 1e9 }

    /// Update adds n uncounted events.
    pub fn update(&self, n: usize) { self.uncounted.fetch_add(n, ORDER); }

    /// Ticks the clock to update the moving average. It assumes it is called
    /// every 5 seconds.
    pub fn tick(&self) {
        let count = self.uncounted.swap(0, ORDER) as f64;
        let instant_rate = count / 5e9;

        if let Err(true) = self.init.compare_exchange(false, true, ORDER, ORDER)
        {
            let mut current_rate = f64::from_bits(self.rate.load(ORDER));
            current_rate += self.alpha * (instant_rate - current_rate);
            self.rate.store(f64::to_bits(current_rate), ORDER);
        } else {
            self.rate.store(f64::to_bits(instant_rate), ORDER);
        }
    }
}