SyncSessions owns Automerge's sync state

5 files changed, 128 insertions, 85 deletions

diff --git a/src/simtest/mod.rs b/src/simtest/mod.rs
index 26092ee..e7fc719 100644
--- a/src/simtest/mod.rs
+++ b/src/simtest/mod.rs
@@ -219,11 +219,11 @@ impl Sim {
                     ago(p.last_seen),
                     ago(p.last_sync_rx),
                     p.is_unreachable(&node.timing, now),
-                    p.sync.state.their_heads.as_ref().map(|heads| heads
+                    p.sync.automerge_view().0.map(|heads| heads
                         .iter()
                         .map(|head| head.to_string()[..8].to_owned())
                         .collect::<Vec<_>>()),
-                    p.sync.state.in_flight,
+                    p.sync.automerge_view().1,
                     p.sync.sending.is_some(),
                     p.sync
                         .reply
diff --git a/src/state/mod.rs b/src/state/mod.rs
index 1c31c61..55a260d 100644
--- a/src/state/mod.rs
+++ b/src/state/mod.rs
@@ -15,7 +15,7 @@ mod timing;
 
 pub(crate) use peer::{HelloState, Incidents, Liveness, PeerState, Standing};
 pub(crate) use read::{ReadSource, StaleView};
-pub(crate) use session::{Sending, SyncSession};
+pub(crate) use session::SyncSession;
 pub(crate) use timing::Timing;
 
 use std::collections::HashMap;
diff --git a/src/state/session.rs b/src/state/session.rs
index 94f9d71..0dcc76d 100644
--- a/src/state/session.rs
+++ b/src/state/session.rs
@@ -1,9 +1,17 @@
 //! One Automerge sync with a peer, from a reset to the next.
+//!
+//! Automerge's sync state (`automerge::sync::State`) is private to this module: the
+//! rest of the node generates, receives and inspects through `SyncSession`'s methods.
+//! The workarounds for its assumption of reliable delivery (resending after a timeout,
+//! forgetting what was sent when the peer still needs it) live here with it, so an
+//! Automerge upgrade that changes the state touches this file only.
 
 use std::time::Duration;
 
-use automerge::sync;
+use automerge::sync::{self, SyncDoc};
+use automerge::{AutoCommit, AutomergeError};
 use tokio::time::Instant;
+use tracing::debug;
 
 use crate::backoff::Backoff;
 use crate::swarm::throughput::SyncBytes;
@@ -27,7 +35,7 @@ const SEND_RETRY_MAX: Duration = Duration::from_secs(30);
 /// forget lives here, so a reset replaces it whole and can't miss a field.
 #[derive(Clone)]
 pub(crate) struct SyncSession {
-    pub(crate) state: sync::State,
+    state: sync::State,
 
     /// A message being handed to our Mycelium.
     pub(crate) sending: Option<Sending>,
@@ -70,6 +78,109 @@ impl SyncSession {
         }
     }
 
+    /// Whether the peer said anything in this session yet (its heads are known).
+    pub(crate) fn has_opened(&self) -> bool {
+        self.state.their_heads.is_some()
+    }
+
+    /// The next sync message for the peer from `doc`, if there is anything to say. A
+    /// message is marked as being handed over (`sending`), and an answer to it is
+    /// expected, unless the peer's heads are ours already: it has everything, and
+    /// Automerge answers only when it has something to say, so waiting for an answer
+    /// would only end in a pointless resend. `bytes`: the sync bytes exchanged with the
+    /// peer so far (`swarm::throughput`).
+    pub(crate) fn generate(
+        &mut self,
+        doc: &mut AutoCommit,
+        now: Instant,
+        bytes: SyncBytes,
+    ) -> Option<sync::Message> {
+        let sent_before = self.state.sent_hashes.clone();
+        let message = doc.sync().generate_sync_message(&mut self.state);
+        let in_step = self.state.their_heads.as_ref() == Some(&doc.get_heads());
+        let sent: Vec<_> = self
+            .state
+            .sent_hashes
+            .difference(&sent_before)
+            .copied()
+            .collect();
+        if in_step {
+            self.reply_settled();
+        } else if message.is_some() {
+            self.expect_reply(sent.clone(), now, bytes);
+        }
+        if message.is_some() {
+            // Marked together with the generation, so no second message is generated
+            // while this one is handed over.
+            self.sending = Some(Sending {
+                sent,
+                regenerate: false,
+            });
+        }
+        message
+    }
+
+    /// Applies a sync message from the peer to `doc`. Doesn't settle the awaited
+    /// answer: the caller measures the exchange first (`reply`).
+    pub(crate) fn receive(
+        &mut self,
+        doc: &mut AutoCommit,
+        message: sync::Message,
+    ) -> Result<(), AutomergeError> {
+        doc.sync().receive_sync_message(&mut self.state, message)
+    }
+
+    /// Whether `doc` holds everything the peer had when it last told us its heads: they
+    /// are known, not empty (empty heads are "contained" in every document, so they
+    /// prove nothing), and all in `doc`.
+    pub(crate) fn have_all_theirs(&self, doc: &mut AutoCommit) -> bool {
+        self.state.their_heads.as_ref().is_some_and(|heads| {
+            !heads.is_empty()
+                && heads
+                    .iter()
+                    .all(|head| doc.get_change_by_hash(head).is_some())
+        })
+    }
+
+    /// The peer has heads we lack (its keepalive said so): ask for what they need,
+    /// even though Automerge thinks a message is still in flight.
+    pub(crate) fn pull_from(&mut self, heads: Vec<automerge::ChangeHash>) {
+        self.state.their_heads = Some(heads);
+        self.state.in_flight = false;
+    }
+
+    /// Automerge never resends a change it believes the peer has received
+    /// (`sent_hashes`). If the message carrying it was lost while later ones arrive, the
+    /// peer keeps asking for it (`need`, from the heads we advertise) and we keep not
+    /// sending it: a livelock the reply timeout can't see (answers keep coming) and the
+    /// keepalive heal can't either (sync never goes idle). So when the peer needs a
+    /// change we think it has, what we think it has is wrong: forget all of it, and the
+    /// next message carries everything its bloom filter lacks. (A reordered message can
+    /// make that send a change twice; cheap.) Call after `receive`.
+    pub(crate) fn resend_if_needed(&mut self, peer: &str) {
+        let Some(need) = &self.state.their_need else {
+            return;
+        };
+        if need
+            .iter()
+            .any(|hash| self.state.sent_hashes.contains(hash))
+        {
+            debug!(
+                peer = %peer,
+                sent = self.state.sent_hashes.len(),
+                "Peer needs changes we had sent; forgetting what was sent"
+            );
+            self.state.sent_hashes.clear();
+        }
+    }
+
+    /// For the simulator's dump: the peer's heads as Automerge knows them, and whether
+    /// it thinks a message is in flight.
+    #[cfg(test)]
+    pub(crate) fn automerge_view(&self) -> (Option<&[automerge::ChangeHash]>, bool) {
+        (self.state.their_heads.as_deref(), self.state.in_flight)
+    }
+
     /// Whether we leave the first message of this session to the peer.
     ///
     /// After a reset, both sides usually open at once, and the openings cross. Between
diff --git a/src/swarm/relay.rs b/src/swarm/relay.rs
index 5053be7..4bfc2b6 100644
--- a/src/swarm/relay.rs
+++ b/src/swarm/relay.rs
@@ -147,8 +147,7 @@ pub(crate) async fn pull_if_behind(state: &AppState, peer: &str, remote: &transp
         match missing {
             Some(heads) => {
                 // Automerge asks for whatever these heads need that we don't have.
-                peer_state.sync.state.their_heads = Some(heads);
-                peer_state.sync.state.in_flight = false;
+                peer_state.sync.pull_from(heads);
             }
             // A fresh session's opening carries our heads and bloom filter; the answer
             // brings what we lack.
diff --git a/src/swarm/sync.rs b/src/swarm/sync.rs
index fc6dc43..07f66a3 100644
--- a/src/swarm/sync.rs
+++ b/src/swarm/sync.rs
@@ -6,13 +6,12 @@ use std::sync::Arc;
 use std::sync::atomic::Ordering;
 
 use anyhow::{Context, Result, anyhow, bail};
-use automerge::sync::SyncDoc;
 use tokio::sync::mpsc;
 use tokio::time::Instant;
 use tracing::{debug, info, trace, warn};
 
 use crate::locks::MutexExt;
-use crate::state::{AppState, Liveness, Sending, Standing};
+use crate::state::{AppState, Liveness, Standing};
 use crate::swarm::peer::{peer_swarm, send_hello_once};
 use crate::swarm::source;
 use crate::swarm::throughput::SyncBytes;
@@ -93,9 +92,7 @@ async fn deliver_sync_to_peer(state: &AppState, peer: &str) -> Result<()> {
             return Ok(());
         }
 
-        if peer_state.sync.state.their_heads.is_none()
-            && !crate::swarm::topology::is_neighbour(state, peer)
-        {
+        if !peer_state.sync.has_opened() && !crate::swarm::topology::is_neighbour(state, peer) {
             // We only answer peers that aren't our neighbours.
             trace!("Not opening sync with a peer that isn't our neighbour");
             return Ok(());
@@ -146,39 +143,12 @@ async fn deliver_sync_to_peer(state: &AppState, peer: &str) -> Result<()> {
                     let mut replica = doc.blocking_write();
                     let mut peer_state = psm.blocking_lock();
                     let (epoch, epoch_seq) = (replica.epoch, replica.epoch_seq);
-                    let sent_before = peer_state.sync.state.sent_hashes.clone();
-                    let message = {
-                        let sync_doc = replica.doc.sync();
-                        sync_doc.generate_sync_message(&mut peer_state.sync.state)
-                    };
-                    // A peer whose heads are ours already has everything; Automerge
-                    // answers only when it has something to say, so it will stay
-                    // silent. Waiting for an answer would only end in a pointless
-                    // resend.
-                    let in_step = peer_state.sync.state.their_heads.as_ref()
-                        == Some(&replica.doc.get_heads());
-                    let sent: Vec<_> = peer_state
-                        .sync
-                        .state
-                        .sent_hashes
-                        .difference(&sent_before)
-                        .copied()
-                        .collect();
-                    if in_step {
-                        peer_state.sync.reply_settled();
-                    } else if message.is_some() {
+                    // Under the same locks as the epoch label: the message belongs to
+                    // this document.
+                    let message =
                         peer_state
                             .sync
-                            .expect_reply(sent.clone(), Instant::now(), bytes_before);
-                    }
-                    if message.is_some() {
-                        // Marked under the same lock as the generation, so no second
-                        // message is generated while this one is handed over.
-                        peer_state.sync.sending = Some(Sending {
-                            sent,
-                            regenerate: false,
-                        });
-                    }
+                            .generate(&mut replica.doc, Instant::now(), bytes_before);
                     (epoch, epoch_seq, message)
                 };
                 let Some(message) = message else {
@@ -348,10 +318,9 @@ pub(crate) async fn handle_sync(
             peer_state.reports_ready = Some(from_ready);
             let heads_before = replica.heads().to_vec();
 
-            replica
-                .doc
-                .sync()
-                .receive_sync_message(&mut peer_state.sync.state, message)
+            peer_state
+                .sync
+                .receive(&mut replica.doc, message)
                 .context("Automerge synchronization failed")?;
             peer_state.last_sync_rx = Some(Instant::now());
             if let Some(wait) = &peer_state.sync.reply {
@@ -361,7 +330,7 @@ pub(crate) async fn handle_sync(
             }
             peer_state.sync.reply_settled();
 
-            resend_if_needed(&peer, &mut peer_state.sync.state);
+            peer_state.sync.resend_if_needed(&peer);
             replica.refresh_heads();
             let heads_after = replica.heads().to_vec();
             stats.heads_seen(heads_after.len());
@@ -386,19 +355,7 @@ pub(crate) async fn handle_sync(
                 // catching up has them, and so does a ready peer whose document is empty
                 // (a misconfigured second swarm, or a genesis node in its first seconds,
                 // before it registered itself).
-                check_caught_up
-                    && from_ready
-                    && peer_state
-                        .sync
-                        .state
-                        .their_heads
-                        .as_ref()
-                        .is_some_and(|heads| {
-                            !heads.is_empty()
-                                && heads
-                                    .iter()
-                                    .all(|h| replica.doc.get_change_by_hash(h).is_some())
-                        }),
+                check_caught_up && from_ready && peer_state.sync.have_all_theirs(&mut replica.doc),
                 changed,
             )))
         })
@@ -436,30 +393,6 @@ pub(crate) async fn handle_sync(
     Ok(())
 }
 
-/// Automerge never resends a change it believes the peer has received (`sent_hashes`).
-/// If the message carrying it was lost while later ones arrive, the peer keeps asking for
-/// it (`need`, from the heads we advertise) and we keep not sending it: a livelock the
-/// reply timeout can't see (answers keep coming) and the keepalive heal can't either
-/// (sync never goes idle). So when the peer needs a change we think it has, what we
-/// think it has is wrong: forget all of it, and the next message carries everything its
-/// bloom filter lacks. (A reordered message can make that send a change twice; cheap.)
-fn resend_if_needed(peer: &str, sync_state: &mut automerge::sync::State) {
-    let Some(need) = &sync_state.their_need else {
-        return;
-    };
-    if need
-        .iter()
-        .any(|hash| sync_state.sent_hashes.contains(hash))
-    {
-        debug!(
-            peer = %peer,
-            sent = sync_state.sent_hashes.len(),
-            "Peer needs changes we had sent; forgetting what was sent"
-        );
-        sync_state.sent_hashes.clear();
-    }
-}
-
 pub(crate) async fn run_sync_worker(state: AppState, mut rx: mpsc::UnboundedReceiver<SyncEvent>) {
     let mut dirty_peers = HashSet::<String>::new();
     let mut document_dirty = false;