1 files changed, 199 insertions, 0 deletions
diff --git a/src/transport/traffic.rs b/src/transport/traffic.rs
new file mode 100644
index 0000000..bb93d48
--- /dev/null
+++ b/src/transport/traffic.rs
@@ -0,0 +1,199 @@
+//! Messages and bytes exchanged with each peer, by kind: counted where every message
+//! passes anyway, at the transport. No messages of our own.
+//!
+//! Bytes are payload bytes, what the Mycelium message carries, without its framing,
+//! the topic, or the base64 of its HTTP API. Counters only grow (from the start of the
+//! process); rates are for whoever reads them.
+//!
+//! A message dropped by the debug firewall still counts as sent (the firewall plays a
+//! lossy network) but not as received.
+
+use std::collections::HashMap;
+use std::sync::{Arc, Mutex, MutexGuard};
+
+use anyhow::Result;
+use async_trait::async_trait;
+use serde::Serialize;
+use utoipa::ToSchema;
+
+use crate::transport::{EncodedMessage, MessageHandler, MessageKind, ParsedMessage, Transport};
+
+#[derive(Debug, Default, Clone, Copy, Serialize, ToSchema)]
+pub struct Counter {
+    pub messages: u64,
+    pub bytes: u64,
+}
+
+impl Counter {
+    fn add(&mut self, bytes: usize) {
+        self.messages += 1;
+        self.bytes += bytes as u64;
+    }
+}
+
+#[derive(Debug, Default, Clone, Copy, Serialize, ToSchema)]
+#[serde(rename_all = "camelCase")]
+pub struct ByKind {
+    pub hello: Counter,
+    pub hello_reply: Counter,
+    pub sync: Counter,
+    pub ping: Counter,
+    pub pong: Counter,
+}
+
+impl ByKind {
+    fn get_mut(&mut self, kind: MessageKind) -> &mut Counter {
+        match kind {
+            MessageKind::Hello => &mut self.hello,
+            MessageKind::HelloReply => &mut self.hello_reply,
+            MessageKind::Sync => &mut self.sync,
+            MessageKind::Ping => &mut self.ping,
+            MessageKind::Pong => &mut self.pong,
+        }
+    }
+}
+
+#[derive(Debug, Default, Clone, Copy, Serialize, ToSchema)]
+#[serde(rename_all = "camelCase")]
+pub struct PeerTraffic {
+    /// Handed to Mycelium (it accepted them).
+    pub sent: ByKind,
+    pub received: ByKind,
+    /// Sends Mycelium refused (e.g. no route to the peer); not counted as sent.
+    pub refused: ByKind,
+}
+
+#[derive(Debug, Default)]
+struct Table {
+    peers: HashMap<String, PeerTraffic>,
+    /// All peers, including ones no longer listed.
+    total: PeerTraffic,
+}
+
+#[derive(Debug, Default)]
+pub struct Traffic {
+    table: Mutex<Table>,
+}
+
+impl Traffic {
+    fn lock(&self) -> MutexGuard<'_, Table> {
+        self.table
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner())
+    }
+
+    fn count(&self, peer: &str, bytes: usize, pick: impl Fn(&mut PeerTraffic) -> &mut Counter) {
+        let mut guard = self.lock();
+        let table = &mut *guard;
+        pick(&mut table.total).add(bytes);
+        match table.peers.get_mut(peer) {
+            Some(traffic) => pick(traffic).add(bytes),
+            None => pick(table.peers.entry(peer.to_owned()).or_default()).add(bytes),
+        }
+    }
+
+    fn sent(&self, peer: &str, kind: MessageKind, bytes: usize) {
+        self.count(peer, bytes, |t| t.sent.get_mut(kind));
+    }
+
+    fn refused(&self, peer: &str, kind: MessageKind, bytes: usize) {
+        self.count(peer, bytes, |t| t.refused.get_mut(kind));
+    }
+
+    fn received(&self, peer: &str, kind: MessageKind, bytes: usize) {
+        self.count(peer, bytes, |t| t.received.get_mut(kind));
+    }
+
+    pub fn peer(&self, peer: &str) -> Option<PeerTraffic> {
+        self.lock().peers.get(peer).copied()
+    }
+
+    pub fn total(&self) -> PeerTraffic {
+        self.lock().total
+    }
+
+    /// Drops the counters of peers not kept (anyone can send us messages; the totals
+    /// keep their share).
+    pub fn retain(&self, mut keep: impl FnMut(&str) -> bool) {
+        self.lock().peers.retain(|peer, _| keep(peer));
+    }
+}
+
+/// Counts what passes through `inner`.
+pub struct CountingTransport {
+    inner: Arc<dyn Transport>,
+    traffic: Arc<Traffic>,
+}
+
+impl CountingTransport {
+    pub fn new(inner: Arc<dyn Transport>, traffic: Arc<Traffic>) -> Self {
+        Self { inner, traffic }
+    }
+}
+
+#[async_trait]
+impl Transport for CountingTransport {
+    fn local_id(&self) -> &str {
+        self.inner.local_id()
+    }
+
+    async fn resolve_seed(&self, seed: &str) -> Result<String> {
+        self.inner.resolve_seed(seed).await
+    }
+
+    fn encode(&self, peer: &str, message: &ParsedMessage) -> Result<EncodedMessage> {
+        self.inner.encode(peer, message)
+    }
+
+    async fn send_encoded(&self, message: &EncodedMessage) -> Result<()> {
+        let result = self.inner.send_encoded(message).await;
+        if result.is_ok() {
+            self.traffic
+                .sent(&message.peer, message.kind, message.payload_len);
+        } else {
+            self.traffic
+                .refused(&message.peer, message.kind, message.payload_len);
+        }
+        result
+    }
+
+    async fn start_consumer(&self, handler: MessageHandler) -> Result<()> {
+        let traffic = Arc::clone(&self.traffic);
+        self.inner
+            .start_consumer(Box::new(move |sender, message, payload_len| {
+                traffic.received(&sender, message.kind(), payload_len);
+                handler(sender, message, payload_len)
+            }))
+            .await
+    }
+
+    async fn wait_for_exit(&self) -> std::io::Result<std::process::ExitStatus> {
+        self.inner.wait_for_exit().await
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn totals_survive_retain() {
+        let traffic = Traffic::default();
+        traffic.sent("a", MessageKind::Ping, 100);
+        traffic.sent("b", MessageKind::Sync, 1000);
+        traffic.received("b", MessageKind::Pong, 90);
+        traffic.refused("a", MessageKind::Sync, 500);
+
+        let b = traffic.peer("b").unwrap();
+        assert_eq!(b.sent.sync.messages, 1);
+        assert_eq!(b.received.pong.bytes, 90);
+
+        traffic.retain(|peer| peer == "b");
+        assert!(traffic.peer("a").is_none());
+
+        let total = traffic.total();
+        assert_eq!(total.sent.ping.bytes, 100);
+        assert_eq!(total.sent.sync.bytes, 1000);
+        assert_eq!(total.refused.sync.bytes, 500);
+    }
+}