add per-node traffic counts

1 files changed, 52 insertions, 0 deletions

diff --git a/src/transport/traffic.rs b/src/transport/traffic.rs
index bf9438f..b7907fb 100644
--- a/src/transport/traffic.rs
+++ b/src/transport/traffic.rs
@@ -30,6 +30,11 @@ impl Counter {
         self.messages += 1;
         self.bytes += bytes as u64;
     }
+
+    fn absorb(&mut self, other: Counter) {
+        self.messages += other.messages;
+        self.bytes += other.bytes;
+    }
 }
 
 #[derive(Debug, Default, Clone, Copy, Serialize)]
@@ -51,6 +56,37 @@ pub(crate) struct ByKind {
 }
 
 impl ByKind {
+    fn absorb(&mut self, other: &ByKind) {
+        for kind in [
+            MessageKind::Hello,
+            MessageKind::HelloReply,
+            MessageKind::Sync,
+            MessageKind::Status,
+            MessageKind::StatusReply,
+            MessageKind::Blob,
+            MessageKind::Who,
+            MessageKind::Ping,
+            MessageKind::Links,
+        ] {
+            let theirs = *other.get(kind);
+            self.get_mut(kind).absorb(theirs);
+        }
+    }
+
+    fn get(&self, kind: MessageKind) -> &Counter {
+        match kind {
+            MessageKind::Hello => &self.hello,
+            MessageKind::HelloReply => &self.hello_reply,
+            MessageKind::Sync => &self.sync,
+            MessageKind::Status => &self.status,
+            MessageKind::StatusReply => &self.status_reply,
+            MessageKind::Blob => &self.blob,
+            MessageKind::Who | MessageKind::Me => &self.identity,
+            MessageKind::Ping | MessageKind::Pong => &self.paths,
+            MessageKind::Links | MessageKind::LinksReply => &self.links,
+        }
+    }
+
     fn get_mut(&mut self, kind: MessageKind) -> &mut Counter {
         match kind {
             MessageKind::Hello => &mut self.hello,
@@ -282,6 +318,22 @@ impl Traffic {
         self.lock().peers.get(peer).copied()
     }
 
+    /// A node's traffic: counted under it, and under `addresses`, where it is (the
+    /// layers below the node count what they send and receive on their own there: WHO
+    /// and ME, the paths layer's PINGs and PONGs).
+    pub(crate) fn node(&self, node: &str, addresses: &[String]) -> PeerTraffic {
+        let table = self.lock();
+        let mut traffic = table.peers.get(node).copied().unwrap_or_default();
+        for address in addresses.iter().filter(|address| address.as_str() != node) {
+            if let Some(there) = table.peers.get(address) {
+                traffic.sent.absorb(&there.sent);
+                traffic.received.absorb(&there.received);
+                traffic.refused.absorb(&there.refused);
+            }
+        }
+        traffic
+    }
+
     pub(crate) fn total(&self) -> PeerTraffic {
         self.lock().total
     }