1 files changed, 21 insertions, 10 deletions

diff --git a/src/swarm/peer.rs b/src/swarm/peer.rs
index 55d0022..a936fa2 100644
--- a/src/swarm/peer.rs
+++ b/src/swarm/peer.rs
@@ -315,8 +315,14 @@ async fn ensure_handshake(state: &AppState, peer: &str) {
     }
 }
 
-/// Answers a HELLO that carried `run_id`.
-async fn send_hello_reply(state: &AppState, peer: &str, run_id: transport::RunId) -> Result<()> {
+/// Answers a HELLO that carried `run_id`; `away` is how long before it we had last heard
+/// from the peer.
+async fn send_hello_reply(
+    state: &AppState,
+    peer: &str,
+    run_id: transport::RunId,
+    away: Option<Duration>,
+) -> Result<()> {
     let swarm_id = state
         .swarm_id()
         .await
@@ -355,6 +361,7 @@ async fn send_hello_reply(state: &AppState, peer: &str, run_id: transport::RunId
         auth,
         ready,
         run_id,
+        away_ms: away.map(|away| away.as_millis() as u64),
     };
 
     debug!(to = %peer, ready, "Sending HELLO_REPLY");
@@ -484,6 +491,7 @@ pub(crate) async fn run_peer_worker(state: AppState) {
         if last_forget_check.is_none_or(|t| t.elapsed() >= FORGET_CHECK_INTERVAL) {
             apply_forget_records(&state).await;
             crate::swarm::members::ensure_registered(&state).await;
+            crate::swarm::stability::tick(&state).await;
             if let Some(change) = topology::refresh(&state).await {
                 apply_neighbour_change(&state, change).await;
             }
@@ -724,13 +732,14 @@ pub(crate) async fn handle_message(
 
     remember_peer(state, &peer).await;
 
-    // Any inbound message proves liveness.
-    let reactivated = state
-        .get_or_create_peer_state(&peer)
-        .await
-        .lock()
-        .await
-        .mark_seen(Instant::now());
+    // 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;
+        (peer_state.mark_seen(Instant::now()), previously_seen)
+    };
     crate::swarm::source::heard_from(state, &peer);
 
     if reactivated {
@@ -757,7 +766,8 @@ pub(crate) async fn handle_message(
             }
             establish_peer(state, &peer, swarm_id, true, hello.ready).await?;
 
-            if let Err(error) = send_hello_reply(state, &peer, hello.run_id).await {
+            let away = previously_seen.map(|seen| seen.elapsed());
+            if let Err(error) = send_hello_reply(state, &peer, hello.run_id, away).await {
                 // The peer retries its HELLO.
                 warn!(peer = %peer, error = %format!("{error:#}"), "HELLO_REPLY failed");
             }
@@ -768,6 +778,7 @@ pub(crate) async fn handle_message(
         }
         transport::ParsedMessage::HelloReply(reply) => {
             state.lock_rtt().hello_reply_received(&peer, Instant::now());
+            crate::swarm::stability::note_away(state, reply.away_ms.map(Duration::from_millis));
             establish_peer(state, &peer, reply.swarm_id.clone(), false, reply.ready).await?;
 
             discover_peers(state, &peer, &reply.peers).await;