SyncSessions owns Automerge's sync state

1 files changed, 11 insertions, 78 deletions

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;