add writer trees

5 files changed, 266 insertions, 27 deletions

diff --git a/src/swarm/mod.rs b/src/swarm/mod.rs
index 322170b..985e968 100644
--- a/src/swarm/mod.rs
+++ b/src/swarm/mod.rs
@@ -21,3 +21,4 @@ pub(crate) mod stats;
 pub(crate) mod sync;
 pub(crate) mod throughput;
 pub(crate) mod topology;
+pub(crate) mod tree;
diff --git a/src/swarm/relay.rs b/src/swarm/relay.rs
index 69d73d3..7dbc5b0 100644
--- a/src/swarm/relay.rs
+++ b/src/swarm/relay.rs
@@ -59,25 +59,33 @@ const PULL_TIMEOUT: Duration = Duration::from_secs(15);
 const ANNOUNCE_DELAY: Duration = Duration::from_millis(200);
 const ANNOUNCE_JITTER: Duration = Duration::from_millis(200);
 
-/// After changes from `from` were applied and our heads changed: tell our other
-/// neighbours, in the next announcement round.
-pub(crate) async fn announce(state: &AppState, from: &str) {
+/// An announcement round scheduled: for changes from these peers, and the neighbours
+/// that were sent them along a writer's tree meanwhile (`swarm::tree`).
+#[derive(Default)]
+pub(crate) struct Pending {
+    senders: HashSet<String>,
+    sent: HashSet<String>,
+}
+
+/// After changes from `from` were applied and our heads changed: passed on at once along
+/// the tree of their `writer` if the message said whose, and our other neighbours told
+/// in the next announcement round.
+pub(crate) async fn announce(state: &AppState, from: &str, writer: Option<u64>) {
     if !state.initial_sync_done.load(Ordering::Acquire) {
         return;
     }
     forward_eagerly(state, from).await;
+    let sent = match writer {
+        Some(writer) => crate::swarm::tree::forward(state, writer, from).await,
+        None => Vec::new(),
+    };
     let schedule = {
         let mut pending = state.announce_pending.lock_or_recover();
-        match pending.as_mut() {
-            Some(senders) => {
-                senders.insert(from.to_owned());
-                false
-            }
-            None => {
-                *pending = Some(HashSet::from([from.to_owned()]));
-                true
-            }
-        }
+        let schedule = pending.is_none();
+        let pending = pending.get_or_insert_with(Pending::default);
+        pending.senders.insert(from.to_owned());
+        pending.sent.extend(sent);
+        schedule
     };
     if !schedule {
         return;
@@ -89,12 +97,12 @@ pub(crate) async fn announce(state: &AppState, from: &str) {
                 .unwrap_or(1)
                 .max(1);
         tokio::time::sleep(ANNOUNCE_DELAY + Duration::from_millis(jitter)).await;
-        let senders = state
+        let pending = state
             .announce_pending
             .lock_or_recover()
             .take()
             .unwrap_or_default();
-        announce_now(&state, &senders).await;
+        announce_now(&state, &pending).await;
     });
 }
 
@@ -115,8 +123,8 @@ async fn forward_eagerly(state: &AppState, from: &str) {
 /// One announcement round. The peer whose changes we announce has them already, so a
 /// round with a single sender skips it; with several, each lacks the others' changes,
 /// so every neighbour is told.
-async fn announce_now(state: &AppState, senders: &HashSet<String>) {
-    let neighbours = unreached(state, senders).await;
+async fn announce_now(state: &AppState, pending: &Pending) {
+    let neighbours = unreached(state, &pending.senders, &pending.sent).await;
     if neighbours.is_empty() {
         return;
     }
@@ -132,8 +140,13 @@ async fn announce_now(state: &AppState, senders: &HashSet<String>) {
 /// announce has them already, so with a single sender it is skipped; with several, each
 /// lacks the others' changes. So is a neighbour whose latest heads (its STATUS, an
 /// announcement of its own, or its STATUS_REPLY) are ours now: it has all we have; and
-/// one we forward to eagerly, which was sent the changes already.
-async fn unreached(state: &AppState, senders: &HashSet<String>) -> Vec<String> {
+/// one we forward to eagerly, or `sent` the changes along a writer's tree, which has
+/// them already.
+async fn unreached(
+    state: &AppState,
+    senders: &HashSet<String>,
+    sent: &HashSet<String>,
+) -> Vec<String> {
     let skip = match senders.len() {
         1 => senders.iter().next().cloned(),
         _ => None,
@@ -144,7 +157,10 @@ async fn unreached(state: &AppState, senders: &HashSet<String>) -> Vec<String> {
     };
     let mut unreached = Vec::new();
     for (peer, peer_state) in state.peers_snapshot().await {
-        if skip.as_deref() == Some(peer.as_str()) || !topology::is_neighbour(state, &peer) {
+        if skip.as_deref() == Some(peer.as_str())
+            || sent.contains(&peer)
+            || !topology::is_neighbour(state, &peer)
+        {
             continue;
         }
         let peer_state = peer_state.lock().await;
diff --git a/src/swarm/sync.rs b/src/swarm/sync.rs
index 37c3b0a..403e0c1 100644
--- a/src/swarm/sync.rs
+++ b/src/swarm/sync.rs
@@ -144,7 +144,7 @@ async fn deliver_sync_to_peer(state: &AppState, peer: &str) -> Result<()> {
             move || -> (u64, Result<Option<transport::EncodedMessage>>) {
                 // Label the message with the epoch of the document it was generated
                 // from, read under the same lock.
-                let (session, epoch, epoch_seq, body) = {
+                let (session, epoch, epoch_seq, body, writer) = {
                     let replica = doc.blocking_read();
                     let mut peer_state = psm.blocking_lock();
                     let (epoch, epoch_seq) = (replica.epoch, replica.epoch_seq);
@@ -171,7 +171,9 @@ async fn deliver_sync_to_peer(state: &AppState, peer: &str) -> Result<()> {
                         );
                         transport::SyncBody::Delta(spec)
                     });
-                    (peer_state.sync.id(), epoch, epoch_seq, body)
+                    // Whose change it passes on, if any: this message carries it.
+                    let writer = peer_state.tree_writer.take();
+                    (peer_state.sync.id(), epoch, epoch_seq, body, writer)
                 };
                 let Some(body) = body else {
                     return (session, Ok(None));
@@ -182,6 +184,7 @@ async fn deliver_sync_to_peer(state: &AppState, peer: &str) -> Result<()> {
                     epoch_seq,
                     body,
                     ready,
+                    writer,
                 };
                 let encoded = transport
                     .encode(&peer_owned, &ParsedMessage::Sync(envelope))
@@ -341,6 +344,7 @@ pub(crate) async fn handle_sync(
 
     let epoch_label = envelope.epoch;
     let from_ready = envelope.ready;
+    let writer = envelope.writer;
     let payload = match envelope.body {
         transport::SyncBody::Message(payload) => payload,
         transport::SyncBody::Delta(spec) => {
@@ -462,7 +466,7 @@ pub(crate) async fn handle_sync(
     }
     // Caught up: our other neighbours may lack what the source had.
     if changed || just_caught_up {
-        crate::swarm::relay::announce(state, peer).await;
+        crate::swarm::relay::announce(state, peer, writer.filter(|_| changed)).await;
     }
 
     state
@@ -490,7 +494,7 @@ fn fetch_delta(state: &AppState, peer: &str, epoch: transport::Epoch, spec: tran
             Ok(changed) => {
                 state.notify_store_changed();
                 if changed {
-                    crate::swarm::relay::announce(&state, &peer).await;
+                    crate::swarm::relay::announce(&state, &peer, None).await;
                 }
             }
             Err(error) => {
@@ -598,12 +602,17 @@ pub(crate) async fn run_sync_worker(state: AppState, mut rx: mpsc::UnboundedRece
                 peers_guard.keys().cloned().collect()
             };
 
-            // Our own changes go to our neighbours; they pass them on. Others only get
-            // answers to what they send us.
+            // Our own changes go to our neighbours, saying they're ours: they pass them
+            // on along our tree (`swarm::tree`). Others only get answers to what they
+            // send us.
+            let own = crate::swarm::topology::pair_key(state.local_id());
             for peer in peer_keys {
                 if peer_swarm(&state, &peer).await.is_some()
                     && crate::swarm::topology::is_neighbour(&state, &peer)
                 {
+                    if let Some(peer_state) = state.get_peer_state(&peer).await {
+                        peer_state.lock().await.tree_writer = Some(own);
+                    }
                     dirty_peers.insert(peer);
                 }
             }
diff --git a/src/swarm/topology/mod.rs b/src/swarm/topology/mod.rs
index 1c1a6c4..b876fd0 100644
--- a/src/swarm/topology/mod.rs
+++ b/src/swarm/topology/mod.rs
@@ -26,6 +26,7 @@
 //! transport. With one transport, it is the one graph as before.
 
 use std::collections::{BTreeMap, BTreeSet};
+use std::sync::Arc;
 
 use tracing::info;
 
@@ -58,6 +59,42 @@ pub(crate) struct NeighbourView {
     neighbours: BTreeSet<String>,
     /// Every reachable member is a neighbour.
     full_mesh: bool,
+    /// The graph the member list gives, as if every member reached every other
+    /// (`member_list_graph`): computed when first asked for, kept while the member list
+    /// stays the same.
+    whole: std::sync::OnceLock<Arc<BTreeMap<String, BTreeSet<String>>>>,
+}
+
+/// Every member's neighbours as the member list gives them, as if every member reached
+/// every other: the same on every node that knows the same member list, whatever each
+/// reaches (`swarm::tree` builds writers' trees on it). On each transport among the
+/// members on it, all together. None until the view is known.
+pub(crate) fn member_list_graph(
+    state: &AppState,
+) -> Option<Arc<BTreeMap<String, BTreeSet<String>>>> {
+    let view = state.neighbours.read_or_recover();
+    let view = view.as_ref()?;
+    Some(Arc::clone(view.whole.get_or_init(|| {
+        let mut whole: BTreeMap<String, BTreeSet<String>> = view
+            .members
+            .iter()
+            .map(|member| (member.clone(), BTreeSet::new()))
+            .collect();
+        let transports: BTreeSet<&String> = view.transports.values().flatten().collect();
+        for transport in transports {
+            let on_it: Vec<String> = view
+                .transports
+                .iter()
+                .filter(|(_, on)| on.contains(transport))
+                .map(|(member, _)| member.clone())
+                .collect();
+            let graph = graph(&on_it, view.k, &|_| false, &|m| view.stable.contains(m));
+            for (member, neighbours) in graph {
+                whole.entry(member).or_default().extend(neighbours);
+            }
+        }
+        Arc::new(whole)
+    })))
 }
 
 /// Whether the neighbours are computed from the member list yet (before, everyone
@@ -342,7 +379,23 @@ pub(crate) async fn refresh(state: &AppState) -> Option<NeighbourChange> {
         }) {
             return None;
         }
-        view.as_ref().map(|view| view.neighbours.clone())
+        view.as_ref().map(|view| {
+            // The member list's graph stays while the member list does.
+            let whole = if view.members == current
+                && view.k == k
+                && view.stable == stable
+                && view.transports == transports
+            {
+                view.whole.clone()
+            } else {
+                Default::default()
+            };
+            (view.neighbours.clone(), whole)
+        })
+    };
+    let (previous, whole) = match previous {
+        Some((previous, whole)) => (Some(previous), whole),
+        None => (None, Default::default()),
     };
 
     // Before the first view everyone known counted as a neighbour.
@@ -400,6 +453,7 @@ pub(crate) async fn refresh(state: &AppState) -> Option<NeighbourChange> {
         unlisted,
         outbound,
         neighbours: computed,
+        whole,
     });
     // New neighbours are kept alive and judged.
     state.peers_due.now();
diff --git a/src/swarm/tree.rs b/src/swarm/tree.rs
new file mode 100644
index 0000000..489bc22
--- /dev/null
+++ b/src/swarm/tree.rs
@@ -0,0 +1,159 @@
+//! Passing a writer's changes on along its tree: every member gets them once, over a
+//! link that works, and as soon as the member before it has them.
+//!
+//! The tree is the same on every node, computed from what they share: the graph the
+//! member list gives (`topology::member_list_graph`: every member's neighbours, as if
+//! all reached all) and what each member tells its neighbours it reaches
+//! (`swarm::reach`). Every member's distance from the writer over that graph orders the
+//! members, by distance and then pair key; a member's parent is the neighbour it reaches
+//! that comes before it, the closest to the writer first, then by pair score
+//! (`topology::pair_score`). Coming before rules out cycles, and taking only a neighbour
+//! it reaches routes around links that don't work: one behind NAT that reaches a single
+//! member takes that one, at whatever distance.
+//!
+//! A sync message that passes a writer's change on says whose (`SyncEnvelope::writer`).
+//! Whoever applies one passes it on to its own children in that writer's tree
+//! (`forward`), however it got it; so does the writer, to all its neighbours, with its
+//! own writes. A member no member it reaches comes before (or whose parent didn't get
+//! the change) gets it from the announcements (`swarm::relay`), as before.
+//!
+//! The graph is computed once per member list (tens of milliseconds for thousands of
+//! members), the distances once per writer and graph (`MAX_WRITERS` kept).
+
+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, topology};
+
+/// The writers whose distances are kept at most.
+const MAX_WRITERS: usize = 256;
+
+type Graph = BTreeMap<String, BTreeSet<String>>;
+
+/// The distances from each writer, over the graph they were computed for.
+#[derive(Default)]
+pub(crate) struct Trees {
+    graph: Option<Arc<Graph>>,
+    /// The members by pair key, of that graph.
+    by_key: HashMap<u64, String>,
+    distances: HashMap<u64, Arc<HashMap<String, u32>>>,
+}
+
+/// The graph, and every member's distance from `writer` (by pair key) over it; None
+/// while the member list isn't known, or `writer` isn't in it.
+fn distances(state: &AppState, writer: u64) -> Option<(Arc<Graph>, Arc<HashMap<String, u32>>)> {
+    let graph = topology::member_list_graph(state)?;
+    let mut trees = state.trees.lock_or_recover();
+    if !trees
+        .graph
+        .as_ref()
+        .is_some_and(|known| Arc::ptr_eq(known, &graph))
+    {
+        trees.by_key = graph
+            .keys()
+            .map(|member| (topology::pair_key(member), member.clone()))
+            .collect();
+        trees.distances.clear();
+        trees.graph = Some(Arc::clone(&graph));
+    }
+    if let Some(known) = trees.distances.get(&writer) {
+        return Some((graph, Arc::clone(known)));
+    }
+    let from = trees.by_key.get(&writer)?.clone();
+    let mut distance = HashMap::from([(from.clone(), 0u32)]);
+    let mut queue = VecDeque::from([from]);
+    while let Some(member) = queue.pop_front() {
+        let next = distance[&member] + 1;
+        for neighbour in graph.get(&member).into_iter().flatten() {
+            if !distance.contains_key(neighbour) {
+                distance.insert(neighbour.clone(), next);
+                queue.push_back(neighbour.clone());
+            }
+        }
+    }
+    if trees.distances.len() >= MAX_WRITERS {
+        trees.distances.clear();
+    }
+    let distance = Arc::new(distance);
+    trees.distances.insert(writer, Arc::clone(&distance));
+    Some((graph, distance))
+}
+
+/// Our children in `writer`'s tree: our neighbours in the graph whose parent we are,
+/// and whom we reach.
+pub(crate) async fn children(state: &AppState, writer: u64) -> Vec<String> {
+    let Some((graph, distance)) = distances(state, writer) else {
+        return Vec::new();
+    };
+    let rank = |member: &str| {
+        distance
+            .get(member)
+            .map(|d| (*d, topology::pair_key(member)))
+    };
+    let me = state.local_id();
+    let Some(my_rank) = rank(me) else {
+        return Vec::new();
+    };
+    let mut children = Vec::new();
+    for member in graph.get(me).into_iter().flatten() {
+        let Some(its_rank) = rank(member) else {
+            continue;
+        };
+        if its_rank <= my_rank {
+            continue;
+        }
+        // Whom it reaches, as it told us; not knowing, as the graph says.
+        let reaches = match state.get_peer_state(member).await {
+            Some(peer_state) => reach::of(state, peer_state.lock().await.reaches.as_ref()),
+            None => None,
+        };
+        let parent = graph
+            .get(member)
+            .into_iter()
+            .flatten()
+            .filter(|candidate| rank(candidate).is_some_and(|r| r < its_rank))
+            .filter(|candidate| {
+                reaches
+                    .as_ref()
+                    .is_none_or(|keys| keys.contains(&topology::pair_key(candidate)))
+            })
+            .min_by_key(|candidate| {
+                (
+                    distance[*candidate],
+                    topology::pair_score(candidate, member),
+                )
+            });
+        if parent.is_some_and(|parent| parent == me) && topology::reaches(state, member).await {
+            children.push(member.clone());
+        }
+    }
+    children
+}
+
+/// We applied a change of `writer`'s that `from` passed on: to our children in its
+/// tree (but `from`), a sync message each, saying whose. Returns them.
+pub(crate) async fn forward(state: &AppState, writer: u64, from: &str) -> Vec<String> {
+    let children: Vec<String> = children(state, writer)
+        .await
+        .into_iter()
+        .filter(|child| child != from)
+        .collect();
+    for child in &children {
+        if let Some(peer_state) = state.get_peer_state(child).await {
+            peer_state.lock().await.tree_writer = Some(writer);
+        }
+        state.request_sync(child);
+    }
+    if !children.is_empty() {
+        trace!(
+            writer,
+            children = children.len(),
+            "Passing a change on along the writer's tree"
+        );
+    }
+    children
+}