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
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
//! [`Ewma`] — per-key exponentially weighted moving average.
use std::{collections::HashMap, hash::Hash, time::Duration};
use crate::Timestamp;
/// Per-key exponentially weighted moving average.
///
/// `alpha ∈ (0, 1]` is the weight applied to each new sample:
///
/// ```text
/// ewma_new = alpha * sample + (1 - alpha) * ewma_old
/// ```
///
/// Smaller `alpha` → smoother / slower-reacting. `alpha = 1.0`
/// makes the EWMA equal the most recent sample.
#[derive(Debug)]
pub struct Ewma<K: Hash + Eq> {
alpha: f64,
entries: HashMap<K, (f64, Timestamp)>,
}
impl<K: Hash + Eq + Clone> Ewma<K> {
pub fn new(alpha: f64) -> Self {
assert!(alpha > 0.0 && alpha <= 1.0, "alpha must be in (0, 1]");
Self {
alpha,
entries: HashMap::new(),
}
}
/// Record `sample` for `key` at `now`. Returns the updated
/// EWMA value.
pub fn record(&mut self, key: K, sample: f64) -> f64 {
self.record_at(key, sample, Timestamp::default())
}
/// Variant of [`Self::record`] that timestamps the sample
/// (used for [`Self::evict_stale`]).
pub fn record_at(&mut self, key: K, sample: f64, now: Timestamp) -> f64 {
let entry = self.entries.entry(key).or_insert((sample, now));
entry.1 = now;
if entry.0.is_nan() {
entry.0 = sample;
} else {
entry.0 = self.alpha * sample + (1.0 - self.alpha) * entry.0;
}
entry.0
}
/// Current EWMA for `key`, or `None` if no samples have been
/// recorded.
pub fn get(&self, key: &K) -> Option<f64> {
self.entries.get(key).map(|(v, _)| *v)
}
/// Iterate `(key, ewma)` pairs.
pub fn iter(&self) -> impl Iterator<Item = (&K, f64)> {
self.entries.iter().map(|(k, (v, _))| (k, *v))
}
/// Drop entries that haven't been touched within `ttl` of
/// `now`. Useful in long-running pipelines to bound memory.
pub fn evict_stale(&mut self, now: Timestamp, ttl: Duration) {
let now_dur = now.to_duration();
self.entries
.retain(|_, (_, last_ts)| now_dur.saturating_sub(last_ts.to_duration()) <= ttl);
}
pub fn len(&self) -> usize {
self.entries.len()
}
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
}
impl<K: Hash + Eq + Clone> crate::correlate::Mergeable for Ewma<K> {
/// Per-key arithmetic-mean fold: for keys present in both,
/// the merged value is `0.5 * (a + b)` and `last_ts` is the
/// later of the two. For keys present in only one side, that
/// side's value is retained.
///
/// **Why this semantic (not weighted-by-sample-count)**:
/// `Ewma` doesn't track sample count, only the exponentially-
/// weighted value. Without sample-count, the per-key average
/// is the only commutative+associative choice. Consumers
/// who need sample-count-weighted merging should track count
/// separately and merge it alongside.
///
/// **Panics** if `alpha` doesn't match — different shards
/// running with different smoothing factors is a config bug.
fn merge(&mut self, other: Self) {
assert_eq!(
self.alpha, other.alpha,
"Ewma::merge requires matching alpha",
);
for (k, (v_other, ts_other)) in other.entries {
match self.entries.get_mut(&k) {
Some((v_self, ts_self)) => {
*v_self = 0.5 * (*v_self + v_other);
if ts_other > *ts_self {
*ts_self = ts_other;
}
}
None => {
self.entries.insert(k, (v_other, ts_other));
}
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn alpha_one_returns_last_sample() {
let mut e: Ewma<u32> = Ewma::new(1.0);
e.record(1, 10.0);
e.record(1, 5.0);
let v = e.get(&1).unwrap();
assert!((v - 5.0).abs() < 1e-9);
}
#[test]
fn alpha_half_averages_last_two() {
let mut e: Ewma<u32> = Ewma::new(0.5);
e.record(1, 10.0);
e.record(1, 20.0);
// After first record: 10.0; second: 0.5*20 + 0.5*10 = 15.
let v = e.get(&1).unwrap();
assert!((v - 15.0).abs() < 1e-9);
}
#[test]
fn per_key_isolation() {
let mut e: Ewma<u32> = Ewma::new(1.0);
e.record(1, 10.0);
e.record(2, 20.0);
assert_eq!(e.get(&1).unwrap(), 10.0);
assert_eq!(e.get(&2).unwrap(), 20.0);
}
#[test]
fn evict_stale_drops_old_entries() {
let mut e: Ewma<u32> = Ewma::new(0.5);
e.record_at(1, 1.0, Timestamp::new(0, 0));
e.record_at(2, 2.0, Timestamp::new(100, 0));
e.evict_stale(Timestamp::new(100, 0), Duration::from_secs(10));
assert!(e.get(&1).is_none());
assert!(e.get(&2).is_some());
}
}