plumtree

6 files changed, 452 insertions, 52 deletions

diff --git a/src/swarm/dissemination.rs b/src/swarm/dissemination.rs
index 1d92b06..93ad4e9 100644
--- a/src/swarm/dissemination.rs
+++ b/src/swarm/dissemination.rs
@@ -10,6 +10,7 @@
 //! first, and so how soon.
 
 use async_trait::async_trait;
+use tracing::trace;
 
 use crate::document::ChangeHash;
 use crate::state::AppState;
@@ -48,4 +49,103 @@ pub(crate) trait Dissemination: Send + Sync + 'static {
         carried: bool,
         authors: &[(String, Vec<ChangeHash>)],
     );
+
+    /// `peer` asks us to send it `writer`'s changes at once (`eager`, a GRAFT) or to
+    /// leave them to announcements (a PRUNE), for a way that keeps such links
+    /// (PlumTree). By default nothing.
+    async fn linked(&self, state: &AppState, peer: &str, writer: u64, eager: bool) {
+        let _ = (state, peer, writer, eager);
+    }
+
+    /// `peer` told us of `heads` we lack (a STATUS: an announcement, or its keepalive's),
+    /// for a way that minds which neighbour has changes first (PlumTree). By default
+    /// nothing.
+    fn announced(&self, peer: &str, heads: &[ChangeHash]) {
+        let _ = (peer, heads);
+    }
+
+    /// The neighbours whose links are lazy, for a way that keeps such links (PlumTree),
+    /// for the benchmark to report. By default none.
+    #[cfg(test)]
+    fn lazy(&self) -> Option<std::collections::HashSet<String>> {
+        None
+    }
+}
+
+/// What every way does with its `plan` for `writer`'s `changes` from `from` (none: we
+/// had them already): those it leaves to others (`Plan::covered`) count as having them
+/// (`SyncSession::known_to_peer`): else, with writers in turn, each change coming along
+/// another way, the next message we pass on to one would carry all we got since our
+/// last to it, mostly what its other parents brought it; should one lack them after
+/// all, it asks, and gets them (`resend_if_needed`). Those it sends to
+/// (`Plan::children`, but `from`) get them at once, with those of them that `changes`
+/// build on, if they were counted so lately (`SyncSession::with_predecessors`).
+pub(crate) async fn pass_on(
+    state: &AppState,
+    writer: u64,
+    from: &str,
+    changes: &[ChangeHash],
+    mut plan: Plan,
+) -> Plan {
+    if !changes.is_empty() {
+        for neighbour in &plan.covered {
+            if let Some(peer_state) = state.get_peer_state(neighbour).await {
+                peer_state
+                    .lock()
+                    .await
+                    .sync
+                    .known_to_peer(changes, tokio::time::Instant::now());
+            }
+        }
+    }
+    plan.children.retain(|child| child != from);
+    let children = &plan.children;
+    let replica = if changes.is_empty() {
+        None
+    } else {
+        Some(state.doc.read().await)
+    };
+    for child in children {
+        if let Some(peer_state) = state.get_peer_state(child).await {
+            let mut peer_state = peer_state.lock().await;
+            peer_state.tree_writer = Some(writer);
+            if let Some(replica) = &replica {
+                peer_state.sync.with_predecessors(
+                    &replica.doc,
+                    changes,
+                    tokio::time::Instant::now(),
+                );
+            }
+        }
+        state.stats.tree_forward();
+        state.request_sync(child);
+    }
+    if !children.is_empty() {
+        trace!(
+            writer,
+            children = children.len(),
+            "Passing a change on along the writer's tree"
+        );
+    }
+    plan
+}
+
+/// The ways to choose from (`--dissemination`); one per swarm.
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
+#[serde(rename_all = "kebab-case")]
+pub(crate) enum Kind {
+    /// Writers' trees (`swarm::tree`).
+    #[default]
+    WritersTrees,
+    /// PlumTree (`swarm::plumtree`).
+    Plumtree,
+}
+
+impl Kind {
+    pub(crate) fn make(self) -> std::sync::Arc<dyn Dissemination> {
+        match self {
+            Kind::WritersTrees => std::sync::Arc::new(crate::swarm::tree::WritersTrees),
+            Kind::Plumtree => std::sync::Arc::new(crate::swarm::plumtree::PlumTree::default()),
+        }
+    }
 }
diff --git a/src/swarm/mod.rs b/src/swarm/mod.rs
index 011807c..e01c6f1 100644
--- a/src/swarm/mod.rs
+++ b/src/swarm/mod.rs
@@ -12,6 +12,7 @@ pub(crate) mod links;
 pub(crate) mod members;
 pub(crate) mod pair;
 pub(crate) mod peer;
+pub(crate) mod plumtree;
 pub(crate) mod reach;
 pub(crate) mod relay;
 pub(crate) mod rtt;
diff --git a/src/swarm/peer/inbound.rs b/src/swarm/peer/inbound.rs
index 7154357..ddcc376 100644
--- a/src/swarm/peer/inbound.rs
+++ b/src/swarm/peer/inbound.rs
@@ -254,6 +254,18 @@ pub(crate) async fn handle_message(
         transport::ParsedMessage::ProbeReply(reply) => {
             crate::swarm::working_links::answered(state, &peer, reply);
         }
+        transport::ParsedMessage::Prune(link) => {
+            state
+                .dissemination
+                .linked(state, &peer, link.writer, false)
+                .await;
+        }
+        transport::ParsedMessage::Graft(link) => {
+            state
+                .dissemination
+                .linked(state, &peer, link.writer, true)
+                .await;
+        }
         transport::ParsedMessage::Who(_)
         | transport::ParsedMessage::Me(_)
         | transport::ParsedMessage::Ping(_)
diff --git a/src/swarm/plumtree.rs b/src/swarm/plumtree.rs
new file mode 100644
index 0000000..9560b9f
--- /dev/null
+++ b/src/swarm/plumtree.rs
@@ -0,0 +1,324 @@
+//! PlumTree (Leitão, Pereira and Rodrigues, "Epidemic Broadcast Trees", 2007) as the way
+//! changes reach every member (`swarm::dissemination`, `--dissemination plumtree`), over
+//! the same neighbours as writers' trees (`swarm::topology`).
+//!
+//! Each node keeps which of its neighbours it pushes changes to at once (eager) and
+//! which it leaves to announcements (lazy), one set for all writers, as in the paper:
+//! at first all are eager. A change that comes along a writer's push with nothing new
+//! (a duplicate: it came another way first) makes the receiver ask the sender to leave
+//! changes to announcements from now on (PRUNE), and both hold the link lazy. A change
+//! that comes another way (pulled after an announcement, PlumTree's IHAVE: the eager
+//! links didn't bring it) makes the receiver ask the one it came from to push at once
+//! (GRAFT), and both hold the link eager. So the eager links settle into a tree made of
+//! the links that delivered first: one that follows latency by itself, where writers'
+//! trees are computed alike on every node from the member list (`swarm::tree`). One
+//! tree for all writers settles over every writer's changes; trees per writer (tried
+//! first) never did with writers in turn, each writing too rarely, every new one
+//! flooding until pruned. PRUNE and GRAFT name the writer, for such a variant; this one
+//! doesn't use it.
+//!
+//! PlumTree's IHAVE is the announcement (`swarm::relay`): lazy neighbours, a few of them
+//! at random (`LAZY_FANOUT`), are told the new heads in the quick round, coalesced, all
+//! of them in the slow round, and count as having the changes
+//! (`Plan::covered`), as their eager links bring them. One that still lacks them a round
+//! trip later pulls (`relay::pull_after_grace`: the graft timer), and the pull, bringing
+//! a writer's changes another way, grafts the link. Told only in the slow round instead
+//! (tried second), repair took seconds; every lazy link told per round is PlumTree's
+//! cost (one STATUS each; a message of its own would only be smaller), so telling all
+//! in the quick round tripled the messages (tried third).
+//!
+//! The paper's optimisation: a lazy neighbour that told us of changes before an eager
+//! one pushed them has the faster link (an announcement waits `relay::ANNOUNCE_DELAY`,
+//! a push doesn't), so it is grafted, once it was so for a few pushes in a row
+//! (`AHEAD_TIMES`); the old parent's pushes then come as duplicates,
+//! and are pruned. A link we just grafted isn't pruned for a duplicate while what it
+//! pushed meanwhile may still be on its way (`settling`). Else a tree stays as it first
+//! formed: a pull beats a push only when
+//! the push is slower by an announcement's delay and two round trips.
+//! Nothing is written to the document.
+
+use std::collections::{HashMap, HashSet};
+use std::sync::Mutex;
+use std::time::Duration;
+
+use async_trait::async_trait;
+use tokio::time::Instant;
+use tracing::debug;
+
+use crate::document::ChangeHash;
+use crate::locks::MutexExt;
+use crate::state::AppState;
+use crate::swarm::dissemination::{Dissemination, pass_on};
+use crate::swarm::topology;
+use crate::swarm::tree::Plan;
+use crate::transport::{ParsedMessage, TreeLink};
+
+/// Lazy links kept at most; past it, all are forgotten (every link eager again, and
+/// pruned again where it delivers twice). More than a node has neighbours.
+const MAX_LAZY: usize = 4096;
+
+/// Lazy neighbours told in the quick round per change passed on, picked at random each
+/// time (`some_of`); the others hear in the slow round. Told all (patch 997), every node
+/// announced to all its lazy neighbours at every hop: about 240 STATUS per write at 50
+/// nodes, three times the messages of PlumTree told only in the slow round.
+const LAZY_FANOUT: usize = 3;
+
+/// Up to `count` of `peers`, picked at random.
+fn some_of(peers: &[String], count: usize) -> Vec<String> {
+    let mut peers = peers.to_vec();
+    let count = count.min(peers.len());
+    for picked in 0..count {
+        let rest = (peers.len() - picked) as u64;
+        let other = picked + (crate::transport::paths::random_nonce() % rest) as usize;
+        peers.swap(picked, other);
+    }
+    peers.truncate(count);
+    peers
+}
+
+/// Heads we lacked when told of them, kept at most; past it, all are forgotten (a
+/// better link is noticed with the next change).
+const MAX_ANNOUNCED: usize = 4096;
+
+/// Who first told us of a head we lacked is noticed for this long.
+const ANNOUNCED_FOR: Duration = Duration::from_secs(10);
+
+/// Pushes in a row a lazy neighbour must have told of first to be grafted (`ahead`).
+/// Grafted at once (patch 1001), trees shifted with every lucky announcement: on Tor,
+/// with one writer, 24 PRUNE and GRAFT per write against 6, and changes along the
+/// eager links took 987 ms against 471.
+const AHEAD_TIMES: u32 = 3;
+
+/// How long after we grafted a link its duplicates don't prune it, in round trips to
+/// the peer (`settling`), and when we know none.
+const SETTLE_ROUND_TRIPS: u32 = 2;
+const SETTLE_UNKNOWN: Duration = Duration::from_secs(2);
+
+/// The neighbours whose links are lazy, who first told us of heads we lacked, how many
+/// pushes in a row a lazy one told of first, and when we grafted each.
+#[derive(Default)]
+pub(crate) struct PlumTree {
+    lazy: Mutex<HashSet<String>>,
+    announced: Mutex<HashMap<ChangeHash, (String, Instant)>>,
+    ahead_times: Mutex<HashMap<String, u32>>,
+    /// When we last grafted each link (`settling`).
+    grafted: Mutex<HashMap<String, Instant>>,
+}
+
+impl PlumTree {
+    /// Holds the link to `peer` lazy, or eager; whether that changed it.
+    fn set(&self, peer: &str, lazy: bool) -> bool {
+        let mut links = self.lazy.lock_or_recover();
+        if lazy {
+            if links.len() >= MAX_LAZY {
+                links.clear();
+            }
+            links.insert(peer.to_owned())
+        } else {
+            links.remove(peer)
+        }
+    }
+
+    /// A lazy neighbour, not `peer`, that told us of any of the changes `peer` just
+    /// pushed before they came (since `ANNOUNCED_FOR`), for the last `AHEAD_TIMES`
+    /// pushes in a row: its link is the faster, by at least an announcement's delay
+    /// (`relay::ANNOUNCE_DELAY`), which a push over it doesn't wait. The changes are
+    /// forgotten as announced.
+    fn ahead(&self, peer: &str, authors: &[(String, Vec<ChangeHash>)]) -> Option<String> {
+        let now = Instant::now();
+        let told: Vec<(String, Instant)> = {
+            let mut announced = self.announced.lock_or_recover();
+            authors
+                .iter()
+                .flat_map(|(_, changes)| changes)
+                .filter_map(|change| announced.remove(change))
+                .collect()
+        };
+        let first = {
+            let lazy = self.lazy.lock_or_recover();
+            told.into_iter()
+                .filter(|(by, at)| {
+                    by != peer && lazy.contains(by) && now.duration_since(*at) < ANNOUNCED_FOR
+                })
+                .min_by_key(|(_, at)| *at)
+                .map(|(by, _)| by)
+        };
+        let mut times = self.ahead_times.lock_or_recover();
+        let Some(first) = first else {
+            times.clear();
+            return None;
+        };
+        times.retain(|by, _| *by == first);
+        let count = times.entry(first.clone()).or_default();
+        *count += 1;
+        if *count < AHEAD_TIMES {
+            return None;
+        }
+        times.clear();
+        Some(first)
+    }
+
+    /// We grafted the link to `peer` now.
+    fn grafted(&self, peer: &str) {
+        let mut grafted = self.grafted.lock_or_recover();
+        if grafted.len() >= MAX_LAZY {
+            grafted.clear();
+        }
+        grafted.insert(peer.to_owned(), Instant::now());
+    }
+
+    /// Whether we grafted the link to `peer` so lately (`SETTLE_ROUND_TRIPS` round trips
+    /// to it) that a duplicate over it doesn't prune it: what it pushed was on its way
+    /// as we pulled the same from it, or from another. Pruned at once, a node no
+    /// neighbour pushed to pulled every change, grafted the one it pulled from, and
+    /// pruned it at its next push, a duplicate of the pull, over and over (patch 1007:
+    /// on Tor, such a node got every change late).
+    fn settling(&self, state: &AppState, peer: &str) -> bool {
+        let Some(at) = self.grafted.lock_or_recover().get(peer).copied() else {
+            return false;
+        };
+        let now = Instant::now();
+        let settle = state
+            .lock_rtt()
+            .effective(peer, now, state.timing.rtt_fresh_for())
+            .map_or(SETTLE_UNKNOWN, |rtt| rtt * SETTLE_ROUND_TRIPS);
+        now.duration_since(at) < settle
+    }
+
+    /// The changes, come another way, are forgotten as announced.
+    fn forget_announced(&self, authors: &[(String, Vec<ChangeHash>)]) {
+        let mut announced = self.announced.lock_or_recover();
+        for change in authors.iter().flat_map(|(_, changes)| changes) {
+            announced.remove(change);
+        }
+    }
+
+    /// Tells `peer` to hold the link for `writer` as we now do: a GRAFT (`eager`) or a
+    /// PRUNE.
+    async fn tell(state: &AppState, peer: &str, writer: u64, eager: bool) {
+        let (epoch, epoch_seq) = state.epoch_stamp();
+        let link = TreeLink {
+            epoch,
+            epoch_seq,
+            writer,
+        };
+        let message = if eager {
+            ParsedMessage::Graft(link)
+        } else {
+            ParsedMessage::Prune(link)
+        };
+        if let Err(error) = state.transport.send_message(peer, &message).await {
+            debug!(peer = %peer, error = %format!("{error:#}"), "PRUNE or GRAFT failed");
+        }
+    }
+}
+
+#[async_trait]
+impl Dissemination for PlumTree {
+    fn name(&self) -> &'static str {
+        "plumtree"
+    }
+
+    /// Eager neighbours get the writer's changes at once (`Plan::children`); lazy ones
+    /// count as having them (`Plan::covered`), `LAZY_FANOUT` of them, picked at random,
+    /// told in the quick round (`Plan::uncovered`: the IHAVE), all in the slow round.
+    async fn plan(&self, state: &AppState, writer: u64) -> Option<Plan> {
+        let _ = writer;
+        let neighbours = topology::our_neighbours(state);
+        let lazy = self.lazy.lock_or_recover().clone();
+        let (lazy, eager): (Vec<String>, Vec<String>) =
+            neighbours.into_iter().partition(|peer| lazy.contains(peer));
+        Some(Plan {
+            children: eager,
+            alone: Vec::new(),
+            uncovered: some_of(&lazy, LAZY_FANOUT),
+            covered: lazy,
+        })
+    }
+
+    async fn forward(
+        &self,
+        state: &AppState,
+        writer: u64,
+        from: &str,
+        changes: &[ChangeHash],
+    ) -> Option<Plan> {
+        let plan = self.plan(state, writer).await?;
+        Some(pass_on(state, writer, from, changes, plan).await)
+    }
+
+    /// A duplicate along a writer's push prunes the link; changes of a writer that
+    /// came another way graft the link to whoever brought them. Told every time, not
+    /// only when our side of the link changed: a duplicate pushed over a link we hold
+    /// lazy, or changes pulled over one we hold eager, say the other end holds it the
+    /// other way (a PRUNE or GRAFT lost, or two crossing), and told only on a change,
+    /// it stayed so: a node every neighbour held lazy while it held them eager pulled
+    /// every change (patch 999: one such on Tor, a few links eager at one end). Counted as the trees'
+    /// counters count (`/admin/stats`): a duplicate as a message from another than the
+    /// parent, changes another way as ones the trees didn't bring.
+    async fn note_arrival(
+        &self,
+        state: &AppState,
+        peer: &str,
+        writer: Option<u64>,
+        carried: bool,
+        authors: &[(String, Vec<ChangeHash>)],
+    ) {
+        match writer {
+            Some(writer) if carried && authors.is_empty() => {
+                state.stats.tree_from_other();
+                if self.settling(state, peer) {
+                    return;
+                }
+                self.set(peer, true);
+                Self::tell(state, peer, writer, false).await;
+            }
+            Some(writer) if !authors.is_empty() => {
+                self.set(peer, false);
+                if let Some(ahead) = self.ahead(peer, authors) {
+                    debug!(peer = %ahead, "A lazy neighbour told of changes before they were pushed, every time lately; grafting it");
+                    self.set(&ahead, false);
+                    self.grafted(&ahead);
+                    Self::tell(state, &ahead, writer, true).await;
+                }
+            }
+            None => {
+                self.forget_announced(authors);
+                for (author, changes) in authors {
+                    if author == state.local_id() {
+                        continue;
+                    }
+                    state.stats.tree_missed(changes.len());
+                    self.set(peer, false);
+                    self.grafted(peer);
+                    Self::tell(state, peer, topology::pair_key(author), true).await;
+                }
+            }
+            Some(_) => {}
+        }
+    }
+
+    async fn linked(&self, state: &AppState, peer: &str, writer: u64, eager: bool) {
+        let _ = (state, writer);
+        self.set(peer, !eager);
+    }
+
+    /// The first to tell us of each head is the one noticed.
+    fn announced(&self, peer: &str, heads: &[ChangeHash]) {
+        let mut announced = self.announced.lock_or_recover();
+        if announced.len() + heads.len() > MAX_ANNOUNCED {
+            announced.clear();
+        }
+        let now = Instant::now();
+        for head in heads {
+            announced
+                .entry(*head)
+                .or_insert_with(|| (peer.to_owned(), now));
+        }
+    }
+
+    #[cfg(test)]
+    fn lazy(&self) -> Option<HashSet<String>> {
+        Some(self.lazy.lock_or_recover().clone())
+    }
+}
diff --git a/src/swarm/relay.rs b/src/swarm/relay.rs
index a8e3d57..bee8687 100644
--- a/src/swarm/relay.rs
+++ b/src/swarm/relay.rs
@@ -333,14 +333,19 @@ pub(crate) async fn lacks(state: &AppState, heads: &HeadsSummary) -> bool {
 /// end (`PeerState::pull_pending`).
 pub(crate) async fn pull_after_grace(state: &AppState, peer: &str, remote: &transport::Status) {
     if HeadsSummary::of(state.stamp.borrow().heads.to_vec()) != remote.heads {
-        let grace = state
-            .lock_rtt()
-            .effective(peer, Instant::now(), state.timing.rtt_fresh_for())
-            .map(|rtt| rtt.min(PULL_GRACE_MAX));
-        if let Some(grace) = grace
-            && lacks(state, &remote.heads).await
-        {
-            tokio::time::sleep(grace).await;
+        let missing = {
+            let replica = state.doc.read().await;
+            crate::document::changes::missing(&replica.doc, remote.heads.listed())
+        };
+        if !missing.is_empty() {
+            state.dissemination.announced(peer, &missing);
+            let grace = state
+                .lock_rtt()
+                .effective(peer, Instant::now(), state.timing.rtt_fresh_for())
+                .map(|rtt| rtt.min(PULL_GRACE_MAX));
+            if let Some(grace) = grace {
+                tokio::time::sleep(grace).await;
+            }
         }
     }
     let latest = match state.get_peer_state(peer).await {
diff --git a/src/swarm/tree.rs b/src/swarm/tree.rs
index de865de..db3e920 100644
--- a/src/swarm/tree.rs
+++ b/src/swarm/tree.rs
@@ -30,8 +30,6 @@
 use std::collections::{BTreeMap, BTreeSet, HashMap, VecDeque};
 use std::sync::Arc;
 
-use tracing::trace;
-
 use crate::locks::MutexExt;
 use crate::state::AppState;
 use crate::swarm::{reach, rtt, topology};
@@ -663,48 +661,8 @@ pub(crate) async fn forward(
     from: &str,
     changes: &[crate::document::ChangeHash],
 ) -> Option<Plan> {
-    let mut plan = plan(state, writer).await?;
-    if !changes.is_empty() {
-        for neighbour in &plan.covered {
-            if let Some(peer_state) = state.get_peer_state(neighbour).await {
-                peer_state
-                    .lock()
-                    .await
-                    .sync
-                    .known_to_peer(changes, tokio::time::Instant::now());
-            }
-        }
-    }
-    plan.children.retain(|child| child != from);
-    let children = &plan.children;
-    let replica = if changes.is_empty() {
-        None
-    } else {
-        Some(state.doc.read().await)
-    };
-    for child in children {
-        if let Some(peer_state) = state.get_peer_state(child).await {
-            let mut peer_state = peer_state.lock().await;
-            peer_state.tree_writer = Some(writer);
-            if let Some(replica) = &replica {
-                peer_state.sync.with_predecessors(
-                    &replica.doc,
-                    changes,
-                    tokio::time::Instant::now(),
-                );
-            }
-        }
-        state.stats.tree_forward();
-        state.request_sync(child);
-    }
-    if !children.is_empty() {
-        trace!(
-            writer,
-            children = children.len(),
-            "Passing a change on along the writer's tree"
-        );
-    }
-    Some(plan)
+    let plan = plan(state, writer).await?;
+    Some(crate::swarm::dissemination::pass_on(state, writer, from, changes, plan).await)
 }
 
 /// Writers' trees as the way changes reach every member