erbridge_core/
stats.rs

1use std::collections::VecDeque;
2use std::fmt;
3use std::fs::{File, OpenOptions};
4use std::io::Write as _;
5use std::path::PathBuf;
6use std::sync::atomic::{AtomicU64, Ordering};
7use std::sync::{Arc, Mutex};
8use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
9
10use serde::Serialize;
11use tokio::sync::broadcast;
12
13pub type ConnId = u64;
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
16#[serde(rename_all = "lowercase")]
17pub enum Protocol {
18    Tcp,
19    Udp,
20}
21
22impl fmt::Display for Protocol {
23    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
24        match self {
25            Protocol::Tcp => write!(f, "tcp"),
26            Protocol::Udp => write!(f, "udp"),
27        }
28    }
29}
30
31/// Live view of one active (or just-closed) connection, cheap to clone since the
32/// byte counters are shared atomics updated by the pump loop as data flows.
33#[derive(Debug, Clone)]
34pub struct ConnectionInfo {
35    pub id: ConnId,
36    pub label: String,
37    pub protocol: Protocol,
38    pub source: String,
39    pub destination: String,
40    pub started_at: Instant,
41    pub bytes_in: Arc<AtomicU64>,
42    pub bytes_out: Arc<AtomicU64>,
43}
44
45impl ConnectionInfo {
46    pub fn duration(&self) -> Duration {
47        self.started_at.elapsed()
48    }
49}
50
51#[derive(Debug, Clone, Serialize)]
52#[serde(tag = "kind", rename_all = "snake_case")]
53pub enum Event {
54    ConnectionOpened {
55        id: ConnId,
56        label: String,
57        protocol: Protocol,
58        source: String,
59        destination: String,
60    },
61    ConnectionClosed {
62        id: ConnId,
63        label: String,
64        protocol: Protocol,
65        source: String,
66        destination: String,
67        bytes_in: u64,
68        bytes_out: u64,
69        duration_ms: u128,
70    },
71    Info {
72        message: String,
73    },
74    Error {
75        message: String,
76    },
77}
78
79impl fmt::Display for Event {
80    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
81        match self {
82            Event::ConnectionOpened {
83                label,
84                protocol,
85                source,
86                destination,
87                ..
88            } => write!(f, "[open] {label} {protocol} {source} -> {destination}"),
89            Event::ConnectionClosed {
90                label,
91                protocol,
92                source,
93                destination,
94                bytes_in,
95                bytes_out,
96                duration_ms,
97                ..
98            } => write!(
99                f,
100                "[close] {label} {protocol} {source} -> {destination} in={bytes_in}B out={bytes_out}B dur={duration_ms}ms"
101            ),
102            Event::Info { message } => write!(f, "[info] {message}"),
103            Event::Error { message } => write!(f, "[error] {message}"),
104        }
105    }
106}
107
108#[derive(Serialize)]
109struct JsonLogLine<'a> {
110    ts_unix_ms: u128,
111    #[serde(flatten)]
112    event: &'a Event,
113}
114
115const MAX_LOG_LINES: usize = 500;
116
117/// Events a slow subscriber may fall behind by before it starts losing the
118/// oldest ones.
119const EVENT_CHANNEL_CAPACITY: usize = 256;
120
121struct Inner {
122    next_id: AtomicU64,
123    total_bytes_in: AtomicU64,
124    total_bytes_out: AtomicU64,
125    total_connections: AtomicU64,
126    live: Mutex<Vec<ConnectionInfo>>,
127    log: Mutex<VecDeque<String>>,
128    log_file: Mutex<Option<File>>,
129    events: broadcast::Sender<Event>,
130}
131
132/// Shared registry of live connections and recent events, read by the TUI and,
133/// when a log file is configured, mirrored to disk as JSON lines for headless use.
134#[derive(Clone)]
135pub struct Registry(Arc<Inner>);
136
137impl Registry {
138    pub fn new() -> Self {
139        Registry(Arc::new(Inner {
140            next_id: AtomicU64::new(1),
141            total_bytes_in: AtomicU64::new(0),
142            total_bytes_out: AtomicU64::new(0),
143            total_connections: AtomicU64::new(0),
144            live: Mutex::new(Vec::new()),
145            log: Mutex::new(VecDeque::with_capacity(MAX_LOG_LINES)),
146            log_file: Mutex::new(None),
147            events: broadcast::channel(EVENT_CHANNEL_CAPACITY).0,
148        }))
149    }
150
151    /// Subscribes to every event from now on, as structured values rather than
152    /// the formatted lines of [`Registry::recent_log`]. A subscriber that falls
153    /// more than a few hundred events behind gets `RecvError::Lagged` and
154    /// skips ahead.
155    pub fn subscribe(&self) -> broadcast::Receiver<Event> {
156        self.0.events.subscribe()
157    }
158
159    pub fn set_log_file(&self, path: PathBuf) -> anyhow::Result<()> {
160        let file = OpenOptions::new().create(true).append(true).open(path)?;
161        *self.0.log_file.lock().unwrap() = Some(file);
162        Ok(())
163    }
164
165    fn next_id(&self) -> ConnId {
166        self.0.next_id.fetch_add(1, Ordering::Relaxed)
167    }
168
169    /// Register a new connection and emit its `ConnectionOpened` event. Returns a
170    /// handle whose `bytes_in`/`bytes_out` counters the caller should update live.
171    pub fn open(
172        &self,
173        label: impl Into<String>,
174        protocol: Protocol,
175        source: impl Into<String>,
176        destination: impl Into<String>,
177    ) -> ConnectionInfo {
178        let info = ConnectionInfo {
179            id: self.next_id(),
180            label: label.into(),
181            protocol,
182            source: source.into(),
183            destination: destination.into(),
184            started_at: Instant::now(),
185            bytes_in: Arc::new(AtomicU64::new(0)),
186            bytes_out: Arc::new(AtomicU64::new(0)),
187        };
188        self.0.total_connections.fetch_add(1, Ordering::Relaxed);
189        self.0.live.lock().unwrap().push(info.clone());
190        self.push_event(Event::ConnectionOpened {
191            id: info.id,
192            label: info.label.clone(),
193            protocol: info.protocol,
194            source: info.source.clone(),
195            destination: info.destination.clone(),
196        });
197        info
198    }
199
200    pub fn close(&self, info: &ConnectionInfo) {
201        let bytes_in = info.bytes_in.load(Ordering::Relaxed);
202        let bytes_out = info.bytes_out.load(Ordering::Relaxed);
203        self.0.total_bytes_in.fetch_add(bytes_in, Ordering::Relaxed);
204        self.0
205            .total_bytes_out
206            .fetch_add(bytes_out, Ordering::Relaxed);
207        self.0.live.lock().unwrap().retain(|c| c.id != info.id);
208        self.push_event(Event::ConnectionClosed {
209            id: info.id,
210            label: info.label.clone(),
211            protocol: info.protocol,
212            source: info.source.clone(),
213            destination: info.destination.clone(),
214            bytes_in,
215            bytes_out,
216            duration_ms: info.duration().as_millis(),
217        });
218    }
219
220    pub fn info(&self, message: impl Into<String>) {
221        self.push_event(Event::Info {
222            message: message.into(),
223        });
224    }
225
226    pub fn error(&self, message: impl Into<String>) {
227        self.push_event(Event::Error {
228            message: message.into(),
229        });
230    }
231
232    fn push_event(&self, event: Event) {
233        let line = event.to_string();
234        {
235            let mut log = self.0.log.lock().unwrap();
236            if log.len() >= MAX_LOG_LINES {
237                log.pop_front();
238            }
239            log.push_back(line);
240        }
241        // `send` errs only when nobody is subscribed, which is fine.
242        let _ = self.0.events.send(event.clone());
243        if let Some(file) = self.0.log_file.lock().unwrap().as_mut() {
244            let ts_unix_ms = SystemTime::now()
245                .duration_since(UNIX_EPOCH)
246                .unwrap_or_default()
247                .as_millis();
248            let json = JsonLogLine {
249                ts_unix_ms,
250                event: &event,
251            };
252            if let Ok(mut s) = serde_json::to_string(&json) {
253                s.push('\n');
254                let _ = file.write_all(s.as_bytes());
255            }
256        }
257    }
258
259    pub fn live_connections(&self) -> Vec<ConnectionInfo> {
260        self.0.live.lock().unwrap().clone()
261    }
262
263    pub fn recent_log(&self) -> Vec<String> {
264        self.0.log.lock().unwrap().iter().cloned().collect()
265    }
266
267    pub fn totals(&self) -> Totals {
268        let live = self.0.live.lock().unwrap();
269        let live_bytes_in: u64 = live
270            .iter()
271            .map(|c| c.bytes_in.load(Ordering::Relaxed))
272            .sum();
273        let live_bytes_out: u64 = live
274            .iter()
275            .map(|c| c.bytes_out.load(Ordering::Relaxed))
276            .sum();
277        Totals {
278            live_connections: live.len(),
279            total_connections: self.0.total_connections.load(Ordering::Relaxed),
280            bytes_in: self.0.total_bytes_in.load(Ordering::Relaxed) + live_bytes_in,
281            bytes_out: self.0.total_bytes_out.load(Ordering::Relaxed) + live_bytes_out,
282        }
283    }
284}
285
286impl Default for Registry {
287    fn default() -> Self {
288        Self::new()
289    }
290}
291
292pub struct Totals {
293    pub live_connections: usize,
294    pub total_connections: u64,
295    pub bytes_in: u64,
296    pub bytes_out: u64,
297}
298
299#[cfg(test)]
300mod tests {
301    use super::*;
302
303    #[test]
304    fn subscribers_receive_events_as_they_are_emitted() {
305        let registry = Registry::new();
306        let mut events = registry.subscribe();
307
308        registry.info("hello");
309        let conn = registry.open("web", Protocol::Tcp, "1.2.3.4:5", "B:web");
310        registry.close(&conn);
311
312        assert!(matches!(events.try_recv(), Ok(Event::Info { message }) if message == "hello"));
313        assert!(matches!(
314            events.try_recv(),
315            Ok(Event::ConnectionOpened { .. })
316        ));
317        assert!(matches!(
318            events.try_recv(),
319            Ok(Event::ConnectionClosed { .. })
320        ));
321    }
322
323    #[test]
324    fn events_without_subscribers_are_not_an_error() {
325        let registry = Registry::new();
326        registry.error("nobody is listening");
327        assert_eq!(registry.recent_log().len(), 1);
328    }
329}
erbridge_core/stats.rs

erbridge_core/
stats.rs
