fix race between alive check and forgetting

3 files changed, 31 insertions, 9 deletions

diff --git a/src/api/peers.rs b/src/api/peers.rs
index e22f3cd..eb522b0 100644
--- a/src/api/peers.rs
+++ b/src/api/peers.rs
@@ -235,7 +235,14 @@ pub(crate) async fn delete_peer(
         .await
         .map_err(|e| store_status(e, "Failed to write forget record"))?;
 
-    if state.peer_states.write().await.remove(&peer).is_some() {
+    // Unless forced, a peer that came alive since the check above stays: it cancels
+    // the record.
+    let forgotten = if force.allows("alive") {
+        state.peer_states.write().await.remove(&peer).is_some()
+    } else {
+        crate::swarm::peer::forget_unless_alive(&state, &peer, Instant::now()).await
+    };
+    if forgotten {
         warn!(peer = %peer, "Peer forgotten swarm-wide by operator");
     }
 
diff --git a/src/swarm/peer/forget.rs b/src/swarm/peer/forget.rs
index ae49a5a..10a676d 100644
--- a/src/swarm/peer/forget.rs
+++ b/src/swarm/peer/forget.rs
@@ -46,16 +46,10 @@ pub(super) async fn apply_forget_records(state: &AppState) {
             continue;
         }
 
-        let Some(peer_state_mutex) = state.get_peer_state(&pk).await else {
-            continue; // already forgotten here
-        };
-
-        let alive = peer_state_mutex.lock().await.is_alive(&state.timing, now);
-
         // An alive peer is kept. Only the peer itself cancels its record: it gets the
         // record from whoever it is alive to, and one cancel is one change instead of
         // one from every node that hears from it.
-        if !alive && state.peer_states.write().await.remove(&pk).is_some() {
+        if forget_unless_alive(state, &pk, now).await {
             info!(peer = %pk, "Forgetting peer per swarm-wide record");
         }
     }
@@ -80,6 +74,27 @@ pub(super) async fn apply_forget_records(state: &AppState) {
         .await;
 }
 
+/// Forgets `pk` here unless it is alive, and returns whether it did. Judged and removed
+/// under the same locks, so a message that arrives meanwhile can't be lost to a peer
+/// state that is already gone. A peer whose state is in use is left for the next check.
+pub(crate) async fn forget_unless_alive(state: &AppState, pk: &str, now: Instant) -> bool {
+    let mut peer_states = state.peer_states.write().await;
+    let Some(peer_state) = peer_states.get(pk) else {
+        return false;
+    };
+    // Not waiting: holding the map, we must not wait for a peer's lock (its holder may
+    // be waiting for the map).
+    let Ok(peer_state) = peer_state.try_lock() else {
+        return false;
+    };
+    if peer_state.is_alive(&state.timing, now) {
+        return false;
+    }
+    drop(peer_state);
+    peer_states.remove(pk);
+    true
+}
+
 /// How long a peer that never became a member is kept once nothing more is heard from
 /// it: one listed to us that never answered, or one that keeps failing the handshake
 /// (another swarm, a wrong key). Anyone can name peers to a node, so without a limit
diff --git a/src/swarm/peer/mod.rs b/src/swarm/peer/mod.rs
index 086f5d2..f0222c4 100644
--- a/src/swarm/peer/mod.rs
+++ b/src/swarm/peer/mod.rs
@@ -17,7 +17,7 @@ mod inbound;
 mod keepalive;
 mod worker;
 
-pub(crate) use forget::forgotten_path;
+pub(crate) use forget::{forget_unless_alive, forgotten_path};
 pub(crate) use handshake::{
     apply_neighbour_change, peer_swarm, rehandshake, remember_peer, send_hello_once,
 };