plumtree

1 files changed, 7 insertions, 0 deletions

diff --git a/src/transport/traffic.rs b/src/transport/traffic.rs
index 2e4266c..501c77b 100644
--- a/src/transport/traffic.rs
+++ b/src/transport/traffic.rs
@@ -55,6 +55,8 @@ pub(crate) struct ByKind {
     pub(crate) links: Counter,
     /// PROBE and PROBE_REPLY (`swarm::working_links`).
     pub(crate) probe: Counter,
+    /// PRUNE and GRAFT (`swarm::plumtree`).
+    pub(crate) plumtree: Counter,
 }
 
 impl ByKind {
@@ -70,6 +72,7 @@ impl ByKind {
             MessageKind::Ping,
             MessageKind::Links,
             MessageKind::Probe,
+            MessageKind::Prune,
         ] {
             let theirs = *other.get(kind);
             self.get_mut(kind).absorb(theirs);
@@ -88,6 +91,7 @@ impl ByKind {
             MessageKind::Ping | MessageKind::Pong => &self.paths,
             MessageKind::Links | MessageKind::LinksReply => &self.links,
             MessageKind::Probe | MessageKind::ProbeReply => &self.probe,
+            MessageKind::Prune | MessageKind::Graft => &self.plumtree,
         }
     }
 
@@ -103,6 +107,7 @@ impl ByKind {
             MessageKind::Ping | MessageKind::Pong => &mut self.paths,
             MessageKind::Links | MessageKind::LinksReply => &mut self.links,
             MessageKind::Probe | MessageKind::ProbeReply => &mut self.probe,
+            MessageKind::Prune | MessageKind::Graft => &mut self.plumtree,
         }
     }
 }
@@ -295,6 +300,8 @@ fn names(kind: MessageKind) -> (&'static str, &'static str) {
         MessageKind::LinksReply => ("LINKS_REPLY", "links"),
         MessageKind::Probe => ("PROBE", "probe"),
         MessageKind::ProbeReply => ("PROBE_REPLY", "probe"),
+        MessageKind::Prune => ("PRUNE", "plumtree"),
+        MessageKind::Graft => ("GRAFT", "plumtree"),
     }
 }