1 files changed, 22 insertions, 7 deletions

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