plumtree

3 files changed, 51 insertions, 2 deletions

diff --git a/src/transport/mod.rs b/src/transport/mod.rs
index 743e713..52b544a 100644
--- a/src/transport/mod.rs
+++ b/src/transport/mod.rs
@@ -247,6 +247,10 @@ pub(crate) enum ParsedMessage {
     Probe(ProbeRequest),
     /// Answers a PROBE.
     ProbeReply(ProbeReply),
+    /// Asks the recipient to leave a writer's changes to announcements (`swarm::plumtree`).
+    Prune(TreeLink),
+    /// Asks the recipient to send a writer's changes at once (`swarm::plumtree`).
+    Graft(TreeLink),
 }
 
 /// A node's proof that it holds the key of its node ID (`identity`), sending from its
@@ -391,6 +395,8 @@ pub(crate) enum MessageKind {
     LinksReply,
     Probe,
     ProbeReply,
+    Prune,
+    Graft,
 }
 
 impl ParsedMessage {
@@ -410,6 +416,8 @@ impl ParsedMessage {
             ParsedMessage::LinksReply(_) => MessageKind::LinksReply,
             ParsedMessage::Probe(_) => MessageKind::Probe,
             ParsedMessage::ProbeReply(_) => MessageKind::ProbeReply,
+            ParsedMessage::Prune(_) => MessageKind::Prune,
+            ParsedMessage::Graft(_) => MessageKind::Graft,
         }
     }
 
@@ -430,6 +438,9 @@ impl ParsedMessage {
             ParsedMessage::LinksReply(reply) => Some((reply.epoch, reply.epoch_seq)),
             ParsedMessage::Probe(request) => Some((request.epoch, request.epoch_seq)),
             ParsedMessage::ProbeReply(reply) => Some((reply.epoch, reply.epoch_seq)),
+            ParsedMessage::Prune(link) | ParsedMessage::Graft(link) => {
+                Some((link.epoch, link.epoch_seq))
+            }
             ParsedMessage::Who(_)
             | ParsedMessage::Me(_)
             | ParsedMessage::Ping(_)
@@ -438,6 +449,18 @@ impl ParsedMessage {
     }
 }
 
+/// PRUNE and GRAFT: a writer's changes left to announcements, or sent at once, between
+/// the sender and the recipient (`swarm::plumtree`).
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub(crate) struct TreeLink {
+    #[serde(with = "crate::transport::wire::codec::epoch")]
+    pub(crate) epoch: Epoch,
+    /// See `SyncEnvelope::epoch_seq`.
+    pub(crate) epoch_seq: u32,
+    /// The writer, by its pair key (`topology::pair_key`).
+    pub(crate) writer: u64,
+}
+
 /// PROBE: whether the recipient reaches `member` (`swarm::working_links`), as SWIM
 /// asks others before it suspects a member. Answered from what the recipient knows,
 /// without contacting the member.
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"),
     }
 }
 
diff --git a/src/transport/wire.rs b/src/transport/wire.rs
index bf46fc9..729acd3 100644
--- a/src/transport/wire.rs
+++ b/src/transport/wire.rs
@@ -14,7 +14,7 @@ use super::{EncodedMessage, MessageKind, ParsedMessage};
 
 /// The wire format's version, the first byte of every payload. Bump it with every change
 /// to a message type or a codec (`golden_bytes` fails then).
-pub(crate) const VERSION: u8 = 36;
+pub(crate) const VERSION: u8 = 37;
 
 /// What the second byte of a payload says.
 fn tag(kind: MessageKind) -> u8 {
@@ -33,6 +33,8 @@ fn tag(kind: MessageKind) -> u8 {
         MessageKind::LinksReply => 11,
         MessageKind::Probe => 12,
         MessageKind::ProbeReply => 13,
+        MessageKind::Prune => 14,
+        MessageKind::Graft => 15,
     }
 }
 
@@ -89,6 +91,9 @@ pub(crate) fn encode(message: &ParsedMessage) -> Result<Vec<u8>, WireError> {
         ParsedMessage::LinksReply(reply) => rmp_serde::encode::write(&mut payload, reply)?,
         ParsedMessage::Probe(request) => rmp_serde::encode::write(&mut payload, request)?,
         ParsedMessage::ProbeReply(reply) => rmp_serde::encode::write(&mut payload, reply)?,
+        ParsedMessage::Prune(link) | ParsedMessage::Graft(link) => {
+            rmp_serde::encode::write(&mut payload, link)?
+        }
     }
     Ok(payload)
 }
@@ -162,6 +167,10 @@ pub(crate) fn decode(payload: &[u8]) -> Result<ParsedMessage, WireError> {
         ParsedMessage::Probe(rmp_serde::from_slice(body)?)
     } else if is(MessageKind::ProbeReply) {
         ParsedMessage::ProbeReply(rmp_serde::from_slice(body)?)
+    } else if is(MessageKind::Prune) {
+        ParsedMessage::Prune(rmp_serde::from_slice(body)?)
+    } else if is(MessageKind::Graft) {
+        ParsedMessage::Graft(rmp_serde::from_slice(body)?)
     } else {
         return Err(WireError::UnknownKind(kind));
     })
@@ -532,6 +541,16 @@ pub(crate) fn samples() -> Vec<ParsedMessage> {
             member: "q".into(),
             reached: None,
         }),
+        ParsedMessage::Prune(crate::transport::TreeLink {
+            epoch,
+            epoch_seq: 1,
+            writer: 5,
+        }),
+        ParsedMessage::Graft(crate::transport::TreeLink {
+            epoch,
+            epoch_seq: 1,
+            writer: 5,
+        }),
     ]
 }
 
@@ -563,7 +582,7 @@ mod tests {
     #[test]
     fn every_kind_round_trips() {
         let samples = samples();
-        assert_eq!(samples.len(), 14, "a kind is missing from `samples`");
+        assert_eq!(samples.len(), 16, "a kind is missing from `samples`");
         for message in samples {
             let bytes = encode(&message).unwrap();
             let decoded = decode(&bytes).unwrap();