geph5_stats/
lib.rs

1use std::{
2    collections::BTreeMap,
3    net::{SocketAddr, UdpSocket},
4    sync::Mutex,
5};
6
7use serde::{Deserialize, Serialize};
8
9/// A single stat datapoint with statsd-style semantics and named tags.
10///
11/// This is the wire type used both for local DogStatsD emission and for
12/// shipping batches from bridges/exits to the broker over RPC.
13#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
14pub struct StatEvent {
15    pub name: String,
16    /// BTreeMap so that serialization is deterministic, which Mac authentication requires.
17    pub tags: BTreeMap<String, String>,
18    pub value: f64,
19    pub kind: StatKind,
20}
21
22#[derive(Serialize, Deserialize, Clone, Copy, Debug, PartialEq, Eq)]
23pub enum StatKind {
24    /// Summed over the flush window (statsd `|c`).
25    Counter,
26    /// Last value wins (statsd `|g`).
27    Gauge,
28    /// Timing in milliseconds; the statsd server derives count/mean/percentiles (statsd `|ms`).
29    TimerMs,
30}
31
32impl StatEvent {
33    pub fn counter(name: &str, tags: &[(&str, &str)], value: f64) -> Self {
34        Self::new(name, tags, value, StatKind::Counter)
35    }
36
37    pub fn gauge(name: &str, tags: &[(&str, &str)], value: f64) -> Self {
38        Self::new(name, tags, value, StatKind::Gauge)
39    }
40
41    pub fn timer_ms(name: &str, tags: &[(&str, &str)], value: f64) -> Self {
42        Self::new(name, tags, value, StatKind::TimerMs)
43    }
44
45    fn new(name: &str, tags: &[(&str, &str)], value: f64, kind: StatKind) -> Self {
46        Self {
47            name: sanitize(name),
48            tags: tags
49                .iter()
50                .map(|(k, v)| (sanitize(k), sanitize(v)))
51                .collect(),
52            value,
53            kind,
54        }
55    }
56
57    /// Encodes as a single DogStatsD line, e.g. `bridge_bytes:1048576|c|#pool:hk,country:CN`.
58    pub fn dogstatsd_line(&self) -> String {
59        let kind = match self.kind {
60            StatKind::Counter => "c",
61            StatKind::Gauge => "g",
62            StatKind::TimerMs => "ms",
63        };
64        let mut line = format!("{}:{}|{}", sanitize(&self.name), self.value, kind);
65        if !self.tags.is_empty() {
66            line.push_str("|#");
67            let mut first = true;
68            for (k, v) in &self.tags {
69                if !first {
70                    line.push(',');
71                }
72                first = false;
73                line.push_str(&sanitize(k));
74                line.push(':');
75                line.push_str(&sanitize(v));
76            }
77        }
78        line
79    }
80}
81
82/// Replaces characters that have meaning in the statsd line protocol.
83fn sanitize(s: &str) -> String {
84    s.replace([':', '|', '#', ',', '\n', '@', '='], "_")
85}
86
87/// Maximum payload of a single statsd datagram. Conservative for loopback.
88const MAX_DATAGRAM: usize = 1400;
89
90/// Fire-and-forget DogStatsD emitter over UDP. Send failures are silently dropped,
91/// matching statsd semantics: stats must never take down or slow the caller.
92pub struct StatsdUdpSink {
93    sock: UdpSocket,
94    dest: SocketAddr,
95}
96
97impl StatsdUdpSink {
98    pub fn new(dest: SocketAddr) -> std::io::Result<Self> {
99        let bind: SocketAddr = if dest.is_ipv4() {
100            "0.0.0.0:0".parse().unwrap()
101        } else {
102            "[::]:0".parse().unwrap()
103        };
104        let sock = UdpSocket::bind(bind)?;
105        sock.set_nonblocking(true)?;
106        Ok(Self { sock, dest })
107    }
108
109    pub fn send_one(&self, event: &StatEvent) {
110        let _ = self
111            .sock
112            .send_to(event.dogstatsd_line().as_bytes(), self.dest);
113    }
114
115    /// Sends events packed into as few datagrams as possible (newline-separated).
116    pub fn send_many<'a>(&self, events: impl IntoIterator<Item = &'a StatEvent>) {
117        let mut buf = String::new();
118        for event in events {
119            let line = event.dogstatsd_line();
120            if !buf.is_empty() && buf.len() + 1 + line.len() > MAX_DATAGRAM {
121                let _ = self.sock.send_to(buf.as_bytes(), self.dest);
122                buf.clear();
123            }
124            if !buf.is_empty() {
125                buf.push('\n');
126            }
127            buf.push_str(&line);
128        }
129        if !buf.is_empty() {
130            let _ = self.sock.send_to(buf.as_bytes(), self.dest);
131        }
132    }
133}
134
135/// Accumulates stats locally so that semi-trusted nodes (bridges, exits) can ship
136/// them to the broker in periodic batches: counters with identical name+tags are
137/// summed, gauges keep the last value, and timers are kept as individual events.
138#[derive(Default)]
139pub struct StatBatcher {
140    inner: Mutex<BatcherInner>,
141}
142
143#[derive(Default)]
144struct BatcherInner {
145    counters: BTreeMap<(String, BTreeMap<String, String>), f64>,
146    gauges: BTreeMap<(String, BTreeMap<String, String>), f64>,
147    timers: Vec<StatEvent>,
148}
149
150impl StatBatcher {
151    pub fn new() -> Self {
152        Self::default()
153    }
154
155    pub fn push(&self, event: StatEvent) {
156        let mut inner = self.inner.lock().unwrap();
157        match event.kind {
158            StatKind::Counter => {
159                *inner
160                    .counters
161                    .entry((event.name, event.tags))
162                    .or_insert(0.0) += event.value;
163            }
164            StatKind::Gauge => {
165                inner.gauges.insert((event.name, event.tags), event.value);
166            }
167            StatKind::TimerMs => inner.timers.push(event),
168        }
169    }
170
171    pub fn counter(&self, name: &str, tags: &[(&str, &str)], value: f64) {
172        self.push(StatEvent::counter(name, tags, value));
173    }
174
175    pub fn gauge(&self, name: &str, tags: &[(&str, &str)], value: f64) {
176        self.push(StatEvent::gauge(name, tags, value));
177    }
178
179    pub fn timer_ms(&self, name: &str, tags: &[(&str, &str)], value: f64) {
180        self.push(StatEvent::timer_ms(name, tags, value));
181    }
182
183    /// Takes all accumulated events, leaving the batcher empty.
184    pub fn drain(&self) -> Vec<StatEvent> {
185        let mut inner = self.inner.lock().unwrap();
186        let counters = std::mem::take(&mut inner.counters);
187        let gauges = std::mem::take(&mut inner.gauges);
188        let timers = std::mem::take(&mut inner.timers);
189        counters
190            .into_iter()
191            .map(|((name, tags), value)| StatEvent {
192                name,
193                tags,
194                value,
195                kind: StatKind::Counter,
196            })
197            .chain(gauges.into_iter().map(|((name, tags), value)| StatEvent {
198                name,
199                tags,
200                value,
201                kind: StatKind::Gauge,
202            }))
203            .chain(timers)
204            .collect()
205    }
206}
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211
212    #[test]
213    fn dogstatsd_line_formats_counter_with_tags() {
214        let event = StatEvent::counter(
215            "bridge_bytes",
216            &[("pool", "hk"), ("country", "CN")],
217            1048576.0,
218        );
219        assert_eq!(
220            event.dogstatsd_line(),
221            "bridge_bytes:1048576|c|#country:CN,pool:hk"
222        );
223    }
224
225    #[test]
226    fn dogstatsd_line_formats_untagged_gauge() {
227        let event = StatEvent::gauge("plus", &[], 17.0);
228        assert_eq!(event.dogstatsd_line(), "plus:17|g");
229    }
230
231    #[test]
232    fn dogstatsd_line_formats_timer() {
233        let event = StatEvent::timer_ms("broker_rpc_calls", &[("method", "get_exits")], 12.5);
234        assert_eq!(
235            event.dogstatsd_line(),
236            "broker_rpc_calls:12.5|ms|#method:get_exits"
237        );
238    }
239
240    #[test]
241    fn sanitize_strips_protocol_characters() {
242        let event = StatEvent::gauge("we|ird:name", &[("ta#g", "va,lue")], 1.0);
243        assert_eq!(event.dogstatsd_line(), "we_ird_name:1|g|#ta_g:va_lue");
244    }
245
246    #[test]
247    fn batcher_sums_counters_and_overwrites_gauges() {
248        let batcher = StatBatcher::new();
249        batcher.counter("bytes", &[("pool", "a")], 10.0);
250        batcher.counter("bytes", &[("pool", "a")], 5.0);
251        batcher.counter("bytes", &[("pool", "b")], 1.0);
252        batcher.gauge("load", &[], 0.5);
253        batcher.gauge("load", &[], 0.7);
254
255        let mut drained = batcher.drain();
256        drained.sort_by(|a, b| (&a.name, &a.tags).cmp(&(&b.name, &b.tags)));
257        assert_eq!(drained.len(), 3);
258        assert_eq!(drained[0].value, 15.0);
259        assert_eq!(drained[1].value, 1.0);
260        assert_eq!(drained[2].value, 0.7);
261        assert!(batcher.drain().is_empty());
262    }
263
264    #[test]
265    fn batcher_keeps_individual_timers() {
266        let batcher = StatBatcher::new();
267        batcher.timer_ms("lat", &[], 1.0);
268        batcher.timer_ms("lat", &[], 2.0);
269        assert_eq!(batcher.drain().len(), 2);
270    }
271
272    #[test]
273    fn send_many_packs_multiple_lines() {
274        let receiver = UdpSocket::bind("127.0.0.1:0").unwrap();
275        let sink = StatsdUdpSink::new(receiver.local_addr().unwrap()).unwrap();
276        sink.send_many(&[
277            StatEvent::gauge("a", &[], 1.0),
278            StatEvent::gauge("b", &[], 2.0),
279        ]);
280        let mut buf = [0u8; 2048];
281        let (n, _) = receiver.recv_from(&mut buf).unwrap();
282        assert_eq!(std::str::from_utf8(&buf[..n]).unwrap(), "a:1|g\nb:2|g");
283    }
284}
geph5_stats/lib.rs

geph5_stats/
lib.rs
