add session IDs

This prevents the regenerate flag from being lost across reset_sync
during an in-flight send.

2 files changed, 87 insertions, 19 deletions

diff --git a/src/state/session.rs b/src/state/session.rs
index 0dcc76d..618d65b 100644
--- a/src/state/session.rs
+++ b/src/state/session.rs
@@ -6,6 +6,7 @@
 //! 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::sync::atomic::{AtomicU64, Ordering};
 use std::time::Duration;
 
 use automerge::sync::{self, SyncDoc};
@@ -26,6 +27,9 @@ pub(crate) struct Sending {
     pub(crate) regenerate: bool,
 }
 
+/// Session IDs, unique within the process.
+static NEXT_SESSION: AtomicU64 = AtomicU64::new(1);
+
 /// First wait before generating again after our Mycelium refused a message; doubles
 /// while it keeps refusing.
 const SEND_RETRY: Duration = Duration::from_millis(500);
@@ -35,6 +39,8 @@ 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 {
+    /// Tells this session from the one a reset puts in its place (`NEXT_SESSION`).
+    id: u64,
     state: sync::State,
 
     /// A message being handed to our Mycelium.
@@ -69,6 +75,7 @@ impl SyncSession {
     /// `reply_backoff`: waits for sync answers (`Timing::sync_reply_backoff`).
     pub(crate) fn new(reply_backoff: Backoff) -> Self {
         Self {
+            id: NEXT_SESSION.fetch_add(1, Ordering::Relaxed),
             state: sync::State::new(),
             sending: None,
             retry_at: None,
@@ -193,8 +200,20 @@ impl SyncSession {
         self.state.their_heads.is_none() && we_ready && peer_ready == Some(false)
     }
 
-    /// Our Mycelium took the message.
-    pub(crate) fn send_succeeded(&mut self) -> bool {
+    /// Which session this is: a message generated from it reports its send's outcome
+    /// with this (`send_succeeded`, `send_failed`).
+    pub(crate) fn id(&self) -> u64 {
+        self.id
+    }
+
+    /// Our Mycelium took the message generated from session `session`. Returns whether
+    /// to generate again: a pass was requested while it was handed over, or a reset
+    /// replaced the session meanwhile (any pass requested then was lost with the old
+    /// session, and the new one hasn't necessarily opened yet).
+    pub(crate) fn send_succeeded(&mut self, session: u64) -> bool {
+        if session != self.id {
+            return true;
+        }
         self.retry_at = None;
         self.send_backoff.reset();
         self.sending
@@ -205,8 +224,11 @@ impl SyncSession {
     /// Our Mycelium refused the message: as far as the peer is concerned, it was never
     /// sent. Takes back its changes (they go out with the next message), clears
     /// `in_flight`, and returns how long to wait before generating again. None if a reset
-    /// already discarded what was being sent.
-    pub(crate) fn send_failed(&mut self, now: Instant) -> Option<Duration> {
+    /// already discarded what was being sent (the message came from another `session`).
+    pub(crate) fn send_failed(&mut self, session: u64, now: Instant) -> Option<Duration> {
+        if session != self.id {
+            return None;
+        }
         let sending = self.sending.take()?;
         for hash in &sending.sent {
             self.state.sent_hashes.remove(hash);
@@ -264,3 +286,41 @@ impl SyncSession {
         Some(wait.sent.len())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn session() -> SyncSession {
+        SyncSession::new(Backoff::new(
+            Duration::from_secs(1),
+            Duration::from_secs(30),
+        ))
+    }
+
+    fn sending() -> Option<Sending> {
+        Some(Sending {
+            sent: Vec::new(),
+            regenerate: false,
+        })
+    }
+
+    #[test]
+    fn a_send_from_a_replaced_session_leaves_the_new_one_alone() {
+        let old = session();
+        let mut new = session();
+        new.sending = sending();
+
+        // A reset discarded the old session while its message was handed over: the
+        // new one's own message stays, and the peer is synced again.
+        assert!(new.send_succeeded(old.id()));
+        assert!(new.sending.is_some());
+        assert_eq!(new.send_failed(old.id(), Instant::now()), None);
+        assert!(new.sending.is_some());
+        assert_eq!(new.retry_at, None);
+
+        // Its own outcome still applies.
+        assert!(!new.send_succeeded(new.id()));
+        assert!(new.sending.is_none());
+    }
+}
diff --git a/src/swarm/sync.rs b/src/swarm/sync.rs
index 805d700..5a4cba4 100644
--- a/src/swarm/sync.rs
+++ b/src/swarm/sync.rs
@@ -126,7 +126,7 @@ async fn deliver_sync_to_peer(state: &AppState, peer: &str) -> Result<()> {
         }
     }
 
-    let encoded = {
+    let (session, encoded) = {
         let doc = Arc::clone(&state.doc);
         let psm = Arc::clone(&peer_state_mutex);
         let transport = Arc::clone(&state.transport);
@@ -135,11 +135,12 @@ async fn deliver_sync_to_peer(state: &AppState, peer: &str) -> Result<()> {
         // Before this message is counted.
         let bytes_before = SyncBytes::of(&state.traffic, peer);
 
-        let generated =
-            tokio::task::spawn_blocking(move || -> Result<Option<transport::EncodedMessage>> {
+        // The session the message comes from, for its send's outcome.
+        let (session, generated) = tokio::task::spawn_blocking(
+            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 (epoch, epoch_seq, message) = {
+                let (session, epoch, epoch_seq, message) = {
                     let mut replica = doc.blocking_write();
                     let mut peer_state = psm.blocking_lock();
                     let (epoch, epoch_seq) = (replica.epoch, replica.epoch_seq);
@@ -149,10 +150,10 @@ async fn deliver_sync_to_peer(state: &AppState, peer: &str) -> Result<()> {
                         peer_state
                             .sync
                             .generate(&mut replica.doc, Instant::now(), bytes_before);
-                    (epoch, epoch_seq, message)
+                    (peer_state.sync.id(), epoch, epoch_seq, message)
                 };
                 let Some(message) = message else {
-                    return Ok(None);
+                    return (session, Ok(None));
                 };
                 let envelope = transport::SyncEnvelope {
                     swarm_id,
@@ -161,12 +162,14 @@ async fn deliver_sync_to_peer(state: &AppState, peer: &str) -> Result<()> {
                     message,
                     ready,
                 };
-                transport
+                let encoded = transport
                     .encode(&peer_owned, &ParsedMessage::Sync(envelope))
-                    .map(Some)
-            })
-            .await
-            .context("Sync generation task panicked")?;
+                    .map(Some);
+                (session, encoded)
+            },
+        )
+        .await
+        .context("Sync generation task panicked")?;
         let generated = match generated {
             Ok(generated) => generated,
             Err(error) => {
@@ -175,7 +178,7 @@ async fn deliver_sync_to_peer(state: &AppState, peer: &str) -> Result<()> {
                     .lock()
                     .await
                     .sync
-                    .send_failed(Instant::now());
+                    .send_failed(session, Instant::now());
                 return Err(error);
             }
         };
@@ -184,13 +187,13 @@ async fn deliver_sync_to_peer(state: &AppState, peer: &str) -> Result<()> {
             return Ok(());
         };
 
-        encoded
+        (session, encoded)
     };
 
     // No state locks are held across the send.
     match state.transport.send_encoded(&encoded).await {
         Ok(()) => {
-            let regenerate = peer_state_mutex.lock().await.sync.send_succeeded();
+            let regenerate = peer_state_mutex.lock().await.sync.send_succeeded(session);
 
             trace!(regenerate, "Outbound sync message accepted by Mycelium");
 
@@ -208,9 +211,14 @@ async fn deliver_sync_to_peer(state: &AppState, peer: &str) -> Result<()> {
                 .lock()
                 .await
                 .sync
-                .send_failed(Instant::now())
+                .send_failed(session, Instant::now())
             else {
+                // The new session may not have opened yet.
                 debug!("Sync message refused, but a reset already discarded it");
+                state
+                    .sync_events
+                    .send(SyncEvent::PeerNeedsSync(peer.to_owned()))
+                    .map_err(|error| anyhow!("Failed to notify sync worker: {error}"))?;
                 return Ok(());
             };