add session IDs

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

1 files changed, 23 insertions, 15 deletions

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(());
             };