1 files changed, 113 insertions, 2 deletions

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