chia_query/peer/pool.rs
1use std::net::SocketAddr;
2use std::sync::atomic::{AtomicU32, AtomicUsize, Ordering};
3use std::sync::Arc;
4use std::time::Duration;
5
6use chia_protocol::{Message, NewPeakWallet, ProtocolMessageTypes};
7use chia_traits::Streamable;
8use futures_util::stream::{FuturesUnordered, StreamExt};
9use tokio::sync::{mpsc, RwLock};
10
11use chia_wallet_sdk::client::Peer;
12use tokio_tungstenite::Connector;
13
14use crate::types::ChiaQueryError;
15use crate::NetworkType;
16
17use super::connect;
18
19// ---------------------------------------------------------------------------
20// Pool entry
21// ---------------------------------------------------------------------------
22
23struct PeerEntry {
24 peer: Peer,
25 address: SocketAddr,
26}
27
28// ---------------------------------------------------------------------------
29// PeerRequirement
30// ---------------------------------------------------------------------------
31
32/// Whether at least one peer must connect for the pool to be considered usable.
33///
34/// A client that can fall back to the coinset HTTP tier is still useful with zero
35/// peers, so failing construction on peer discovery would deny a keyless reader over a
36/// peer-tier problem it does not need (dig_ecosystem#2210).
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub enum PeerRequirement {
39 /// Peer discovery failing is fatal.
40 Required,
41 /// An empty pool is acceptable; it refills in the background.
42 Optional,
43}
44
45// ---------------------------------------------------------------------------
46// PeerPool
47// ---------------------------------------------------------------------------
48
49pub struct PeerPool {
50 entries: RwLock<Vec<PeerEntry>>,
51 next_idx: AtomicUsize,
52 max_peers: usize,
53 tls: Connector,
54 network: NetworkType,
55 connect_timeout: Duration,
56 /// Latest peak height observed from any connected peer's NewPeakWallet
57 /// messages. Updated in the background by receiver handler tasks.
58 peak_height: Arc<AtomicU32>,
59}
60
61impl PeerPool {
62 /// Spin up the pool by connecting to `max_peers` random full-node peers
63 /// concurrently. Under [`PeerRequirement::Required`] at least one peer must
64 /// succeed, otherwise we return [`ChiaQueryError::PeerDiscoveryFailed`]; under
65 /// [`PeerRequirement::Optional`] an empty pool is returned and refills later.
66 pub async fn new(
67 network: NetworkType,
68 tls: Connector,
69 max_peers: usize,
70 requirement: PeerRequirement,
71 connect_timeout: Duration,
72 ) -> Result<Self, ChiaQueryError> {
73 let peak_height = Arc::new(AtomicU32::new(0));
74
75 // Connect to peers concurrently.
76 let mut futures = FuturesUnordered::new();
77 for _ in 0..max_peers {
78 let t = tls.clone();
79 futures.push(async move {
80 connect::connect_random_peer(network, &t, connect_timeout).await
81 });
82 }
83
84 let mut initial: Vec<PeerEntry> = Vec::new();
85 let mut receivers = Vec::new();
86 while let Some(result) = futures.next().await {
87 match result {
88 Ok((peer, addr, receiver)) => {
89 initial.push(PeerEntry {
90 peer,
91 address: addr,
92 });
93 receivers.push(receiver);
94 }
95 Err(e) => log::debug!("initial peer connect failed: {e}"),
96 }
97 }
98
99 if initial.is_empty() {
100 if requirement == PeerRequirement::Required {
101 return Err(ChiaQueryError::PeerDiscoveryFailed);
102 }
103 log::warn!("no peers connected; serving from the coinset fallback until one does");
104 }
105
106 let pool = Self {
107 entries: RwLock::new(initial),
108 next_idx: AtomicUsize::new(0),
109 max_peers,
110 tls,
111 network,
112 connect_timeout,
113 peak_height,
114 };
115
116 // Spawn receiver handlers for initial peers (must happen after pool
117 // construction so peak_height Arc is available).
118 for receiver in receivers {
119 pool.spawn_receiver_handler(receiver);
120 }
121
122 Ok(pool)
123 }
124
125 /// Latest peak height observed across all connected peers.
126 /// Returns 0 if no peak has been received yet.
127 pub fn peak_height(&self) -> u32 {
128 self.peak_height.load(Ordering::Relaxed)
129 }
130
131 /// Round-robin select a peer from the pool.
132 /// Returns `None` when the pool is empty.
133 pub async fn select_peer(&self) -> Option<(Peer, SocketAddr)> {
134 let entries = self.entries.read().await;
135 if entries.is_empty() {
136 return None;
137 }
138 let idx = self.next_idx.fetch_add(1, Ordering::Relaxed) % entries.len();
139 let entry = &entries[idx];
140 Some((entry.peer.clone(), entry.address))
141 }
142
143 /// Remove a peer from the pool and asynchronously connect a replacement.
144 pub async fn eject_peer(&self, addr: SocketAddr) {
145 {
146 let mut entries = self.entries.write().await;
147 entries.retain(|e| e.address != addr);
148 }
149 log::debug!(
150 "peer ejected from pool; will refill on next request (network={:?})",
151 self.network,
152 );
153 }
154
155 /// Whether the pool has at least one usable peer.
156 pub async fn has_peers(&self) -> bool {
157 !self.entries.read().await.is_empty()
158 }
159
160 /// How many peers the pool HOLDS right now.
161 ///
162 /// This is a live count of the connections currently in the pool, not
163 /// [`max_peers`](Self::new)'s target: a pool that is still filling reports what it has, and
164 /// reports the target only once it has reached it. A caller showing this number to a user is
165 /// stating a fact about the machine, so a configured intention must never stand in for it.
166 ///
167 /// A peer is removed by [`eject_peer`](Self::eject_peer), which runs when a request to it
168 /// FAILS. So the count is of peers held and believed usable; a connection that has died
169 /// silently is still counted until something tries to use it. That is the same liveness
170 /// standard [`has_peers`](Self::has_peers) has always answered by, made countable.
171 pub async fn peer_count(&self) -> usize {
172 self.entries.read().await.len()
173 }
174
175 /// If the pool is under capacity, try to connect one new peer.
176 /// Also spawns a background task to handle its inbound `NewPeakWallet`
177 /// messages.
178 pub async fn try_refill(&self) {
179 let current = self.entries.read().await.len();
180 if current >= self.max_peers {
181 return;
182 }
183 match connect::connect_random_peer(self.network, &self.tls, self.connect_timeout).await {
184 Ok((peer, addr, receiver)) => {
185 self.spawn_receiver_handler(receiver);
186 let mut entries = self.entries.write().await;
187 if entries.len() < self.max_peers {
188 entries.push(PeerEntry {
189 peer,
190 address: addr,
191 });
192 log::debug!("replacement peer connected: {addr}");
193 }
194 }
195 Err(e) => log::warn!("replacement peer connect failed: {e}"),
196 }
197 }
198
199 // -----------------------------------------------------------------------
200 // Receiver helpers (handle NewPeakWallet from peers)
201 // -----------------------------------------------------------------------
202
203 /// Spawn a background task that reads inbound messages from a peer's
204 /// receiver channel and updates the shared peak height. This mirrors
205 /// the pattern used by chia-block-listener.
206 pub fn spawn_receiver_handler(&self, mut receiver: mpsc::Receiver<Message>) {
207 let peak = Arc::clone(&self.peak_height);
208 tokio::spawn(async move {
209 while let Some(msg) = receiver.recv().await {
210 if msg.msg_type == ProtocolMessageTypes::NewPeakWallet {
211 if let Ok(new_peak) = NewPeakWallet::from_bytes(&msg.data) {
212 let prev = peak.fetch_max(new_peak.height, Ordering::Relaxed);
213 if new_peak.height > prev {
214 log::debug!("new peak from peer: {}", new_peak.height);
215 }
216 }
217 }
218 }
219 });
220 }
221}
222
223#[cfg(test)]
224mod tests {
225 use super::*;
226 use crate::peer::connect::create_generated_tls;
227
228 /// `max_peers: 0` attempts no connection at all, so the pool is deterministically
229 /// empty offline — an exact, network-free fixture for the empty-pool branch.
230 async fn pool_with_no_connection_attempts(
231 requirement: PeerRequirement,
232 ) -> Result<PeerPool, ChiaQueryError> {
233 PeerPool::new(
234 NetworkType::Mainnet,
235 create_generated_tls().expect("generate a TLS identity"),
236 0,
237 requirement,
238 Duration::from_millis(1),
239 )
240 .await
241 }
242
243 /// The control: an empty pool is still fatal when nothing can serve in its place.
244 #[tokio::test]
245 async fn empty_pool_is_fatal_when_peers_are_required() {
246 assert!(matches!(
247 pool_with_no_connection_attempts(PeerRequirement::Required).await,
248 Err(ChiaQueryError::PeerDiscoveryFailed)
249 ));
250 }
251
252 /// The fix: with a fallback able to serve, an empty pool must not deny the client.
253 #[tokio::test]
254 async fn empty_pool_is_tolerated_when_peers_are_optional() {
255 let pool = pool_with_no_connection_attempts(PeerRequirement::Optional)
256 .await
257 .expect("an optional peer pool must construct with zero peers");
258 assert!(!pool.has_peers().await);
259 }
260
261 /// **The count is what is HELD, never what was asked for.**
262 ///
263 /// Built by hand rather than through [`PeerPool::new`] so `max_peers` can be a realistic 5
264 /// while the pool provably holds nothing — the one shape that separates a measurement from a
265 /// configured intention. A `peer_count` that returned `max_peers` would satisfy every
266 /// assertion reachable through the offline constructor, whose `max_peers` is necessarily 0,
267 /// and would then report "5 peers" on a machine holding none.
268 #[tokio::test]
269 async fn an_unfilled_pool_counts_what_it_holds_not_the_target_it_was_given() {
270 let pool = PeerPool {
271 entries: RwLock::new(Vec::new()),
272 next_idx: AtomicUsize::new(0),
273 max_peers: 5,
274 tls: create_generated_tls().expect("generate a TLS identity"),
275 network: NetworkType::Mainnet,
276 connect_timeout: Duration::from_millis(1),
277 peak_height: Arc::new(AtomicU32::new(0)),
278 };
279
280 assert_eq!(
281 pool.peer_count().await,
282 0,
283 "held is 0 while the target is 5"
284 );
285 assert!(!pool.has_peers().await);
286 }
287}