2 files changed, 35 insertions, 23 deletions

diff --git a/src/swarm/peer/inbound.rs b/src/swarm/peer/inbound.rs
index d4df6c7..b4a96e4 100644
--- a/src/swarm/peer/inbound.rs
+++ b/src/swarm/peer/inbound.rs
@@ -4,14 +4,16 @@ use std::time::Duration;
 
 use anyhow::{Result, anyhow, bail};
 use tokio::time::Instant;
-use tracing::{debug, info, warn};
+use tracing::{debug, warn};
 
 use super::admit::admit;
 use super::epoch::observe_peer_epoch;
 use super::handshake::{
     discover_peers, ensure_handshake, establish_peer, peer_swarm, remember_peer, send_hello_reply,
 };
-use super::keepalive::{heal_if_diverged, local_status, note_alive, note_status, status_received};
+use super::keepalive::{
+    heal_if_diverged, heard_from, local_status, note_alive, note_status, status_received,
+};
 use super::lanes::SyncLanes;
 use crate::state::AppState;
 use crate::swarm::sync::SyncEvent;
@@ -58,21 +60,9 @@ pub(crate) async fn handle_message(
 
     // Any inbound message proves liveness. When we heard from it before tells a
     // restarted peer how long it was away (in our HELLO_REPLY).
-    let (reactivated, previously_seen) = {
-        let peer_state_mutex = state.get_or_create_peer_state(&peer).await;
-        let mut peer_state = peer_state_mutex.lock().await;
-        let previously_seen = peer_state.last_seen;
-        if peer_state.unreachable_since.is_some() {
-            // Reachable again (`mark_seen` clears it).
-            state.topology_changed.notify_one();
-        }
-        (peer_state.mark_seen(Instant::now()), previously_seen)
-    };
-    crate::swarm::source::heard_from(state, &peer);
-
-    if reactivated {
-        info!(peer = %peer, "Dormant peer is active again; re-handshaking");
-    }
+    let peer_state = state.get_or_create_peer_state(&peer).await;
+    let (reactivated, previously_seen) =
+        heard_from(state, &peer, &peer_state, Instant::now(), |_| {}).await;
 
     let run = match &message {
         transport::ParsedMessage::Hello(hello) => Some(hello.run_id),
diff --git a/src/swarm/peer/keepalive.rs b/src/swarm/peer/keepalive.rs
index 9037fdd..0e80980 100644
--- a/src/swarm/peer/keepalive.rs
+++ b/src/swarm/peer/keepalive.rs
@@ -16,7 +16,7 @@ use tokio::time::Instant;
 use tracing::{debug, info, warn};
 
 use super::handshake::ensure_handshake;
-use crate::state::AppState;
+use crate::state::{AppState, PeerState};
 use crate::swarm::sync::SyncEvent;
 use crate::swarm::topology;
 use crate::transport::{self, ParsedMessage};
@@ -64,23 +64,45 @@ pub(super) async fn note_alive(state: &AppState, peer: &str, message: &ParsedMes
     let Some(peer_state) = state.get_peer_state(peer).await else {
         return;
     };
-    let reactivated = {
-        let mut peer_state = peer_state.lock().await;
-        if matches!(message, ParsedMessage::Ping(_)) {
+    let ping = matches!(message, ParsedMessage::Ping(_));
+    let (reactivated, _) = heard_from(state, peer, &peer_state, now, |peer_state| {
+        if ping {
             // Its keepalive (or a path probe): the next is ours.
             peer_state.keepalive_received_at = Some(now);
         }
+    })
+    .await;
+    if reactivated {
+        ensure_handshake(state, peer).await;
+    }
+}
+
+/// A message from `peer` came in: it is seen (`PeerState::mark_seen`: active, and
+/// reachable, again), and the catch-up source hears of it. `also` runs on its state
+/// under the same lock. Returns whether it was dormant (the caller re-handshakes), and
+/// when it was last seen before.
+pub(super) async fn heard_from(
+    state: &AppState,
+    peer: &str,
+    peer_state: &tokio::sync::Mutex<PeerState>,
+    now: Instant,
+    also: impl FnOnce(&mut PeerState),
+) -> (bool, Option<Instant>) {
+    let (reactivated, previously_seen) = {
+        let mut peer_state = peer_state.lock().await;
+        also(&mut peer_state);
+        let previously_seen = peer_state.last_seen;
         if peer_state.unreachable_since.is_some() {
             // Reachable again (`mark_seen` clears it).
             state.topology_changed.notify_one();
         }
-        peer_state.mark_seen(now)
+        (peer_state.mark_seen(now), previously_seen)
     };
     crate::swarm::source::heard_from(state, peer);
     if reactivated {
         info!(peer = %peer, "Dormant peer is active again; re-handshaking");
-        ensure_handshake(state, peer).await;
     }
+    (reactivated, previously_seen)
 }
 
 /// How long shutting down waits for the last STATUSes (`farewell`).