late messages don't flip peer's epoch back

1 files changed, 46 insertions, 7 deletions

diff --git a/src/swarm/peer.rs b/src/swarm/peer.rs
index daf1af4..8c1e7eb 100644
--- a/src/swarm/peer.rs
+++ b/src/swarm/peer.rs
@@ -110,11 +110,12 @@ pub(crate) async fn send_ping(state: &AppState, peer: &str) -> Result<()> {
 pub(crate) async fn local_keepalive(state: &AppState, peer: &str) -> transport::Keepalive {
     // Heads and epoch come from the same locked replica, so they always match. The heads
     // are cached (`Replica`), so the read lock is enough.
-    let (heads, epoch) = {
+    let (heads, epoch, epoch_seq) = {
         let replica = state.doc.read().await;
         (
             transport::HeadsSummary::of(replica.heads().to_vec()),
             replica.epoch,
+            replica.epoch_seq,
         )
     };
     let syncs_with_you = match state.get_peer_state(peer).await {
@@ -129,6 +130,7 @@ pub(crate) async fn local_keepalive(state: &AppState, peer: &str) -> transport::
     transport::Keepalive {
         heads,
         epoch,
+        epoch_seq,
         ready: state.is_ready().await,
         syncs_with_you,
     }
@@ -276,6 +278,7 @@ async fn transmit_hello(state: &AppState, peer: &str) -> Result<()> {
         .as_ref()
         .map(|key| key.hello_tag(swarm_id.as_deref().unwrap_or(""), state.local_id(), peer))
         .unwrap_or_default();
+    let (epoch, epoch_seq) = state.epoch_stamp().await;
     state.lock_rtt().hello_sent(peer, Instant::now());
     let result = state
         .transport
@@ -283,7 +286,8 @@ async fn transmit_hello(state: &AppState, peer: &str) -> Result<()> {
             peer,
             &ParsedMessage::Hello(transport::HelloRequest {
                 swarm_id,
-                epoch: state.local_epoch(),
+                epoch,
+                epoch_seq,
                 auth,
                 ready: state.is_ready().await,
                 run_id: state.run_id,
@@ -354,14 +358,17 @@ async fn send_hello_reply(
         .map(|key| key.hello_reply_tag(&swarm_id, state.local_id(), peer))
         .unwrap_or_default();
     let ready = state.is_ready().await;
+    let (epoch, epoch_seq) = state.epoch_stamp().await;
     let reply = transport::HelloReply {
         swarm_id,
         peers,
-        epoch: state.local_epoch(),
+        epoch,
+        epoch_seq,
         auth,
         ready,
         run_id,
         away_ms: away.map(|away| away.as_millis() as u64),
+        sender_run_id: state.run_id,
     };
 
     debug!(to = %peer, ready, "Sending HELLO_REPLY");
@@ -582,9 +589,20 @@ pub(crate) async fn peer_swarm(state: &AppState, peer: &str) -> Option<String> {
 
 // --- epochs ------------------------------------------------------------------
 
-/// Records the epoch a peer advertised and compares it with ours. Returns true when both
-/// are on the same epoch, the only case in which the documents may be synced.
-async fn observe_peer_epoch(state: &AppState, peer: &str, remote: Epoch) -> bool {
+/// Records the epoch a peer advertised and compares it with ours. Returns true when the
+/// message's epoch is ours, the only case in which its documents may be synced.
+///
+/// `seq` is the peer's `epoch_seq` with it. A message with a lower one than we have was
+/// sent before the peer's last epoch change and delivered late: its epoch isn't
+/// recorded, so it can't flip our view of the peer back. The count starts over when the
+/// peer restarts, which its HELLO or HELLO_REPLY tells by a new run (`run`).
+async fn observe_peer_epoch(
+    state: &AppState,
+    peer: &str,
+    remote: Epoch,
+    seq: u32,
+    run: Option<transport::RunId>,
+) -> bool {
     let local = state.local_epoch();
     let relation = local.relation_to(&remote);
 
@@ -594,6 +612,21 @@ async fn observe_peer_epoch(state: &AppState, peer: &str, remote: Epoch) -> bool
 
     let previous = {
         let mut peer_state = peer_state_mutex.lock().await;
+        if let Some(run) = run.filter(|run| peer_state.run_id != Some(*run)) {
+            peer_state.run_id = Some(run);
+            peer_state.epoch_seq = None;
+        }
+        if peer_state.epoch_seq.is_some_and(|newest| seq < newest) {
+            debug!(
+                peer = %peer,
+                epoch = %remote,
+                seq,
+                newest = ?peer_state.epoch_seq,
+                "Ignoring a late message's epoch"
+            );
+            return relation == EpochRelation::Same;
+        }
+        peer_state.epoch_seq = Some(seq);
         let previous = peer_state.epoch.replace(remote);
         if previous.is_some_and(|previous| previous != remote) {
             // The peer switched documents; our sync state describes its old one. Count
@@ -758,7 +791,13 @@ pub(crate) async fn handle_message(
         info!(peer = %peer, "Dormant peer is active again; re-handshaking");
     }
 
-    let same_epoch = observe_peer_epoch(state, &peer, message.epoch()).await;
+    let run = match &message {
+        transport::ParsedMessage::Hello(hello) => Some(hello.run_id),
+        transport::ParsedMessage::HelloReply(reply) => Some(reply.sender_run_id),
+        _ => None,
+    };
+    let same_epoch =
+        observe_peer_epoch(state, &peer, message.epoch(), message.epoch_seq(), run).await;
 
     //debug!(from = %peer, message = ?message, "Received message");