unify peers snapshot logic

6 files changed, 24 insertions, 52 deletions

diff --git a/src/api/peers.rs b/src/api/peers.rs
index fac3440..c7b3c50 100644
--- a/src/api/peers.rs
+++ b/src/api/peers.rs
@@ -1,14 +1,11 @@
 //! /peers handlers.
 
-use std::sync::Arc;
-
 use axum::{
     Json,
     extract::{Path, Query, State},
     http::StatusCode,
 };
 use serde::{Deserialize, Serialize};
-use tokio::sync::Mutex;
 use tokio::time::Instant;
 use tracing::warn;
 use utoipa::ToSchema;
@@ -17,7 +14,7 @@ use super::store_status;
 use crate::document::replica::Epoch;
 use crate::document::store;
 use crate::locks::MutexExt;
-use crate::state::{AppState, HelloState, Incidents, Liveness, PeerState};
+use crate::state::{AppState, HelloState, Incidents, Liveness};
 use crate::swarm::peer::forgotten_path;
 use crate::transport::traffic::PeerTraffic;
 
@@ -120,13 +117,7 @@ pub(crate) struct DeletePeerQuery {
 pub(crate) async fn list_peers(
     State(state): State<AppState>,
 ) -> Result<Json<PeerListResponse>, StatusCode> {
-    let snapshot: Vec<(String, Arc<Mutex<PeerState>>)> = state
-        .peer_states
-        .read()
-        .await
-        .iter()
-        .map(|(pk, m)| (pk.clone(), Arc::clone(m)))
-        .collect();
+    let snapshot = state.peers_snapshot().await;
 
     let now = Instant::now();
     let local_epoch = state.local_epoch();
diff --git a/src/api/stats/prometheus.rs b/src/api/stats/prometheus.rs
index 07a6cc8..3709ea4 100644
--- a/src/api/stats/prometheus.rs
+++ b/src/api/stats/prometheus.rs
@@ -120,13 +120,7 @@ struct Peer {
 
 async fn peers(state: &AppState) -> Vec<Peer> {
     let now = Instant::now();
-    let snapshot: Vec<_> = state
-        .peer_states
-        .read()
-        .await
-        .iter()
-        .map(|(key, peer_state)| (key.clone(), peer_state.clone()))
-        .collect();
+    let snapshot = state.peers_snapshot().await;
     let mut aliases = crate::swarm::members::aliases(&state.doc.read().await.doc);
     let mut peers = Vec::with_capacity(snapshot.len());
     for (key, peer_state) in snapshot {
diff --git a/src/simtest/mod.rs b/src/simtest/mod.rs
index 408f53e..26092ee 100644
--- a/src/simtest/mod.rs
+++ b/src/simtest/mod.rs
@@ -200,13 +200,7 @@ impl Sim {
                     .load(std::sync::atomic::Ordering::Acquire),
             );
             let _ = writeln!(out, "  view: {:?}", *node.neighbours.read_or_recover());
-            let peers: Vec<_> = node
-                .peer_states
-                .read()
-                .await
-                .iter()
-                .map(|(peer, state)| (peer.clone(), Arc::clone(state)))
-                .collect();
+            let peers = node.peers_snapshot().await;
             for (peer, peer_state) in peers {
                 let p = peer_state.lock().await;
                 let _ = writeln!(
diff --git a/src/state/host.rs b/src/state/host.rs
index 9b149b8..fd40bf3 100644
--- a/src/state/host.rs
+++ b/src/state/host.rs
@@ -7,11 +7,11 @@ use std::time::Duration;
 
 use anyhow::Result;
 use automerge::{AutoCommit, ChangeHash};
-use tokio::sync::{Mutex, broadcast};
+use tokio::sync::broadcast;
 use tokio::time::Instant;
 use tracing::debug;
 
-use super::{AppState, HelloState, Liveness, NodeLifecycle, PeerState, StaleView, Standing};
+use super::{AppState, HelloState, Liveness, NodeLifecycle, StaleView, Standing};
 use crate::document::replica::Epoch;
 use crate::locks::MutexExt;
 use crate::remake::{Host, PeerSummary, RemakeState};
@@ -22,13 +22,7 @@ impl AppState {
     /// document, so start over with all of them. Peers still on the old epoch are
     /// skipped by the sync worker until they switch too.
     async fn after_epoch_switch(&self) {
-        let peers: Vec<(String, Arc<Mutex<PeerState>>)> = self
-            .peer_states
-            .read()
-            .await
-            .iter()
-            .map(|(id, peer)| (id.clone(), Arc::clone(peer)))
-            .collect();
+        let peers = self.peers_snapshot().await;
 
         // Sync states were reset under the document lock by the swap itself.
         self.publish_change();
@@ -126,13 +120,7 @@ impl Host for AppState {
     fn peer_summaries(&self) -> impl Future<Output = Vec<PeerSummary>> + Send {
         let state = self.clone();
         async move {
-            let peers: Vec<(String, Arc<Mutex<PeerState>>)> = state
-                .peer_states
-                .read()
-                .await
-                .iter()
-                .map(|(id, peer)| (id.clone(), Arc::clone(peer)))
-                .collect();
+            let peers = state.peers_snapshot().await;
 
             let mut summaries = Vec::with_capacity(peers.len());
             for (id, peer) in peers {
diff --git a/src/state/mod.rs b/src/state/mod.rs
index 3336b00..1c31c61 100644
--- a/src/state/mod.rs
+++ b/src/state/mod.rs
@@ -280,6 +280,19 @@ impl AppState {
         Ok(())
     }
 
+    /// Every known peer with its state, copied out of the peer map. Lock a peer's state
+    /// only from such a copy, never while holding the map's lock: the map would then
+    /// wait for the peer, and with it everything that only needs the map (a lookup, a
+    /// new peer).
+    pub(crate) async fn peers_snapshot(&self) -> Vec<(String, Arc<Mutex<PeerState>>)> {
+        self.peer_states
+            .read()
+            .await
+            .iter()
+            .map(|(peer, state)| (peer.clone(), Arc::clone(state)))
+            .collect()
+    }
+
     /// Looks up a peer without creating it. Background work (sync, retries,
     /// keepalives) must use this so a forgotten peer isn't resurrected.
     pub(crate) async fn get_peer_state(&self, pk: &str) -> Option<Arc<Mutex<PeerState>>> {
diff --git a/src/swarm/peer.rs b/src/swarm/peer.rs
index ec6543b..79e958b 100644
--- a/src/swarm/peer.rs
+++ b/src/swarm/peer.rs
@@ -351,12 +351,10 @@ async fn send_hello_reply(
         .ok_or_else(|| anyhow!("Cannot send HELLO_REPLY without an established swarm ID"))?;
 
     let candidates: Vec<(String, Arc<Mutex<PeerState>>)> = state
-        .peer_states
-        .read()
+        .peers_snapshot()
         .await
-        .iter()
+        .into_iter()
         .filter(|(known, _)| known.as_str() != peer)
-        .map(|(known, m)| (known.clone(), Arc::clone(m)))
         .collect();
 
     let mut peers = Vec::new();
@@ -476,13 +474,7 @@ pub(crate) async fn run_peer_worker(state: AppState) {
         }
         was_ready = ready;
 
-        let peers: Vec<(String, Arc<Mutex<PeerState>>)> = state
-            .peer_states
-            .read()
-            .await
-            .iter()
-            .map(|(pk, m)| (pk.clone(), Arc::clone(m)))
-            .collect();
+        let peers = state.peers_snapshot().await;
 
         for (peer, peer_state_mutex) in &peers {
             let now = Instant::now();