add farewell ping

2 files changed, 22 insertions, 2 deletions

diff --git a/src/swarm/peer/keepalive.rs b/src/swarm/peer/keepalive.rs
index a593ac9..b22d954 100644
--- a/src/swarm/peer/keepalive.rs
+++ b/src/swarm/peer/keepalive.rs
@@ -2,6 +2,7 @@
 //! documents differ while sync is idle.
 
 use std::sync::atomic::Ordering;
+use std::time::Duration;
 
 use anyhow::Result;
 use tokio::time::Instant;
@@ -23,6 +24,25 @@ pub(crate) async fn send_ping(state: &AppState, peer: &str) -> Result<()> {
         .await
 }
 
+/// How long shutting down waits for the last PINGs (`farewell`).
+const FAREWELL_TIMEOUT: Duration = Duration::from_secs(2);
+
+/// At shutdown: one last PING to each neighbour, saying we aren't ready. A neighbour
+/// catching up from us picks another source now, instead of once we go dormant. Bounded
+/// in time: our Mycelium may be gone already.
+pub(crate) async fn farewell(state: &AppState) {
+    state.stopping.store(true, Ordering::Release);
+    let mut pings = tokio::task::JoinSet::new();
+    for (peer, _) in state.peers_snapshot().await {
+        if topology::is_neighbour(state, &peer) {
+            let state = state.clone();
+            pings.spawn(async move { send_ping(&state, &peer).await });
+        }
+    }
+    let all_sent = async { while pings.join_next().await.is_some() {} };
+    let _ = tokio::time::timeout(FAREWELL_TIMEOUT, all_sent).await;
+}
+
 /// Our keepalive for `peer`.
 pub(crate) async fn local_keepalive(state: &AppState, peer: &str) -> transport::Keepalive {
     // Heads and epoch are published together (`Replica::stamps`), so they always match,
@@ -48,7 +68,7 @@ pub(crate) async fn local_keepalive(state: &AppState, peer: &str) -> transport::
         heads,
         epoch,
         epoch_seq,
-        ready: state.is_ready().await,
+        ready: state.is_ready().await && !state.stopping.load(Ordering::Acquire),
         syncs_with_you,
     }
 }
diff --git a/src/swarm/peer/mod.rs b/src/swarm/peer/mod.rs
index 2e5c1dc..0249f3b 100644
--- a/src/swarm/peer/mod.rs
+++ b/src/swarm/peer/mod.rs
@@ -24,6 +24,6 @@ pub(crate) use handshake::{
     apply_neighbour_change, peer_swarm, rehandshake, remember_peer, send_hello_once,
 };
 pub(crate) use inbound::handle_message;
-pub(crate) use keepalive::send_ping;
+pub(crate) use keepalive::{farewell, send_ping};
 pub(crate) use lanes::SyncLanes;
 pub(crate) use worker::run_peer_worker;