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    /// If the pool is under capacity, try to connect one new peer.
161    /// Also spawns a background task to handle its inbound `NewPeakWallet`
162    /// messages.
163    pub async fn try_refill(&self) {
164        let current = self.entries.read().await.len();
165        if current >= self.max_peers {
166            return;
167        }
168        match connect::connect_random_peer(self.network, &self.tls, self.connect_timeout).await {
169            Ok((peer, addr, receiver)) => {
170                self.spawn_receiver_handler(receiver);
171                let mut entries = self.entries.write().await;
172                if entries.len() < self.max_peers {
173                    entries.push(PeerEntry {
174                        peer,
175                        address: addr,
176                    });
177                    log::debug!("replacement peer connected: {addr}");
178                }
179            }
180            Err(e) => log::warn!("replacement peer connect failed: {e}"),
181        }
182    }
183
184    // -----------------------------------------------------------------------
185    // Receiver helpers (handle NewPeakWallet from peers)
186    // -----------------------------------------------------------------------
187
188    /// Spawn a background task that reads inbound messages from a peer's
189    /// receiver channel and updates the shared peak height.  This mirrors
190    /// the pattern used by chia-block-listener.
191    pub fn spawn_receiver_handler(&self, mut receiver: mpsc::Receiver<Message>) {
192        let peak = Arc::clone(&self.peak_height);
193        tokio::spawn(async move {
194            while let Some(msg) = receiver.recv().await {
195                if msg.msg_type == ProtocolMessageTypes::NewPeakWallet {
196                    if let Ok(new_peak) = NewPeakWallet::from_bytes(&msg.data) {
197                        let prev = peak.fetch_max(new_peak.height, Ordering::Relaxed);
198                        if new_peak.height > prev {
199                            log::debug!("new peak from peer: {}", new_peak.height);
200                        }
201                    }
202                }
203            }
204        });
205    }
206}
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211    use crate::peer::connect::create_generated_tls;
212
213    /// `max_peers: 0` attempts no connection at all, so the pool is deterministically
214    /// empty offline — an exact, network-free fixture for the empty-pool branch.
215    async fn pool_with_no_connection_attempts(
216        requirement: PeerRequirement,
217    ) -> Result<PeerPool, ChiaQueryError> {
218        PeerPool::new(
219            NetworkType::Mainnet,
220            create_generated_tls().expect("generate a TLS identity"),
221            0,
222            requirement,
223            Duration::from_millis(1),
224        )
225        .await
226    }
227
228    /// The control: an empty pool is still fatal when nothing can serve in its place.
229    #[tokio::test]
230    async fn empty_pool_is_fatal_when_peers_are_required() {
231        assert!(matches!(
232            pool_with_no_connection_attempts(PeerRequirement::Required).await,
233            Err(ChiaQueryError::PeerDiscoveryFailed)
234        ));
235    }
236
237    /// The fix: with a fallback able to serve, an empty pool must not deny the client.
238    #[tokio::test]
239    async fn empty_pool_is_tolerated_when_peers_are_optional() {
240        let pool = pool_with_no_connection_attempts(PeerRequirement::Optional)
241            .await
242            .expect("an optional peer pool must construct with zero peers");
243        assert!(!pool.has_peers().await);
244    }
245}
chia_query/peer/pool.rs

chia_query/peer/
pool.rs
