diff --git a/dash-spv-bench/src/main.rs b/dash-spv-bench/src/main.rs index f3be43593..4584f8082 100644 --- a/dash-spv-bench/src/main.rs +++ b/dash-spv-bench/src/main.rs @@ -152,9 +152,7 @@ async fn main() -> Result<()> { let wallet_probe = wallet.clone(); let handler = Arc::new(BenchEventHandler::new(dashboard.clone())); - let network = dash_spv::network::PeerNetworkManager::new(&config) - .await - .map_err(|e| anyhow!("network new: {e}"))?; + let network = dash_spv::network::PeerNetworkManager::new(&config).await; let client = DashSpvClient::new( config, diff --git a/dash-spv-ffi/src/callbacks.rs b/dash-spv-ffi/src/callbacks.rs index c9fc45ff5..360144252 100644 --- a/dash-spv-ffi/src/callbacks.rs +++ b/dash-spv-ffi/src/callbacks.rs @@ -349,6 +349,7 @@ impl FFISyncEventCallbacks { start_height, end_height, tip_height, + .. } => { if let Some(cb) = self.on_filter_headers_stored { cb(*start_height, *end_height, *tip_height, self.user_data); @@ -550,17 +551,13 @@ impl FFINetworkEventCallbacks { use dash_spv::network::NetworkEvent; match event { - NetworkEvent::PeerConnected { - address, - } => { + NetworkEvent::PeerConnected(address) => { if let Some(cb) = self.on_peer_connected { let c_addr = CString::new(address.to_string()).unwrap_or_default(); cb(c_addr.as_ptr(), self.user_data); } } - NetworkEvent::PeerDisconnected { - address, - } => { + NetworkEvent::PeerDisconnected(address) => { if let Some(cb) = self.on_peer_disconnected { let c_addr = CString::new(address.to_string()).unwrap_or_default(); cb(c_addr.as_ptr(), self.user_data); @@ -569,10 +566,9 @@ impl FFINetworkEventCallbacks { NetworkEvent::PeersUpdated { connected_count, best_height, - .. } => { if let Some(cb) = self.on_peers_updated { - cb(*connected_count as u32, best_height.unwrap_or(0), self.user_data); + cb(*connected_count, *best_height, self.user_data); } } } diff --git a/dash-spv-ffi/src/client.rs b/dash-spv-ffi/src/client.rs index 33b334ef9..2fea970e4 100644 --- a/dash-spv-ffi/src/client.rs +++ b/dash-spv-ffi/src/client.rs @@ -82,15 +82,15 @@ pub unsafe extern "C" fn dash_spv_ffi_client_new( let client_result = runtime.block_on(async move { // Construct concrete implementations for generics - let network = dash_spv::network::PeerNetworkManager::new(&client_config).await; let storage = DiskStorageManager::new(&client_config).await; let wallet = key_wallet_manager::WalletManager::< key_wallet::wallet::managed_wallet_info::ManagedWalletInfo, >::new(client_config.network); let wallet = std::sync::Arc::new(tokio::sync::RwLock::new(wallet)); - match (network, storage) { - (Ok(network), Ok(storage)) => { + match storage { + Ok(storage) => { + let network = dash_spv::network::PeerNetworkManager::new(&client_config).await; DashSpvClient::new( client_config, network, @@ -100,8 +100,7 @@ pub unsafe extern "C" fn dash_spv_ffi_client_new( ) .await } - (Err(e), _) => Err(e), - (_, Err(e)) => Err(dash_spv::SpvError::Storage(e)), + Err(e) => Err(dash_spv::SpvError::Storage(e)), } }); diff --git a/dash-spv/Cargo.toml b/dash-spv/Cargo.toml index 0c80f5b36..c8b470944 100644 --- a/dash-spv/Cargo.toml +++ b/dash-spv/Cargo.toml @@ -10,7 +10,7 @@ rust-version = "1.89" [dependencies] # Core Dash libraries -dashcore = { path = "../dash", features = ["serde", "core-block-hash-use-x11", "message_verification", "bls", "quorum_validation"] } +dashcore = { path = "../dash", features = ["serde", "core-block-hash-use-x11", "message_verification", "bls", "quorum_validation", "tokio"] } dashcore_hashes = { path = "../hashes" } dash-network-seeds = { path = "../dash-network-seeds" } key-wallet = { path = "../key-wallet" } diff --git a/dash-spv/examples/filter_sync.rs b/dash-spv/examples/filter_sync.rs index 203896ce6..3b95f4975 100644 --- a/dash-spv/examples/filter_sync.rs +++ b/dash-spv/examples/filter_sync.rs @@ -26,7 +26,7 @@ async fn main() -> Result<(), Box> { .without_masternodes(); // Skip masternode sync for this example // Create network manager - let network_manager = PeerNetworkManager::new(&config).await?; + let network_manager = PeerNetworkManager::new(&config).await; // Create storage manager let storage_manager = DiskStorageManager::new(&config).await?; diff --git a/dash-spv/examples/simple_sync.rs b/dash-spv/examples/simple_sync.rs index 0568768fe..8c2cd20e7 100644 --- a/dash-spv/examples/simple_sync.rs +++ b/dash-spv/examples/simple_sync.rs @@ -21,7 +21,7 @@ async fn main() -> Result<(), Box> { .without_masternodes(); // Skip masternode sync for this example // Create network manager - let network_manager = PeerNetworkManager::new(&config).await?; + let network_manager = PeerNetworkManager::new(&config).await; // Create storage manager let storage_manager = DiskStorageManager::new(&config).await?; diff --git a/dash-spv/examples/spv_with_wallet.rs b/dash-spv/examples/spv_with_wallet.rs index 2d2c661d8..3cf81ca45 100644 --- a/dash-spv/examples/spv_with_wallet.rs +++ b/dash-spv/examples/spv_with_wallet.rs @@ -21,7 +21,7 @@ async fn main() -> Result<(), Box> { .with_validation_mode(dash_spv::ValidationMode::Full); // Create network manager - let network_manager = PeerNetworkManager::new(&config).await?; + let network_manager = PeerNetworkManager::new(&config).await; // Create storage manager - use disk storage for persistence let storage_manager = DiskStorageManager::new(&config).await?; diff --git a/dash-spv/src/client/core.rs b/dash-spv/src/client/core.rs index 09e8bb712..fcf0e6651 100644 --- a/dash-spv/src/client/core.rs +++ b/dash-spv/src/client/core.rs @@ -105,7 +105,7 @@ pub(super) type PersistentSyncCoordinator = SyncCoordinator< /// The generic design is an intentional, beneficial architectural choice for a library. pub struct DashSpvClient { pub(super) config: Arc>, - pub(super) network: Arc>, + pub(super) network: Arc, pub(super) storage: Arc>, /// External wallet implementation (required) pub(super) wallet: Arc>, @@ -114,6 +114,7 @@ pub struct DashSpvClient>, + pub(super) stop_requested: Arc, pub(super) event_handlers: Arc>>, } @@ -127,6 +128,7 @@ impl Clone for DashSpv masternode_engine: self.masternode_engine.clone(), sync_coordinator: Arc::clone(&self.sync_coordinator), running: Arc::clone(&self.running), + stop_requested: Arc::clone(&self.stop_requested), event_handlers: Arc::clone(&self.event_handlers), } } diff --git a/dash-spv/src/client/event_handler.rs b/dash-spv/src/client/event_handler.rs index 1e4b57860..b2282d64c 100644 --- a/dash-spv/src/client/event_handler.rs +++ b/dash-spv/src/client/event_handler.rs @@ -327,8 +327,7 @@ mod tests { handler.on_sync_event(&event); handler.on_network_event(&NetworkEvent::PeersUpdated { connected_count: 0, - addresses: vec![], - best_height: None, + best_height: 0, }); handler.on_progress(&SyncProgress::default()); handler.on_error("test error"); @@ -533,14 +532,8 @@ mod tests { ); let addr: SocketAddr = "127.0.0.1:9999".parse().unwrap(); - tx.send(NetworkEvent::PeerConnected { - address: addr, - }) - .unwrap(); - tx.send(NetworkEvent::PeerDisconnected { - address: addr, - }) - .unwrap(); + tx.send(NetworkEvent::PeerConnected(addr)).unwrap(); + tx.send(NetworkEvent::PeerDisconnected(addr)).unwrap(); tokio::time::sleep(std::time::Duration::from_millis(50)).await; shutdown.cancel(); diff --git a/dash-spv/src/client/events.rs b/dash-spv/src/client/events.rs index 6ed109ed9..78d072fde 100644 --- a/dash-spv/src/client/events.rs +++ b/dash-spv/src/client/events.rs @@ -32,6 +32,6 @@ impl DashSpvClient broadcast::Receiver { - self.network.lock().await.subscribe_network_events() + self.network.events() } } diff --git a/dash-spv/src/client/lifecycle.rs b/dash-spv/src/client/lifecycle.rs index 46e26f71c..93b6ad296 100644 --- a/dash-spv/src/client/lifecycle.rs +++ b/dash-spv/src/client/lifecycle.rs @@ -156,12 +156,13 @@ impl DashSpvClient DashSpvClient = self.network.clone(); + if let Err(e) = self.sync_coordinator.lock().await.start(&network).await { tracing::error!("Failed to start sync coordinator: {}", e); return Err(SpvError::Sync(e)); } - // Connect to network - self.network.lock().await.connect().await?; - - // Only mark as running after all startup operations succeed. - // `send_replace` always stores the value regardless of receiver count, - // so this is correct even when `run()` has not subscribed yet. - self.running.send_replace(true); + self.network.start(); + + // Only mark as running after all startup operations succeed — and only if + // no `stop()` raced in while we were connecting. The check runs inside the + // watch lock (via `send_if_modified`), and `stop()` sets `stop_requested` + // before it flips `running`, so the two orderings are both safe: + // - we win the lock first: set running=true; a later stop() flips it false. + // - stop() won: `stop_requested` is already true here, so we leave running + // false and the run loop tears down immediately instead of syncing forever. + // `send_if_modified` stores the value regardless of receiver count, so this + // is correct even when `run()` has not subscribed yet. + self.running.send_if_modified(|running| { + if self.stop_requested.load(std::sync::atomic::Ordering::SeqCst) { + false + } else { + *running = true; + true + } + }); Ok(()) } /// Stop the SPV client. pub async fn stop(&self) -> Result<()> { - // Check if already stopped - if !*self.running.borrow() { - return Ok(()); - } + // Record the stop request BEFORE flipping `running`, so a `start()` still + // connecting observes it (under the watch lock) and declines to mark the + // client running. Otherwise a stop that arrives mid-startup would be lost: + // `start()` would flip running true afterwards and the run task would sync + // forever, hanging `run_handle.await`. + self.stop_requested.store(true, std::sync::atomic::Ordering::SeqCst); // Flip the running state before tearing anything down so a concurrent // `run()` loop wakes immediately and breaks out before it can lock the @@ -216,14 +238,23 @@ impl DashSpvClient DashSpvClient usize { - self.network.lock().await.peer_count() - } - - /// Disconnect a specific peer. - pub async fn disconnect_peer(&self, addr: &std::net::SocketAddr, reason: &str) -> Result<()> { - Ok(self.network.lock().await.disconnect_peer(addr, reason).await?) + self.network.connected_count().await as usize } // ============ Masternode Queries ============ diff --git a/dash-spv/src/client/transactions.rs b/dash-spv/src/client/transactions.rs index 94847cea5..ae998912c 100644 --- a/dash-spv/src/client/transactions.rs +++ b/dash-spv/src/client/transactions.rs @@ -30,21 +30,19 @@ impl DashSpvClient Result<()> { - let network_guard = self.network.lock().await; - - if network_guard.peer_count() == 0 { + if self.network.connected_count().await == 0 { return Err(SpvError::Network(NetworkError::NotConnected)); } if !self.config.read().await.enable_mempool_tracking { // Legacy untracked path: fan out to every peer. - network_guard.broadcast(NetworkMessage::Tx(tx.clone())).await?; + self.network.broadcast(NetworkMessage::Tx(tx.clone())); } // Inject locally so the mempool manager picks it up through handle_tx. // With tracking enabled the manager performs the actual (targeted) // network send when it processes this message. - network_guard.dispatch_local(NetworkMessage::Tx(tx.clone())).await; + self.network.dispatch_local(NetworkMessage::Tx(tx.clone())).await; Ok(()) } diff --git a/dash-spv/src/lib.rs b/dash-spv/src/lib.rs index b0263f08e..6dce0560a 100644 --- a/dash-spv/src/lib.rs +++ b/dash-spv/src/lib.rs @@ -28,7 +28,7 @@ //! .with_storage_path("./.tmp/example-storage"); //! //! // Create the required components -//! let network = PeerNetworkManager::new(&config).await?; +//! let network = PeerNetworkManager::new(&config).await; //! let storage = DiskStorageManager::new(&config).await?; //! let wallet = Arc::new(RwLock::new(WalletManager::::new(config.network))); //! diff --git a/dash-spv/src/main.rs b/dash-spv/src/main.rs index eca7d7a6a..f0a52c944 100644 --- a/dash-spv/src/main.rs +++ b/dash-spv/src/main.rs @@ -259,13 +259,7 @@ async fn run() -> Result<(), Box> { let wallet = Arc::new(tokio::sync::RwLock::new(wallet_manager)); // Create network manager - let network_manager = match dash_spv::network::manager::PeerNetworkManager::new(&config).await { - Ok(nm) => nm, - Err(e) => { - eprintln!("Failed to create network manager: {}", e); - process::exit(1); - } - }; + let network_manager = dash_spv::network::PeerNetworkManager::new(&config).await; let storage_manager = match dash_spv::storage::DiskStorageManager::new(&config).await { Ok(sm) => sm, @@ -382,7 +376,7 @@ fn parse_llmq_devnet_params(raw: &str) -> Result { async fn run_client( config: ClientConfig, - network_manager: dash_spv::network::manager::PeerNetworkManager, + network_manager: dash_spv::network::PeerNetworkManager, storage_manager: S, wallet: Arc>>, ) -> Result<(), Box> { @@ -390,7 +384,7 @@ async fn run_client( let client = match DashSpvClient::< WalletManager, - dash_spv::network::manager::PeerNetworkManager, + dash_spv::network::PeerNetworkManager, S, >::new( config.clone(), network_manager, storage_manager, wallet.clone(), Vec::new() diff --git a/dash-spv/src/network/addrv2.rs b/dash-spv/src/network/addrv2.rs deleted file mode 100644 index 839f8a0d9..000000000 --- a/dash-spv/src/network/addrv2.rs +++ /dev/null @@ -1,236 +0,0 @@ -//! AddrV2 message handling for modern peer exchange protocol - -use rand::prelude::*; -use std::collections::{HashMap, HashSet}; -use std::net::SocketAddr; -use std::sync::Arc; -use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use tokio::sync::RwLock; - -use dashcore::network::address::{AddrV2, AddrV2Message}; -use dashcore::network::constants::ServiceFlags; -use dashcore::network::message::NetworkMessage; - -use crate::network::constants::{MAX_ADDR_TO_SEND, MAX_ADDR_TO_STORE}; - -const ONE_WEEK: u32 = 7 * 24 * 60 * 60; -const TEN_MINUTES: u32 = 600; - -/// Evict oldest entries if the map exceeds capacity, keeping the freshest addresses. -fn evict_if_needed(peers: &mut HashMap) { - if peers.len() > MAX_ADDR_TO_STORE { - let mut entries: Vec<_> = peers.drain().collect(); - entries.sort_by_key(|(_, msg)| std::cmp::Reverse(msg.time)); - entries.truncate(MAX_ADDR_TO_STORE); - peers.extend(entries); - } -} - -/// Handler for AddrV2 peer exchange protocol -pub struct AddrV2Handler { - /// Known peer addresses from AddrV2 messages - known_peers: Arc>>, - /// Peers that support AddrV2 - supports_addrv2: Arc>>, -} - -impl AddrV2Handler { - /// Create a new AddrV2 handler - pub fn new() -> Self { - Self { - known_peers: Arc::new(RwLock::new(HashMap::new())), - supports_addrv2: Arc::new(RwLock::new(HashSet::new())), - } - } - - /// Handle SendAddrV2 message indicating peer support - pub async fn handle_sendaddrv2(&self, peer_addr: SocketAddr) { - self.supports_addrv2.write().await.insert(peer_addr); - tracing::debug!("Peer {} supports AddrV2", peer_addr); - } - - /// Handle incoming AddrV2 messages - pub async fn handle_addrv2(&self, messages: Vec) { - let mut known_peers = self.known_peers.write().await; - let now = SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_else(|e| { - tracing::error!("System time error in handle_addrv2: {}", e); - Duration::from_secs(0) - }) - .as_secs() as u32; - - let received = messages.len(); - let mut added = 0; - let mut updated = 0; - - for msg in messages { - // Accept addresses seen within the last week. Older addresses are likely stale. - // Also, reject timestamps more than 10 minutes in the future which are invalid. - if msg.time < now.saturating_sub(ONE_WEEK) || msg.time > now + TEN_MINUTES { - tracing::trace!("Ignoring AddrV2 with invalid timestamp: {}", msg.time); - continue; - } - - let Ok(socket_addr) = msg.socket_addr() else { - continue; - }; - - // Only update if new or has fresher timestamp - match known_peers.get(&socket_addr) { - Some(existing) if existing.time >= msg.time => continue, - Some(_) => updated += 1, - None => added += 1, - } - known_peers.insert(socket_addr, msg); - } - - evict_if_needed(&mut known_peers); - - tracing::info!( - "Processed AddrV2 messages: received {}, added {}, updated {}, total known peers: {}", - received, - added, - updated, - known_peers.len() - ); - } - - /// Get addresses to share with a peer - pub async fn get_addresses_for_peer(&self, count: usize) -> Vec { - let known_peers = self.known_peers.read().await; - - if known_peers.is_empty() { - return vec![]; - } - - // Select random subset - let mut rng = thread_rng(); - let count = count.min(MAX_ADDR_TO_SEND).min(known_peers.len()); - - let addresses: Vec = - known_peers.values().choose_multiple(&mut rng, count).into_iter().cloned().collect(); - - addresses - } - - /// Check if a peer supports AddrV2 - pub async fn peer_supports_addrv2(&self, addr: &SocketAddr) -> bool { - self.supports_addrv2.read().await.contains(addr) - } - - /// Get all known socket addresses - pub async fn get_known_addresses(&self) -> Vec { - self.known_peers.read().await.values().cloned().collect() - } - - /// Add a known peer address - pub async fn add_known_address(&self, addr: SocketAddr, services: ServiceFlags) { - let now = SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_else(|e| { - tracing::error!("System time error in add_known_address: {}", e); - Duration::from_secs(0) - }) - .as_secs() as u32; - - let addr_v2 = match addr.ip() { - std::net::IpAddr::V4(ipv4) => AddrV2::Ipv4(ipv4), - std::net::IpAddr::V6(ipv6) => AddrV2::Ipv6(ipv6), - }; - - let addr_msg = AddrV2Message { - time: now, - services, - addr: addr_v2, - port: addr.port(), - }; - - let mut known_peers = self.known_peers.write().await; - known_peers.insert(addr, addr_msg); - evict_if_needed(&mut known_peers); - } - - /// Build a GetAddr response message - pub async fn build_addr_response(&self) -> NetworkMessage { - let addresses = self.get_addresses_for_peer(23).await; // Bitcoin typically sends ~23 addresses - NetworkMessage::AddrV2(addresses) - } -} - -impl Default for AddrV2Handler { - fn default() -> Self { - Self::new() - } -} - -#[cfg(test)] -mod tests { - use super::*; - use dashcore::network::address::AddrV2; - - #[tokio::test] - async fn test_addrv2_handler_basic() { - let handler = AddrV2Handler::new(); - - // Test SendAddrV2 support tracking - let peer = "127.0.0.1:9999".parse().expect("Failed to parse test peer address"); - handler.handle_sendaddrv2(peer).await; - assert!(handler.peer_supports_addrv2(&peer).await); - - // Test adding known address - let addr = "192.168.1.1:9999".parse().expect("Failed to parse test address"); - handler.add_known_address(addr, ServiceFlags::NETWORK).await; - - let known = handler.get_known_addresses().await; - assert_eq!(known.len(), 1); - assert_eq!(known[0].socket_addr().unwrap(), addr); - } - - #[tokio::test] - async fn test_addrv2_timestamp_validation() { - let handler = AddrV2Handler::new(); - let now = SystemTime::now() - .duration_since(UNIX_EPOCH) - .expect("Failed to get system time in test") - .as_secs() as u32; - - // Create test messages with various timestamps - let addr: SocketAddr = - "127.0.0.1:9999".parse().expect("Failed to parse test socket address"); - let ipv4_addr = match addr.ip() { - std::net::IpAddr::V4(v4) => v4, - _ => panic!("Test expects IPv4 address but got IPv6"), - }; - - let messages = vec![ - // Valid: current time - AddrV2Message { - time: now, - services: ServiceFlags::NETWORK, - addr: AddrV2::Ipv4(ipv4_addr), - port: addr.port(), - }, - // Invalid: too old (4 hours ago) - AddrV2Message { - time: now.saturating_sub(14400), - services: ServiceFlags::NETWORK, - addr: AddrV2::Ipv4(ipv4_addr), - port: addr.port(), - }, - // Invalid: too far in future (20 minutes) - AddrV2Message { - time: now + 1200, - services: ServiceFlags::NETWORK, - addr: AddrV2::Ipv4(ipv4_addr), - port: addr.port(), - }, - ]; - - handler.handle_addrv2(messages).await; - - // Only the valid message should be stored - let known = handler.get_known_addresses().await; - assert_eq!(known.len(), 1); - } -} diff --git a/dash-spv/src/network/constants.rs b/dash-spv/src/network/constants.rs deleted file mode 100644 index 30928f70e..000000000 --- a/dash-spv/src/network/constants.rs +++ /dev/null @@ -1,26 +0,0 @@ -//! Network constants for peer support - -use std::time::Duration; - -// Timeouts -pub const CONNECTION_TIMEOUT: Duration = Duration::from_secs(30); -pub const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10); -pub const MESSAGE_TIMEOUT: Duration = Duration::from_secs(120); -pub const PING_INTERVAL: Duration = Duration::from_secs(120); - -// Reconnection -pub const RECONNECT_DELAY: Duration = Duration::from_secs(5); -pub const MAX_RECONNECT_ATTEMPTS: u32 = 3; - -// Peer exchange -pub const MAX_ADDR_TO_SEND: usize = 1000; -pub const MAX_ADDR_TO_STORE: usize = 2000; - -// Connection maintenance -pub const MAINTENANCE_INTERVAL: Duration = Duration::from_secs(10); // Check more frequently -pub const PEER_DISCOVERY_INTERVAL: Duration = Duration::from_secs(60); // Discover more frequently - -// DNS and polling intervals -pub const DNS_DISCOVERY_DELAY: Duration = Duration::from_secs(10); -pub const MESSAGE_POLL_INTERVAL: Duration = Duration::from_millis(10); -pub const MESSAGE_RECEIVE_TIMEOUT: Duration = Duration::from_millis(100); diff --git a/dash-spv/src/network/discovery.rs b/dash-spv/src/network/discovery.rs index 035bc7f28..e43c4cdb6 100644 --- a/dash-spv/src/network/discovery.rs +++ b/dash-spv/src/network/discovery.rs @@ -1,111 +1,98 @@ -//! Peer discovery for Dash network. -//! -//! Peer discovery is seeded from two sources, in priority order: -//! -//! 1. A hardcoded masternode IP list for the network, embedded at compile time -//! from `dash-spv/seeds/.txt`. This file is regenerated weekly by -//! CI from a live Dash Core node (see `masternode-seeds-fetcher`). -//! 2. DNS seed queries as a backup. DNS resolution failures are logged but are -//! not fatal — as long as the embedded list yields at least one peer, the -//! client can bootstrap. -//! -//! Results from both sources are merged and deduplicated. +use std::net::SocketAddr; use dashcore::Network; -use std::net::SocketAddr; +use rand::seq::SliceRandom; -/// DNS discovery for finding initial peers. -/// -/// Despite the name (kept for backwards compatibility), this type also returns -/// hardcoded masternode seeds embedded at compile time; DNS is used as a -/// fallback. -#[derive(Default)] -pub struct DnsDiscovery {} - -impl DnsDiscovery { - /// Create a new DNS discovery instance - pub fn new() -> Self { - Self {} - } +use crate::network::peer::DisconnectedPeer; +use crate::ClientConfig; + +const DNS_LOOKUP_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5); - /// Discover peers for the given network. +pub struct PeerDiscoverer { + network: Network, + // Empty means "discover" from the compiled-in seeds, then DNS. + fixed: Vec, + restrict_to_configured_peers: bool, + /// Discovered addresses, resolved once and then kept. /// - /// Returns the union of the embedded hardcoded masternode seeds and any - /// addresses resolved via DNS. DNS resolution failures are logged at warn - /// level but do not cause this function to fail — the embedded list acts - /// as the primary source and DNS is a best-effort backup. - pub async fn discover_peers(&self, network: Network) -> Vec { - let seeds = network.dns_seeds(); - let port = network.default_p2p_port(); - let mut addresses = dash_network_seeds::addresses(network); + /// Deliberately not consumed as it is handed out: the reconnector comes back here + /// every time the peer set drops, and a pool that drained itself would leave a client + /// with one known peer unable to reconnect after its second disconnect. + discovered: Option>, +} - let embedded_count = addresses.len(); - tracing::info!("Loaded {} hardcoded masternode seed(s) for {:?}", embedded_count, network); +impl PeerDiscoverer { + pub fn new(config: &ClientConfig) -> PeerDiscoverer { + PeerDiscoverer { + network: config.network, + fixed: config.peers.clone(), + restrict_to_configured_peers: config.restrict_to_configured_peers, + discovered: None, + } + } + + /// Up to `count` addresses to try, sampled at random from whatever source applies. + pub async fn get(&mut self, count: usize) -> Vec { + let pool = if !self.fixed.is_empty() { + &self.fixed + } else if self.restrict_to_configured_peers { + return Vec::new(); + } else { + if self.discovered.is_none() { + let found = Self::discover(self.network).await; + self.discovered = Some(found); + } + self.discovered.as_ref().expect("just set") + }; - for seed in seeds { - tracing::debug!("Querying DNS seed: {}", seed); + pool.choose_multiple(&mut rand::thread_rng(), count) + .map(|addr| DisconnectedPeer::new(*addr, self.network)) + .collect() + } + + /// Addresses to try: the compiled-in seeds, then DNS + async fn discover(network: Network) -> Vec { + let mut addresses = dash_network_seeds::addresses(network); - match tokio::net::lookup_host((*seed, port)).await { - Ok(iter) => { + let port = network.default_p2p_port(); + for seed in network.dns_seeds() { + match tokio::time::timeout(DNS_LOOKUP_TIMEOUT, tokio::net::lookup_host((*seed, port))) + .await + { + Ok(Ok(iter)) => { let resolved: Vec = iter.collect(); tracing::info!("DNS seed {} returned {} addresses", seed, resolved.len()); addresses.extend(resolved); } - Err(e) => { - // DNS is a best-effort backup; do not propagate the error. + Ok(Err(e)) => { tracing::warn!("Failed to resolve DNS seed {} (backup source): {}", seed, e); } + Err(_) => { + tracing::warn!( + "DNS seed {} did not resolve within {:?} (backup source)", + seed, + DNS_LOOKUP_TIMEOUT + ); + } } } addresses.sort(); addresses.dedup(); - tracing::info!( - "Discovered {} unique peer addresses for {:?} ({} from embedded seeds + DNS)", - addresses.len(), - network, - embedded_count - ); addresses } - - /// Discover peers with a limit on the number returned - pub async fn discover_peers_limited(&self, network: Network, limit: usize) -> Vec { - let mut peers = self.discover_peers(network).await; - peers.truncate(limit); - peers - } } #[cfg(test)] mod tests { use super::*; - #[tokio::test] - #[ignore] // Requires network access - async fn test_dns_discovery_mainnet() { - let discovery = DnsDiscovery::new(); - let peers = discovery.discover_peers(Network::Mainnet).await; - - // Print discovered peers for debugging - println!("Discovered {} mainnet peers:", peers.len()); - for peer in &peers { - println!(" {}", peer); - } - - // All peers should use the correct port - for peer in &peers { - assert_eq!(peer.port(), Network::Mainnet.default_p2p_port()); - } - } - #[tokio::test] async fn test_dns_discovery_testnet_returns_embedded_when_dns_fails() { // This test does not require network access: even if DNS resolution // fails, the embedded seed file must yield peers. - let discovery = DnsDiscovery::new(); - let peers = discovery.discover_peers(Network::Testnet).await; + let peers = PeerDiscoverer::discover(Network::Testnet).await; assert!( peers.len() >= 29, @@ -119,8 +106,7 @@ mod tests { #[tokio::test] async fn test_dns_discovery_regtest() { - let discovery = DnsDiscovery::new(); - let peers = discovery.discover_peers(Network::Regtest).await; + let peers = PeerDiscoverer::discover(Network::Regtest).await; // Should return empty for regtest (no DNS seeds and no embedded list) assert!(peers.is_empty()); diff --git a/dash-spv/src/network/event.rs b/dash-spv/src/network/event.rs deleted file mode 100644 index 397ebd862..000000000 --- a/dash-spv/src/network/event.rs +++ /dev/null @@ -1,62 +0,0 @@ -//! Network event system for peer connection state changes. -//! -//! This module provides events for network layer changes that sync managers -//! need to react to, such as peer connections and disconnections. - -use dashcore::prelude::CoreBlockHeight; -use std::fmt; -use std::net::SocketAddr; - -/// Events emitted by the network layer. -/// -/// These events inform sync managers about network state changes, -/// allowing them to wait for connections before sending requests. -#[derive(Debug, Clone)] -pub enum NetworkEvent { - /// A peer has connected. - PeerConnected { - /// Socket address of the connected peer. - address: SocketAddr, - }, - - /// A peer has disconnected. - PeerDisconnected { - /// Socket address of the disconnected peer. - address: SocketAddr, - }, - - /// Summary of connected peers (emitted after connect/disconnect). - /// - /// This event provides the current state of connections after any change. - PeersUpdated { - /// Number of currently connected peers. - connected_count: usize, - /// Addresses of all connected peers. - addresses: Vec, - /// Best height of connected peers. - best_height: Option, - }, -} - -impl fmt::Display for NetworkEvent { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - NetworkEvent::PeerConnected { - address, - } => write!(f, "PeerConnected({})", address), - NetworkEvent::PeerDisconnected { - address, - } => write!(f, "PeerDisconnected({})", address), - NetworkEvent::PeersUpdated { - connected_count, - addresses: _, - best_height, - } => write!( - f, - "PeersUpdated(connected={}, best_height={})", - connected_count, - best_height.unwrap_or(0) - ), - } - } -} diff --git a/dash-spv/src/network/handshake.rs b/dash-spv/src/network/handshake.rs deleted file mode 100644 index 3fc71bd5c..000000000 --- a/dash-spv/src/network/handshake.rs +++ /dev/null @@ -1,320 +0,0 @@ -//! Network handshake management. - -use std::net::SocketAddr; -use std::time::{Duration, SystemTime, UNIX_EPOCH}; - -use dashcore::network::constants; -use dashcore::network::constants::{ServiceFlags, NODE_HEADERS_COMPRESSED}; -use dashcore::network::message::NetworkMessage; -use dashcore::network::message_network::VersionMessage; -use dashcore::Network; -// Hash trait not needed in current implementation - -use crate::error::{NetworkError, NetworkResult}; -use crate::network::peer::Peer; -use crate::network::Message; - -/// Handshake state. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum HandshakeState { - /// Initial state. - Init, - /// Version message sent. - VersionSent, - /// Version received and verack sent. - VersionReceivedVerackSent, - /// Verack received. - VerackReceived, - /// Handshake complete. - Complete, -} - -/// Manages the network handshake process. -pub struct HandshakeManager { - _network: Network, - state: HandshakeState, - our_version: u32, - peer_version: Option, - peer_services: Option, - version_received: bool, - verack_received: bool, - version_sent: bool, - user_agent: Option, -} - -impl HandshakeManager { - /// Create a new handshake manager. - pub fn new(network: Network, user_agent: Option) -> Self { - Self { - _network: network, - state: HandshakeState::Init, - our_version: constants::PROTOCOL_VERSION, - peer_version: None, - peer_services: None, - version_received: false, - verack_received: false, - version_sent: false, - user_agent, - } - } - - /// Perform the handshake with a peer. - pub async fn perform_handshake(&mut self, connection: &mut Peer) -> NetworkResult<()> { - use tokio::time::{timeout, Duration}; - - // Send version message - self.send_version(connection).await?; - self.version_sent = true; - self.state = HandshakeState::VersionSent; - tracing::info!("Handshake initiated - version message sent to peer"); - - // Define timeout for the entire handshake process - const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10); - const MESSAGE_POLL_INTERVAL: Duration = Duration::from_millis(100); - - let start_time = tokio::time::Instant::now(); - - // Wait for responses with timeout - loop { - // Check if we've exceeded the overall handshake timeout - if start_time.elapsed() > HANDSHAKE_TIMEOUT { - tracing::error!( - "Handshake timeout after {}s - version_received={}, verack_received={}", - HANDSHAKE_TIMEOUT.as_secs(), - self.version_received, - self.verack_received - ); - return Err(NetworkError::Timeout); - } - - // Try to receive a message with a short timeout - match timeout(MESSAGE_POLL_INTERVAL, connection.receive_message()).await { - Ok(Ok(Some(message))) => { - tracing::debug!("Received message during handshake: {:?}", message.cmd()); - match self.handle_handshake_message(connection, &message).await? { - Some(HandshakeState::Complete) => { - self.state = HandshakeState::Complete; - break; - } - _ => { - // Continue immediately to check for more messages in the buffer - // Don't add any delays here as multiple messages may be waiting - continue; - } - } - } - Ok(Ok(None)) => { - // No message available, continue immediately - // The read timeout already provides the necessary delay - continue; - } - Ok(Err(e)) => { - tracing::error!("Error receiving message during handshake: {}", e); - return Err(e); - } - Err(_) => { - // Timeout on receive_message, continue to check overall timeout - continue; - } - } - } - - tracing::info!( - "Handshake completed successfully - version_received={}, verack_received={}", - self.version_received, - self.verack_received - ); - Ok(()) - } - - /// Reset the handshake state. - pub fn reset(&mut self) { - self.state = HandshakeState::Init; - self.peer_version = None; - self.version_received = false; - self.verack_received = false; - self.version_sent = false; - } - - /// Handle a handshake message. - async fn handle_handshake_message( - &mut self, - connection: &mut Peer, - message: &Message, - ) -> NetworkResult> { - match message.inner() { - NetworkMessage::Version(version_msg) => { - tracing::debug!( - "Peer {} sent version message: {:?}", - message.peer_address(), - version_msg - ); - self.peer_version = Some(version_msg.version); - self.peer_services = Some(version_msg.services); - self.version_received = true; - - // Update connection's peer information - connection.update_peer_info(version_msg); - - // If we haven't sent our version yet (peer initiated), send it now - if !self.version_sent { - tracing::debug!( - "Peer {} initiated handshake, sending our version", - message.peer_address() - ); - self.send_version(connection).await?; - self.version_sent = true; - } - - // Send SendAddrV2 first to signal support (must be before verack!) - tracing::debug!("Sending sendaddrv2 to signal AddrV2 support"); - connection.send_message(NetworkMessage::SendAddrV2).await?; - - // Then send verack - tracing::debug!("Sending verack in response to version"); - connection.send_message(NetworkMessage::Verack).await?; - tracing::debug!( - "Sent verack, version_received={}, verack_received={}", - self.version_received, - self.verack_received - ); - - // Update state - self.state = HandshakeState::VersionReceivedVerackSent; - - // Check if handshake is complete (both version and verack received) - if self.version_received && self.verack_received { - tracing::info!("Handshake complete - both version and verack exchanged!"); - - // Negotiate headers2 support - self.negotiate_headers2(connection).await?; - - return Ok(Some(HandshakeState::Complete)); - } - - Ok(None) - } - NetworkMessage::Verack => { - tracing::debug!("Received verack message, current state: {:?}", self.state); - self.verack_received = true; - - // Update state - if self.state == HandshakeState::VersionSent { - self.state = HandshakeState::VerackReceived; - } - - // Check if handshake is complete (both version and verack received) - if self.version_received && self.verack_received { - tracing::info!("Handshake complete - both version and verack exchanged!"); - - // Negotiate headers2 support - self.negotiate_headers2(connection).await?; - - return Ok(Some(HandshakeState::Complete)); - } else { - tracing::debug!( - "Verack received but handshake not complete: version_received={}, verack_received={}", - self.version_received, self.verack_received - ); - } - Ok(None) - } - NetworkMessage::Ping(nonce) => { - // Respond to ping during handshake - tracing::debug!("Responding to ping during handshake: {}", nonce); - connection.send_message(NetworkMessage::Pong(*nonce)).await?; - Ok(None) - } - NetworkMessage::SendAddrV2 => { - // Peer supports AddrV2 - tracing::debug!("Peer signaled AddrV2 support"); - Ok(None) - } - _ => { - // Ignore other messages during handshake - tracing::debug!("Ignoring message during handshake: {:?}", message); - Ok(None) - } - } - } - - /// Send version message. - async fn send_version(&mut self, connection: &mut Peer) -> NetworkResult<()> { - let version_message = self.build_version_message(connection.address())?; - connection.send_message(NetworkMessage::Version(version_message)).await?; - tracing::debug!("Sent version message"); - Ok(()) - } - - /// Build version message. - fn build_version_message(&self, address: SocketAddr) -> NetworkResult { - let timestamp = SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or(Duration::from_secs(0)) - .as_secs() as i64; - - // Advertise headers2 support (NODE_HEADERS_COMPRESSED) - let services = ServiceFlags::NONE | NODE_HEADERS_COMPRESSED; - - // Parse the local address safely - let local_addr = "127.0.0.1:0" - .parse() - .map_err(|_| NetworkError::AddressParse("Failed to parse local address".to_string()))?; - - // Determine user agent: prefer configured value, else default to crate/version. - let default_agent = format!("/rust-dash-spv:{}/", env!("CARGO_PKG_VERSION")); - let mut ua = self.user_agent.clone().unwrap_or(default_agent); - // Normalize: ensure it starts and ends with '/'; trim if excessively long. - if !ua.starts_with('/') { - ua.insert(0, '/'); - } - if !ua.ends_with('/') { - ua.push('/'); - } - // Keep within a reasonable bound (match peer validation bound of 256) - if ua.len() > 256 { - ua.truncate(256); - } - - Ok(VersionMessage { - version: self.our_version, - services, - timestamp, - receiver: dashcore::network::address::Address::new(&address, ServiceFlags::NETWORK), - sender: dashcore::network::address::Address::new(&local_addr, services), - nonce: rand::random(), - user_agent: ua, - start_height: 0, // SPV client starts at 0 - relay: false, // relay enabled on demand via filterload/filterclear - mn_auth_challenge: [0; 32], // Not a masternode - masternode_connection: false, // Not connecting to masternode - }) - } - - /// Get current handshake state. - pub fn state(&self) -> &HandshakeState { - &self.state - } - - /// Get peer version if available. - pub fn peer_version(&self) -> Option { - self.peer_version - } - - /// Check if peer supports headers2 compression. - pub fn peer_supports_headers2(&self) -> bool { - self.peer_services.map(|services| services.has(NODE_HEADERS_COMPRESSED)).unwrap_or(false) - } - - /// Negotiate headers2 support with the peer after handshake completion. - async fn negotiate_headers2(&self, connection: &mut Peer) -> NetworkResult<()> { - if self.peer_supports_headers2() { - tracing::info!("Peer supports headers2 - sending SendHeaders2"); - connection.send_message(NetworkMessage::SendHeaders2).await?; - } else { - tracing::info!("Peer does not support headers2 - sending SendHeaders"); - connection.send_message(NetworkMessage::SendHeaders).await?; - } - Ok(()) - } -} diff --git a/dash-spv/src/network/manager.rs b/dash-spv/src/network/manager.rs index dbfccb443..9fd413c79 100644 --- a/dash-spv/src/network/manager.rs +++ b/dash-spv/src/network/manager.rs @@ -1,1558 +1,1986 @@ -//! Peer network manager for SPV client - -use std::collections::{HashMap, HashSet}; -use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4}; -use std::path::PathBuf; -use std::sync::atomic::{AtomicUsize, Ordering}; -use std::sync::Arc; -use std::time::Duration; -use tokio::sync::{broadcast, Mutex, RwLock}; -use tokio::task::JoinSet; -use tokio::time; - -use crate::client::ClientConfig; -use crate::error::{NetworkError, NetworkResult, SpvError as Error}; -use crate::network::addrv2::AddrV2Handler; -use crate::network::constants::*; -use crate::network::discovery::DnsDiscovery; -use crate::network::pool::PeerPool; -use crate::network::reputation::{ChangeReason, PeerReputationManager, ReputationAware}; -use crate::network::{ - HandshakeManager, Message, MessageDispatcher, MessageType, NetworkEvent, NetworkManager, - NetworkRequest, Peer, RequestSender, +use std::{ + collections::{HashMap, HashSet, VecDeque}, + net::SocketAddr, + sync::{ + atomic::{AtomicU32, AtomicU64, AtomicUsize, Ordering}, + Arc, + }, + time::{Duration, Instant}, }; -use crate::storage::{PeerStorage, PersistentPeerStorage, PersistentStorage}; -use async_trait::async_trait; -use dashcore::network::address::{AddrV2, AddrV2Message}; + use dashcore::network::constants::ServiceFlags; use dashcore::network::message::NetworkMessage; -use dashcore::network::message_headers2::CompressionState; -use dashcore::Network; -use tokio::sync::mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender}; -use tokio::time::Instant; +use dashcore::network::message_blockdata::Inventory; +use futures::future::join_all; +use tokio::sync::mpsc::{self, UnboundedReceiver, UnboundedSender}; +use tokio::sync::{broadcast, Mutex, Notify}; +use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; -const DEFAULT_NETWORK_EVENT_CAPACITY: usize = 10000; +/// Bounded concurrent handshakes per probe round. +const CONNECT_CHUNK: usize = 16; + +/// Candidates probed per improve round (at cap). Small to bound connect/close churn +/// while still discovering a better peer over time. +const IMPROVE_PROBE: usize = 4; + +/// Handshake ping below which a peer is "decent": preferred when filling, and the +/// bar a candidate must clear to be allowed to displace a slow connected peer. +const DECENT_LAG_MS: u32 = 100; + +/// Handshake ping at/above which a peer is "very bad": taken only as a last resort, +/// when nothing better is connectable and the set would otherwise be empty. +const BAD_LAG_MS: u32 = 1000; + +/// A candidate displaces the worst connected peer only if its ping is at most this +/// fraction of the worst peer's — i.e. clearly, not marginally, better. +const SWAP_IMPROVEMENT: u32 = 2; + +/// How often the supervisor re-checks a below-cap set and probes to keep filling. +const FILL_TICK: Duration = Duration::from_secs(2); + +/// How often the supervisor probes for a better peer once the set is at capacity. +const IMPROVE_TICK: Duration = Duration::from_secs(5); + +/// Cap on remembered (ranked) backup addresses, as a multiple of `max_peers`. +const BACKUP_MULTIPLE: usize = 8; + +/// How long the router sleeps on a full-capacity stall before re-evaluating. It +/// wakes early whenever a response frees a slot or the timeout monitor kicks a +/// peer; this is just a backstop so it never sleeps on a notify that never comes. +const STALL_CHECK: Duration = Duration::from_secs(5); + +/// A request unanswered for this long is treated as dead: the timeout monitor +/// re-queues it to another peer and kicks the peer sitting on it. +/// +/// Deliberately aggressive. A cfilter batch normally completes in well under a +/// second, so a peer holding a request for 10s is slow enough to be worth +/// dropping — the reconnector refills the slot with a fresh peer faster than a +/// laggard recovers, and we would rather churn a slow peer than let it drag its +/// in-flight slots. (The per-peer AIMED already throttles a merely-slow-but- +/// responding peer to the floor; this catches the ones that stop answering.) +const REQUEST_TIMEOUT: Duration = Duration::from_secs(10); + +/// How often the timeout monitor scans the outstanding-request registry. +const TIMEOUT_CHECK: Duration = Duration::from_secs(1); + +/// How long a retired peer (displaced during startup, see [`retire_drained`]) is +/// kept alive to drain its in-flight responses before being force-closed. +const RETIRE_DRAIN_CAP: Duration = Duration::from_secs(90); + +/// Poll interval for draining a retired peer's in-flight requests (see +/// [`retire_drained`]). The drain is capped at [`RETIRE_DRAIN_CAP`]. +const DRAIN_POLL: Duration = Duration::from_secs(1); +use crate::{ + network::{ + discovery::PeerDiscoverer, + peer::{ConnectedPeer, DisconnectedPeer, PeerEvent}, + }, + ClientConfig, +}; + +/// An inbound message on its way to the managers that subscribed to its type. +/// +/// Shared rather than cloned: a `block` or `cfilter` carries its whole payload, and the +/// pump would otherwise deep-copy it for every subscriber. See `spawn_pump`'s fan-out. +pub type Inbound = (SocketAddr, Arc); +type Subscribers = Arc>>>>; + +/// Every pipeline request the broker is handling, keyed by its identity, from the +/// moment `send` accepts it until the owning manager reports it answered +/// (`request_answered`) or cancels it (`cancel`). This is the single home of +/// request state: the pipelines hold no download coordinator, they just declare +/// what they want and the broker de-duplicates, paces, times out and retries. +/// +/// - Membership is the de-dup set: `send` ignores a key already present. +/// - `OnWire` entries carry the message so the timeout monitor can re-inject it +/// (retry) after dropping the peer that ignored it. +type Registry = Arc>>; + +/// The kinds of peer message a sync manager can subscribe to. Replaces the +/// stringly-typed command names: managers declare interest with these variants +/// and the pump routes incoming messages by mapping `cmd()` back to one. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub enum MessageType { + Headers, + Inv, + CfHeaders, + CFilter, + Block, + MnListDiff, + QrInfo, + Tx, + IsDLock, + ChainLock, +} + +impl MessageType { + /// The wire command string this type corresponds to. + pub fn cmd(self) -> &'static str { + match self { + MessageType::Headers => "headers", + MessageType::Inv => "inv", + MessageType::CfHeaders => "cfheaders", + MessageType::CFilter => "cfilter", + MessageType::Block => "block", + MessageType::MnListDiff => "mnlistdiff", + MessageType::QrInfo => "qrinfo", + MessageType::Tx => "tx", + MessageType::IsDLock => "isdlock", + MessageType::ChainLock => "clsig", + } + } + + /// Map an incoming message's command back to a subscribed type, if any. + pub fn from_cmd(cmd: &str) -> Option { + Some(match cmd { + "headers" => MessageType::Headers, + "inv" => MessageType::Inv, + "cfheaders" => MessageType::CfHeaders, + "cfilter" => MessageType::CFilter, + "block" => MessageType::Block, + "mnlistdiff" => MessageType::MnListDiff, + "qrinfo" => MessageType::QrInfo, + "tx" => MessageType::Tx, + "isdlock" => MessageType::IsDLock, + "clsig" => MessageType::ChainLock, + _ => return None, + }) + } +} + +#[derive(Clone, Debug)] +pub enum NetworkEvent { + /// The connected peer set changed. + PeersUpdated { + /// How many peers are currently connected. + connected_count: u32, + /// Best tip height advertised across those peers. + best_height: u32, + }, + PeerConnected(SocketAddr), + PeerDisconnected(SocketAddr), +} + +/// Identifies a pipeline request the network manager tracks from send to +/// response, so it can time the request out and re-queue it (and drop the peer +/// that ignored it). One variant per router-paced request type. Used as a map +/// key in the outstanding-request registry, hence `Hash`/`Eq`. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub enum RequestKey { + /// `getheaders` — keyed by the locator's first hash (the segment tip). + Headers(dashcore::BlockHash), + /// `getcfheaders` — keyed by the stop hash. + CfHeaders(dashcore::BlockHash), + /// `getcfilters` — keyed by the start height. + CFilters(u32), + /// `getmnlistdiff` — keyed by the target block hash. Tracked for de-dup and + /// retry only; unlike the others it is not counted against a peer's in-flight + /// budget (masternode diffs are few and not throughput-paced). + MnListDiff(dashcore::BlockHash), + /// Block `getdata` — keyed by the requested block hash. + Block(dashcore::BlockHash), +} + +/// Where a broker-tracked request is in its lifecycle. +enum ReqState { + /// Accepted by `send` and sitting in the `MsgQueue`, not yet on the wire. + /// Held here only for de-dup; the message itself lives in the queue. + Queued, + /// Sent to a peer and awaiting a response. Carries the message so the timeout + /// monitor can re-inject it (retry) after dropping the peer. Boxed: it holds a + /// whole `NetworkMessage`, far larger than the `Queued` variant. + OnWire(Box), +} + +/// A request currently on the wire, awaiting a response or a timeout. +struct OnWire { + /// The peer the request was routed to. + peer: SocketAddr, + /// When this request last made progress: set on send, refreshed by every piece of + /// a streaming response, so a peer streaming steadily is not judged stalled. + last_progress: Instant, + /// The exact message, re-queued verbatim on timeout (requests are + /// self-contained, so re-sending the same bytes is a valid retry). + msg: NetworkMessage, +} -/// Peer network manager pub struct PeerNetworkManager { - /// Peer pool - pool: Arc, + connected_peers: Arc>>, + other_peers: Arc>>, + discoverer: Arc>, + msg_queue: Arc, + inbound_tx: UnboundedSender, + subscribers: Subscribers, + /// Broker state for every request in play (queued or on the wire). + requests: Registry, + events_tx: broadcast::Sender, + /// Best tip advertised by the peers, learned in `start`. Shared with the + /// reconnector so its `PeersUpdated` carries the real height. + best_tip: Arc, max_peers: usize, - /// DNS discovery - discovery: Arc, - /// AddrV2 handler - addrv2_handler: Arc, - /// Peer persistence - peer_store: Arc, - /// Peer reputation manager - reputation_manager: Arc, - /// Network type - network: Network, - /// Shutdown token - shutdown_token: CancellationToken, - /// Background tasks - tasks: Arc>>, - /// Initial peer addresses - initial_peers: Vec, - /// Data directory for storage - data_dir: PathBuf, - /// Optional user agent to advertise - user_agent: Option, - /// Exclusive mode: restrict to configured peers only (no DNS or peer store) - exclusive_mode: bool, - /// Service flags connected peers must advertise. NONE disables capability churn. + /// Service flags a peer must advertise to be kept (e.g. COMPACT_FILTERS when + /// filters are enabled). Checked at handshake; `NONE` keeps every peer. required_services: ServiceFlags, - /// Addresses evicted for lacking required services. Excluded from top-up candidates. - /// TODO: remove once peer session outcomes track why sessions ended and drive reconnect policy. - capability_rejected: Arc>>, - /// Cached count of currently connected peers for fast, non-blocking queries - connected_peer_count: Arc, - /// Disable headers2 after decompression failure - headers2_disabled: Arc>>, - /// Dispatcher for unbounded and message-type filtered message distribution. - message_dispatcher: Arc>, - /// Request queue sender, cloneable handle for sending requests to the network manager. - request_tx: UnboundedSender, - /// Request queue receiver (consumed by send loop). - request_rx: Arc>>>, - /// Round-robin counter for distributing requests across peers. - round_robin_counter: Arc, - /// Network event bus for notifying about network/peer related changes. - network_event_sender: broadcast::Sender, + /// Total bytes read from all peers. Held so `start` can hand it to the peers it connects. + bytes: Arc, + // Cancelled by `stop()` to tear down the router, pump and every peer reader. + shutdown: CancellationToken, } -const CAPABILITY_REJECTED_TTL: Duration = Duration::from_secs(30 * 60); +/// Scheduling class of a queued message, in strict-priority order (see +/// [`MsgQueue::pop_n`]). Splitting by type lets the router drain them by priority +/// rather than FIFO, so a big backlog of one type never blocks another behind it. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +enum MsgClass { + /// Control traffic (mempool, tx, ping, chainlock/islock `getdata`, + /// `filterload`…): always first, few, and usually something is blocked on it. + Other, + /// Block `getdata` — the wallet's matched blocks. Highest bulk priority: a + /// matched block gates the gap-limit cascade and completes in a single reply, + /// so getting it out first keeps the scan advancing instead of stalling behind + /// the streaming filter backlog. + Blocks, + /// `getcfilters` — the bulk of the bytes. + CFilters, + /// `getcfheaders` — filter headers. + CfHeaders, + /// `getheaders` — block headers. + Headers, +} -fn required_services_from_config(config: &ClientConfig, exclusive_mode: bool) -> ServiceFlags { - if exclusive_mode { - return ServiceFlags::NONE; - } - let mut flags = ServiceFlags::NONE; - if config.enable_filters { - flags |= ServiceFlags::COMPACT_FILTERS; +fn classify(msg: &NetworkMessage) -> MsgClass { + match msg { + NetworkMessage::GetData(inv) + if !inv.is_empty() && inv.iter().all(|i| matches!(i, Inventory::Block(_))) => + { + MsgClass::Blocks + } + NetworkMessage::GetCFilters(_) => MsgClass::CFilters, + NetworkMessage::GetCFHeaders(_) => MsgClass::CfHeaders, + NetworkMessage::GetHeaders(_) | NetworkMessage::GetHeaders2(_) => MsgClass::Headers, + _ => MsgClass::Other, } - flags } -impl PeerNetworkManager { - /// Create a new peer network manager - pub async fn new(config: &ClientConfig) -> Result { - let discovery = DnsDiscovery::new(); - let data_dir = config.storage_path.clone(); +/// One queue per class, each behind its own lock: a pipeline enqueuing a burst +/// only contends with itself. The router drains them by strict priority, so no +/// class is ever starved behind another's backlog. +struct MsgQueue { + other: Mutex>, + blocks: Mutex>, + cfilters: Mutex>, + cfheaders: Mutex>, + headers: Mutex>, + len: AtomicUsize, + notify: Notify, +} - let peer_store = PersistentPeerStorage::open(data_dir.clone()).await?; +/// Strict drain priority: control traffic first, then blocks (they gate the scan +/// and complete fast), then filters (the bytes), then filter headers, then block +/// headers. +const DRAIN_PRIORITY: [MsgClass; 5] = + [MsgClass::Other, MsgClass::Blocks, MsgClass::CFilters, MsgClass::CfHeaders, MsgClass::Headers]; - let reputation_manager = Arc::new(PeerReputationManager::new()); +struct State {} - if let Err(e) = reputation_manager.load_from_storage(&peer_store).await { - tracing::warn!("Failed to load peer reputation data: {}", e); +impl PeerNetworkManager { + pub async fn new(config: &ClientConfig) -> Self { + let discoverer = Arc::new(Mutex::new(PeerDiscoverer::new(config))); + let max_peers = config.max_peers.max(1) as usize; + // NETWORK is unconditional: the block pipeline asks for full blocks at + // arbitrary historical heights (and the gap-limit rescan far more so), which + // a NETWORK_LIMITED peer stops serving past its last ~288 blocks. + let mut required_services = ServiceFlags::NETWORK; + // A filter-syncing client can only use peers that serve compact filters + // (BIP157 — the same flag also covers compact filter headers). + if config.enable_filters { + required_services |= ServiceFlags::COMPACT_FILTERS; + } + // Mempool tracking sends `mempool` to every activated peer, and `filterload` + // too under the BloomFilter strategy. A peer with bloom filters disabled + // answers the first by dropping us and the second, per BIP111, by banning us. + if config.enable_mempool_tracking { + required_services |= ServiceFlags::BLOOM; } - // Determine exclusive mode: either explicitly requested or peers were provided - let exclusive_mode = config.restrict_to_configured_peers || !config.peers.is_empty(); - let required_services = required_services_from_config(config, exclusive_mode); + let connected_peers = Arc::new(Mutex::new(Vec::with_capacity(30))); + let other_peers = Arc::new(Mutex::new(Vec::with_capacity(30))); + let msg_queue = Arc::new(MsgQueue::new()); + + let (inbound_tx, inbound_rx) = mpsc::unbounded_channel(); + let subscribers: Subscribers = Arc::new(Mutex::new(HashMap::new())); + let requests: Registry = Arc::new(Mutex::new(HashMap::new())); + // Sized generously: peer-connect churn plus one `RequestTimedOut` per + // dead request during a bad-peer storm. + let (events_tx, _) = broadcast::channel(4096); + let shutdown = CancellationToken::new(); + // Total bytes read from all peers (download-only) and the global in-flight + // budget. The budget is NOT a fixed number: it starts at a tiny bootstrap + // just large enough to begin measuring, then `spawn_bandwidth_controller` + // sizes it from the measured download capacity (Little's Law). + let bytes = Arc::new(AtomicU64::new(0)); + let global_cap = Arc::new(AtomicUsize::new(max_peers.saturating_mul(4).max(8))); + let best_tip = Arc::new(AtomicU32::new(0)); + + // Detached like the bandwidth controller and reconnector below: torn down + // via the shutdown token, not by holding their handles. + spawn_pump( + inbound_rx, + subscribers.clone(), + connected_peers.clone(), + events_tx.clone(), + msg_queue.clone(), + requests.clone(), + shutdown.clone(), + ); - // Create request queue for outgoing messages - let (request_tx, request_rx) = unbounded_channel(); + spawn_router( + msg_queue.clone(), + connected_peers.clone(), + shutdown.clone(), + global_cap.clone(), + requests.clone(), + ); - let max_peers = config.max_peers.max(1) as usize; + spawn_timeout_monitor( + requests.clone(), + connected_peers.clone(), + msg_queue.clone(), + shutdown.clone(), + ); - Ok(Self { - pool: Arc::new(PeerPool::new(max_peers)), + spawn_bandwidth_controller( + bytes.clone(), + global_cap.clone(), + connected_peers.clone(), + shutdown.clone(), + ); + + // The peer supervisor is spawned by `start()`, not here: it must not emit + // `PeersUpdated` until the sync managers have subscribed (see `start`). + + PeerNetworkManager { + connected_peers, + other_peers, + discoverer, + msg_queue, + inbound_tx, + subscribers, + requests, + events_tx, + best_tip, max_peers, - discovery: Arc::new(discovery), - addrv2_handler: Arc::new(AddrV2Handler::new()), - peer_store: Arc::new(peer_store), - reputation_manager, - network: config.network, - shutdown_token: CancellationToken::new(), - tasks: Arc::new(Mutex::new(JoinSet::new())), - initial_peers: config.peers.clone(), - data_dir, - user_agent: config.user_agent.clone(), - exclusive_mode, required_services, - capability_rejected: Arc::new(RwLock::new(HashMap::new())), - connected_peer_count: Arc::new(AtomicUsize::new(0)), - headers2_disabled: Arc::new(Mutex::new(HashSet::new())), - message_dispatcher: Arc::new(Mutex::new(MessageDispatcher::default())), - request_tx, - request_rx: Arc::new(Mutex::new(Some(request_rx))), - round_robin_counter: Arc::new(AtomicUsize::new(0)), - network_event_sender: broadcast::Sender::new(DEFAULT_NETWORK_EVENT_CAPACITY), - }) + bytes, + shutdown, + } } - /// Creates and returns a receiver that yields only messages of the matching the provided message types. - pub async fn message_receiver( - &mut self, - message_types: &[MessageType], - ) -> UnboundedReceiver { - self.message_dispatcher.lock().await.message_receiver(message_types) + /// Connect to peers and announce them. + /// + /// Split out of `new` on purpose: connecting there meant the one-shot `PeersUpdated` + /// (and every `PeerConnected`) fired before any sync manager had subscribed, so those + /// events were simply lost. Managers that track the peer set — the mempool, which must + /// send `filterload` to enable transaction relay — ended up with an empty set and never + /// activated. Build the manager, let the coordinator spawn and subscribe its managers, + /// then call this. + pub fn start(&self) { + // Non-blocking: spawn the supervisor and return. It probes peers, connects + // the decent ones, and emits `PeerConnected`/`PeersUpdated` as they arrive — + // so sync begins the moment the first decent peer is up, without `start` + // waiting on peer discovery. Spawned here rather than in `new` so it runs + // only after the coordinator has subscribed its managers; otherwise the first + // `PeersUpdated` (and the mempool's `filterload` trigger) would fire into the + // void. + spawn_peer_supervisor( + self.discoverer.clone(), + self.connected_peers.clone(), + self.other_peers.clone(), + self.inbound_tx.clone(), + self.shutdown.clone(), + self.bytes.clone(), + self.events_tx.clone(), + self.best_tip.clone(), + self.max_peers, + self.required_services, + ); } - /// Get a RequestSender for queueing outgoing network requests. - pub fn request_sender(&self) -> RequestSender { - RequestSender::new(self.request_tx.clone()) + /// Tear down the network layer: stop the router and pump, and cancel every + /// peer reader so no more messages arrive. Called on client shutdown. + pub fn stop(&self) { + tracing::info!(target: "dash_spv::network", "network manager stopping: cancelling tasks and peers"); + self.shutdown.cancel(); } - /// Get the network event bus for sharing with other components. - pub fn network_event_sender(&self) -> &broadcast::Sender { - &self.network_event_sender + /// Ask the broker to make a request. De-duplicated by request identity: if the + /// same request is already queued or on the wire, this is a no-op. Pipelines + /// exploit that to re-declare what they want each tick without tracking what + /// they already sent — the broker paces it, times it out and retries it. + /// + /// Non-pipeline messages (tx, mempool, control `getdata`) carry no request key + /// and are neither de-duplicated nor tracked; they just go on the queue. + pub async fn send(&self, msg: NetworkMessage) { + let keys = request_keys(&msg); + if keys.is_empty() { + self.msg_queue.push(msg).await; + return; + } + { + let mut reqs = self.requests.lock().await; + if keys.iter().any(|k| reqs.contains_key(k)) { + return; // already in play + } + for key in keys { + reqs.insert(key, ReqState::Queued); + } + } + self.msg_queue.push(msg).await; } - /// Start the network manager - pub async fn start(&self) -> Result<(), Error> { - tracing::info!("Starting peer network manager for {:?}", self.network); - - let mut peer_addresses: Vec = self - .initial_peers - .iter() - .map(|addr| AddrV2Message::new(*addr, ServiceFlags::NETWORK)) - .collect(); + /// Send a message to one specific peer, bypassing the router's round-robin. + /// + /// `send` hands a message to whichever peer has capacity, which is right for a request + /// any peer can answer. It is wrong for a message that sets state ON the remote node: + /// `filterload`/`filterclear` (and the `mempool` that follows) turn transaction relay + /// on for THAT peer, so routing them to "whoever is free" leaves the intended peer + /// silent — and, with several peers, can enable relay on the same one twice. + /// + /// Returns false if the peer is not connected (or the write failed). + pub async fn send_to(&self, addr: SocketAddr, msg: NetworkMessage) -> bool { + let peers = self.connected_peers.lock().await; + let Some((peer, _)) = peers.iter().find(|(p, _)| p.addr() == addr) else { + return false; + }; - if self.exclusive_mode { - tracing::info!( - "Exclusive peer mode: connecting ONLY to {} specified peer(s)", - self.initial_peers.len() - ); - } else { - // Load saved peers from disk - let saved_peers = self.peer_store.load_peers().await.unwrap_or_else(|e| { - tracing::warn!("Failed to load peers: {}", e); - Vec::new() - }); - peer_addresses.extend(saved_peers); - - // If we still have no peers, immediately discover via DNS - if peer_addresses.is_empty() { - tracing::info!( - "No peers configured, performing immediate DNS discovery for {:?}", - self.network - ); - let dns_peers = self.discovery.discover_peers(self.network).await; - let dns_peers_found = dns_peers.len(); - peer_addresses.extend( - dns_peers - .into_iter() - .take(self.max_peers) - .map(|addr| AddrV2Message::new(addr, ServiceFlags::NETWORK)), - ); - tracing::info!( - "DNS discovery found {} peers, using {} for startup", - dns_peers_found, - peer_addresses.len() - ); - } else { - tracing::info!( - "Starting with {} peers from disk (DNS discovery will be used later if needed)", - peer_addresses.len() - ); + match peer.send(&msg).await { + Ok(()) => true, + Err(e) => { + tracing::warn!(target: "dash_spv::network", "send to {addr} failed: {e}"); + false } } - - self.addrv2_handler.handle_addrv2(peer_addresses.clone()).await; - - // Start maintenance loop - self.start_maintenance_loop().await; - - // Start request processing task for managers to queue outgoing messages - self.start_request_processor().await; - - Ok(()) } - /// Connect to a specific peer - async fn connect_to_peer(&self, addr: SocketAddr) { - // Check reputation first - if !self.reputation_manager.should_connect_to_peer(&addr).await { - tracing::warn!("Not connecting to {} due to bad reputation", addr); + /// Note that `n` streaming requests served by `peer` have fully completed + /// (e.g. a `getcfilters` batch whose last `cfilter` just arrived), freeing + /// that peer's in-flight units and waking the router. Single-response + /// requests are freed in the peer's own reader instead. + pub async fn request_completed(&self, peer: SocketAddr, n: usize) { + if n == 0 { return; } - - // Check if already connected or connecting - if self.pool.is_connected(&addr).await || self.pool.is_connecting(&addr).await { - return; + if let Some((p, _)) = + self.connected_peers.lock().await.iter().find(|(p, _)| p.addr() == peer) + { + p.response_completed(n).await; } + self.msg_queue.notify.notify_one(); + } - // Mark as connecting - if !self.pool.mark_connecting(addr).await { - return; // Already being connected to - } + /// Report that a request's response arrived, so the broker stops tracking it + /// (no timeout, no retry, and its key is free to be requested again). The + /// owning manager calls this once it has correlated a response back to the + /// request key — the broker can't do it generically, since some responses + /// (e.g. an empty `headers`) carry nothing to match on. No-op if the key is + /// already gone (timed out first). + pub async fn request_answered(&self, key: RequestKey) { + self.requests.lock().await.remove(&key); + } - // Record connection attempt - self.reputation_manager.record_connection_attempt(addr).await; - - let pool = self.pool.clone(); - let network = self.network; - let addrv2_handler = self.addrv2_handler.clone(); - let shutdown_token = self.shutdown_token.clone(); - let reputation_manager = self.reputation_manager.clone(); - let user_agent = self.user_agent.clone(); - let required_services = self.required_services; - let capability_rejected = self.capability_rejected.clone(); - let connected_peer_count = self.connected_peer_count.clone(); - let headers2_disabled = self.headers2_disabled.clone(); - let message_dispatcher = self.message_dispatcher.clone(); - let network_event_sender = self.network_event_sender.clone(); - - // Spawn connection task — use select to avoid blocking on the lock during shutdown - let mut tasks = tokio::select! { - guard = self.tasks.lock() => guard, - _ = self.shutdown_token.cancelled() => { - self.pool.remove_peer(&addr).await; - return; + pub fn broadcast(&self, msg: NetworkMessage) { + let peers = self.connected_peers.clone(); + tokio::spawn(async move { + let guard = peers.lock().await; + for (peer, _) in guard.iter() { + let _ = peer.send(&msg).await; } - }; - tasks.spawn(async move { - tracing::debug!("Attempting to connect to {}", addr); - - let connect_result = tokio::select! { - result = Peer::connect(addr, CONNECTION_TIMEOUT.as_secs(), network) => result, - _ = shutdown_token.cancelled() => { - tracing::debug!("Connection to {} cancelled by shutdown", addr); - pool.remove_peer(&addr).await; - return; - } - }; - - match connect_result { - Ok(mut peer) => { - // Perform handshake - let mut handshake_manager = HandshakeManager::new(network, user_agent); - match handshake_manager.perform_handshake(&mut peer).await { - Ok(_) => { - if PeerNetworkManager::should_reject_after_handshake( - &pool, - &peer, - required_services, - ) - .await - { - tracing::info!( - "Rejecting peer {} during handshake - missing required services ({}) while a capable peer is connected", - addr, - required_services - ); - PeerNetworkManager::record_capability_rejection_in( - &capability_rejected, - addr, - ) - .await; - pool.remove_peer(&addr).await; - return; - } - tracing::info!("Successfully connected to {}", addr); + }); + } - // Request addresses from the peer for discovery - if let Err(e) = peer.send_message(NetworkMessage::GetAddr).await { - tracing::warn!("Failed to send GetAddr to {}: {}", addr, e); - } + /// Inject a message into the local pump as if it arrived from a peer, so + /// managers process it through the same path. Uses the `0.0.0.0:0` sentinel + /// address that managers treat as locally-originated. + pub async fn dispatch_local(&self, msg: NetworkMessage) { + let local: SocketAddr = ([0, 0, 0, 0], 0).into(); + let _ = self.inbound_tx.send(PeerEvent::Message(local, msg)); + } - // Record successful connection - reputation_manager.record_successful_connection(addr).await; + pub async fn subscribe(&self, kinds: &[MessageType]) -> UnboundedReceiver { + let (tx, rx) = mpsc::unbounded_channel(); + let mut subscribers = self.subscribers.lock().await; + for kind in kinds.iter().copied().collect::>() { + subscribers.entry(kind).or_default().push(tx.clone()); + } + rx + } - // Add to pool - if let Err(e) = pool.add_peer(addr, peer).await { - tracing::error!("Failed to add peer to pool: {}", e); - return; - } + pub fn tip(&self) -> u32 { + self.best_tip.load(Ordering::Relaxed) + } - // Increment connected peer counter on successful add - connected_peer_count.fetch_add(1, Ordering::Relaxed); + /// How many peers are currently connected. + pub async fn connected_count(&self) -> u32 { + self.connected_peers.lock().await.len() as u32 + } - // Emit peer connected event - let count = connected_peer_count.load(Ordering::Relaxed); - let addresses = pool.get_connected_addresses().await; - let best_height = pool.get_best_height().await; - let _ = network_event_sender.send(NetworkEvent::PeerConnected { - address: addr, - }); - let _ = network_event_sender.send(NetworkEvent::PeersUpdated { - connected_count: count, - addresses, - best_height, - }); + pub fn events(&self) -> broadcast::Receiver { + self.events_tx.subscribe() + } +} - // Add to known addresses - addrv2_handler.add_known_address(addr, ServiceFlags::NETWORK).await; - - // // Start message reader for this peer - Self::start_peer_reader( - addr, - pool.clone(), - addrv2_handler, - shutdown_token, - reputation_manager.clone(), - connected_peer_count.clone(), - headers2_disabled.clone(), - message_dispatcher, - network_event_sender.clone(), - ) - .await; - } - Err(e) => { - tracing::warn!("Handshake failed with {}: {}", addr, e); - // Only clears connecting set. Peer was never added, so no count/event needed. - pool.remove_peer(&addr).await; - // Update reputation for handshake failure - reputation_manager - .update_reputation(addr, ChangeReason::HandshakeFailed) - .await; - // For handshake failures, try again later - tokio::time::sleep(RECONNECT_DELAY).await; - } - } +fn spawn_router( + queue: Arc, + connected: Arc>>, + shutdown: CancellationToken, + global_cap: Arc, + requests: Registry, +) -> JoinHandle<()> { + tokio::spawn(async move { + loop { + if shutdown.is_cancelled() { + break; + } + // Wait for work. `notify` fires both when a message is queued and + // when a response frees a peer slot. + if queue.len() == 0 { + tokio::select! { + _ = shutdown.cancelled() => break, + _ = queue.notify.notified() => continue, } - Err(e) => { - tracing::debug!("Failed to connect to {}: {}", addr, e); - // Only clears connecting set. Peer was never added, so no count/event needed. - pool.remove_peer(&addr).await; - // Minor reputation penalty for connection failure - reputation_manager - .update_reputation(addr, ChangeReason::ConnectionFailed) - .await; + } + + let peers = connected.lock().await; + let sent = + route_tick(&queue, &peers, global_cap.load(Ordering::Relaxed), &requests).await; + drop(peers); + + if sent == 0 { + // Queue non-empty but every peer is at its in-flight cap: wait for + // a response to free a slot (the pump notifies on each response) or + // for the timeout monitor to kick a stalled peer (which notifies + // too). The sleep is only a backstop against a missed wake — dead + // in-flight slots are reclaimed by the monitor dropping the peer + // that holds them, not here. + tokio::select! { + _ = shutdown.cancelled() => break, + _ = queue.notify.notified() => {}, + _ = tokio::time::sleep(STALL_CHECK) => {}, } } - }); - } + } + }) +} - /// Decrement the connected count and emit PeerDisconnected / PeersUpdated events. - async fn notify_peer_removed( - pool: &PeerPool, - addr: &SocketAddr, - connected_peer_count: &AtomicUsize, - network_event_sender: &broadcast::Sender, - ) { - let sub_result = - connected_peer_count - .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |c| c.checked_sub(1)); - if sub_result.is_err() { - tracing::warn!("Peer count already zero when removing {}", addr); +/// Extract the pipeline keys of a router-paced request, so the router can record +/// it in the outstanding-request registry. Returns empty for non-pipeline +/// messages (mirrors `is_pipeline_request` in `peer`: only these count toward a +/// peer's in-flight and are timed out). +fn request_keys(msg: &NetworkMessage) -> Vec { + match msg { + NetworkMessage::GetHeaders(m) | NetworkMessage::GetHeaders2(m) => { + m.locator_hashes.first().map(|h| RequestKey::Headers(*h)).into_iter().collect() } - let count = connected_peer_count.load(Ordering::Relaxed); - let addresses = pool.get_connected_addresses().await; - let best_height = pool.get_best_height().await; - let _ = network_event_sender.send(NetworkEvent::PeerDisconnected { - address: *addr, - }); - let _ = network_event_sender.send(NetworkEvent::PeersUpdated { - connected_count: count, - addresses, - best_height, - }); + NetworkMessage::GetCFHeaders(m) => vec![RequestKey::CfHeaders(m.stop_hash)], + NetworkMessage::GetCFilters(m) => vec![RequestKey::CFilters(m.start_height)], + NetworkMessage::GetMnListD(m) => vec![RequestKey::MnListDiff(m.block_hash)], + // One `getdata` may name several blocks; each is its own tracked request. + NetworkMessage::GetData(inv) => inv + .iter() + .filter_map(|i| match i { + Inventory::Block(h) => Some(RequestKey::Block(*h)), + _ => None, + }) + .collect(), + _ => Vec::new(), + } +} + +async fn route_tick( + queue: &MsgQueue, + peers: &[(ConnectedPeer, State)], + global_cap: usize, + requests: &Registry, +) -> usize { + if peers.is_empty() { + return 0; } - /// Remove a peer from the pool, decrement the connected count, and emit - /// PeerDisconnected / PeersUpdated events. - async fn remove_peer_and_notify( - pool: &PeerPool, - addr: &SocketAddr, - connected_peer_count: &AtomicUsize, - network_event_sender: &broadcast::Sender, - ) { - if pool.remove_peer(addr).await.is_some() { - Self::notify_peer_removed(pool, addr, connected_peer_count, network_event_sender).await; + // Free capacity this round = min(sum of per-peer room, global room). Each + // peer's cap is its MEASURED serving capacity (its bandwidth-delay product), + // sized by the controller from that peer's own completion rate and service + // time — fast peers carry more, slow peers less, with no fixed constant. The + // global cap is our measured download capacity. Whichever binds first limits + // this round, so we ride each peer's real ceiling without over-committing. + let total_in_flight: usize = peers.iter().map(|(p, _)| p.in_flight()).sum(); + let per_peer_room: usize = + peers.iter().map(|(p, _)| p.cap().saturating_sub(p.in_flight())).sum(); + let global_room = global_cap.saturating_sub(total_in_flight); + let capacity = per_peer_room.min(global_room); + if capacity == 0 { + return 0; + } + + let msgs = queue.pop_n(capacity).await; + let mut sent = 0; + // Anything popped that we could not put on the wire goes BACK on the queue. + // Dropping it would strand the owning pipeline forever: it has already marked + // the request as handed to the network, and its response timeout only starts + // when the router reports the request on the wire, so a dropped message is + // never re-sent and never times out. + let mut unsent: Vec = Vec::new(); + // Messages that made it onto the wire this round, recorded in the broker in one + // lock acquisition after the send loop. + let mut on_wire: Vec<(NetworkMessage, SocketAddr)> = Vec::new(); + let mut msgs = msgs.into_iter(); + for msg in msgs.by_ref() { + // Send to the peer with the most free measured capacity. + let Some((peer, _)) = peers + .iter() + .filter(|(p, _)| p.in_flight() < p.cap()) + .max_by_key(|(p, _)| p.cap().saturating_sub(p.in_flight())) + else { + unsent.push(msg); // every peer is at its measured cap + break; + }; + if peer.send(&msg).await.is_ok() { + sent += 1; + // Record which peer got it, so the monitor can attribute a stall to it. + on_wire.push((msg, peer.addr())); + } else { + tracing::warn!(target: "dash_spv::network", "router: send to {} failed", peer.addr()); + unsent.push(msg); + } + } + unsent.extend(msgs); // whatever the loop never reached + queue.push_front_all(unsent).await; + + if !on_wire.is_empty() { + let mut reqs = requests.lock().await; + for (msg, peer) in on_wire { + for key in request_keys(&msg) { + // Transition Queued -> OnWire, keeping the message for retry. Skip + // keys no longer present (cancelled while queued): the request went + // out but we don't track it, so its response is simply ignored. + if let Some(slot) = reqs.get_mut(&key) { + *slot = ReqState::OnWire(Box::new(OnWire { + peer, + msg: msg.clone(), + last_progress: Instant::now(), + })); + } + } } } - /// Start reading messages from a peer - #[allow(clippy::too_many_arguments)] // TODO: refactor to reduce arguments - async fn start_peer_reader( - addr: SocketAddr, - pool: Arc, - addrv2_handler: Arc, - shutdown_token: CancellationToken, - reputation_manager: Arc, - connected_peer_count: Arc, - headers2_disabled: Arc>>, - message_dispatcher: Arc>, - network_event_sender: broadcast::Sender, - ) { - tokio::spawn(async move { - tracing::debug!("Starting peer reader loop for {}", addr); - let mut loop_iteration = 0; - let mut headers2_state = CompressionState::default(); + if sent > 0 { + tracing::debug!( + target: "dash_spv::network", + "router: sent {} | peers={} queue={}", + sent, + peers.len(), + queue.len(), + ); + } + sent +} - loop { - loop_iteration += 1; +/// Per-peer state the bandwidth controller carries across windows to size each +/// connection's in-flight cap independently, by Little's Law over that peer's OWN +/// completion stream (rather than an even split of the global budget). +#[derive(Default, Clone, Copy)] +struct PeerCapState { + /// Cumulative completions/service-ns at the last window, to diff against. + last_count: u64, + last_total_ns: u64, + /// Cumulative bytes downloaded from this peer at the last window, for its + /// per-window throughput. + last_bytes: u64, + /// Uncongested service-time baseline in seconds (the peer's min `W`), 0 until + /// first measured. Little's Law targets `L = λ · min_W`. + min_w: f64, + /// Windows since `min_w` last took a new low (BBR-style min-filter age). Lets a + /// stale baseline expire and re-track the current cost, instead of one cheap + /// early sample pinning it forever. + min_w_age: u32, + /// Smoothed cap, so it doesn't jitter window to window. + cap_ema: f64, +} - // Check shutdown signal first with detailed logging - if shutdown_token.is_cancelled() { - tracing::info!("Breaking peer reader loop for {} - shutdown signal received (iteration {})", addr, loop_iteration); - break; - } +/// Sizes the host's GLOBAL in-flight budget from MEASURED download throughput and +/// each PEER's cap from its own completion stream — no fixed magic number, per the +/// design goal of estimating how many requests the current network can absorb +/// before it saturates. +/// +/// Two levels, both measured: +/// - The GLOBAL budget (host downlink) hill-climbs on bytes/s read off the sockets +/// (see below). It is the ceiling the host can reach with all peers combined. +/// - Each PEER's cap is `L = λ · min_W` (Little's Law) from THAT peer's completion +/// rate `λ` and uncongested service time `min_W`, so a fast peer earns a high cap +/// and a slow one a low cap. `route_tick` binds `min(Σ per-peer room, global +/// room)`, so whichever is the real bottleneck — the peers or the host link — +/// limits each round. +/// +/// The estimate is Little's Law applied to downloads: the number of requests in +/// flight that sustains a completion rate `λ` at an uncongested per-request +/// service time `W` is `L = λ · W`. We measure both from the peers' response +/// stream — `λ` = requests completed per second, `W` = average time from send to +/// the response that completes the request (for `getcfilters`, dominated by the +/// download time of its ~1000 `cfilter`s, i.e. our downlink). We track the +/// minimum `W` as the uncongested baseline (the pipe's true latency at the front +/// of the knee) and target `L = λ · min_W`, probing slightly past it. +/// +/// Saturation is detected by `W` INFLATION, not by throughput: once in-flight +/// exceeds the bandwidth-delay product, extra requests just queue at the peers, +/// so `W` climbs while `λ` (and the download rate) plateaus. That inflation is +/// visible even though a raw throughput meter can't see the knee (a prior +/// throughput hill-climb ran away and regressed sync ~8x because the backlog +/// kept bytes flowing). When `W > min_W · INFLATE` we stop probing and shrink +/// back toward the sustaining level, so we ride just below saturation. +fn spawn_bandwidth_controller( + bytes: Arc, + cap: Arc, + connected: Arc>>, + shutdown: CancellationToken, +) -> JoinHandle<()> { + const WINDOW: Duration = Duration::from_millis(500); + const FLOOR_PER_PEER: usize = 2; // global floor = peers · this + const PEER_CEIL: usize = 32; // per-connection sanity bound on in-flight + const RISE: f64 = 1.05; // throughput must climb 5% to justify a bigger cap + const DROP: f64 = 0.85; // throughput below this·last => over-commit, back off + const REPROBE: u32 = 8; // plateau windows to hold before nudging the cap up + const IDLE_BPS: f64 = 1.0e6; // downlink under 1 MB/s = idle, hold the cap + const EMA_ALPHA: f64 = 0.5; // smoothing for the noisy per-window rate + // Per-peer cap: AIMED driven purely by THIS peer's service-time (lag). There is + // NO hard per-peer request limit — a peer with headroom keeps growing, so we + // fill the peers we have instead of recruiting more. It only backs off when its + // own lag inflates past the uncongested baseline. + const CAP_GROW: f64 = 1.0; // additive increase per window while lag is flat + const CAP_BACKOFF: f64 = 0.8; // multiplicative decrease when lag inflates + const W_INFLATE: f64 = 1.5; // W above min_W·this => this peer is backing up + const MIN_W_WINDOW: u32 = 20; // windows before a stale min_W baseline is re-tracked + let window_s = WINDOW.as_secs_f64(); + + tokio::spawn(async move { + let mut last_bytes = bytes.load(Ordering::Relaxed); + let mut rate_ema = 0.0f64; // smoothed downlink bytes/s + let mut last_rate = 0.0f64; // smoothed rate at the previous cap adjustment + let mut hold = 0u32; // consecutive plateau windows + // Per-peer cap state across windows, keyed by peer address. + let mut peer_caps: HashMap = HashMap::new(); + let mut ticker = tokio::time::interval(WINDOW); + loop { + tokio::select! { + _ = shutdown.cancelled() => break, + _ = ticker.tick() => {} + } - // Get peer - let peer = match pool.get_peer(&addr).await { - Some(peer) => peer, - None => { - tracing::warn!("Breaking peer reader loop for {} - peer no longer in pool (iteration {})", addr, loop_iteration); - break; - } - }; - - // Read message with minimal lock time - let msg_result = { - // Try to get a read lock first to check if peer is available - let peer_guard = peer.read().await; - if !peer_guard.is_connected() { - tracing::warn!("Breaking peer reader loop for {} - peer no longer connected (iteration {})", addr, loop_iteration); - drop(peer_guard); - break; - } - drop(peer_guard); - - // Now get write lock only for the duration of the read - let mut peer_guard = peer.write().await; - tokio::select! { - message = peer_guard.receive_message() => { - message - }, - _ = tokio::time::sleep(MESSAGE_POLL_INTERVAL) => { - Ok(None) - }, - _ = shutdown_token.cancelled() => { - tracing::info!("Breaking peer reader loop for {} - shutdown signal received while reading (iteration {})", addr, loop_iteration); - break; - } - } - }; - - match msg_result { - Ok(Some(msg)) => { - // Log all received messages at debug level to help troubleshoot - tracing::trace!("Received {:?} from {}", msg.cmd(), addr); - - // Handle some messages directly - match &msg.inner() { - NetworkMessage::SendAddrV2 => { - addrv2_handler.handle_sendaddrv2(addr).await; - continue; // Don't forward to client - } - NetworkMessage::SendHeaders2 => { - // Peer is indicating they will send us compressed headers - tracing::info!( - "Peer {} sent SendHeaders2 - they will send compressed headers", - addr - ); - let mut peer_guard = peer.write().await; - peer_guard.set_peer_sent_sendheaders2(true); - drop(peer_guard); - continue; // Don't forward to client - } - NetworkMessage::AddrV2(addresses) => { - addrv2_handler.handle_addrv2(addresses.clone()).await; - continue; // Don't forward to client - } - NetworkMessage::GetAddr => { - tracing::trace!( - "Received GetAddr from {}, sending known addresses", - addr - ); - // Send our known addresses - let response = addrv2_handler.build_addr_response().await; - let mut peer_guard = peer.write().await; - if let Err(e) = peer_guard.send_message(response).await { - tracing::error!( - "Failed to send addr response to {}: {}", - addr, - e - ); - } - continue; // Don't forward GetAddr to client - } - NetworkMessage::Ping(nonce) => { - // Handle ping directly - let mut peer_guard = peer.write().await; - if let Err(e) = peer_guard.handle_ping(*nonce).await { - tracing::error!("Failed to handle ping from {}: {}", addr, e); - // If we can't send pong, connection is likely broken - if matches!(e, NetworkError::ConnectionFailed(_)) { - tracing::warn!("Breaking peer reader loop for {} - failed to send pong response (iteration {})", addr, loop_iteration); - break; - } - } - continue; // Don't forward ping to client - } - NetworkMessage::Pong(nonce) => { - // Handle pong directly - let mut peer_guard = peer.write().await; - if let Err(e) = peer_guard.handle_pong(*nonce) { - tracing::error!("Failed to handle pong from {}: {}", addr, e); - } - continue; // Don't forward pong to client - } - NetworkMessage::Version(_) | NetworkMessage::Verack => { - // These are handled during handshake, ignore here - tracing::trace!( - "Ignoring handshake message {:?} from {}", - msg.cmd(), - addr - ); - continue; - } - NetworkMessage::Addr(addresses) => { - // Convert legacy addr messages to AddrV2 format - let converted: Vec = addresses - .iter() - .filter_map(|(time, a)| { - let socket = a.socket_addr().ok()?; - let addr_v2 = match socket.ip() { - std::net::IpAddr::V4(v4) => AddrV2::Ipv4(v4), - std::net::IpAddr::V6(v6) => AddrV2::Ipv6(v6), - }; - Some(AddrV2Message { - time: *time, - services: a.services, - addr: addr_v2, - port: socket.port(), - }) - }) - .collect(); - if !converted.is_empty() { - tracing::debug!( - "Converted {} legacy addr entries from {}", - converted.len(), - addr - ); - addrv2_handler.handle_addrv2(converted).await; - } - continue; - } - NetworkMessage::Headers(headers) => { - // Log headers messages specifically - tracing::info!( - "📨 Received Headers message from {} with {} headers! (regular uncompressed)", - addr, - headers.len() - ); - // Check if peer supports headers2 - let peer_guard = peer.read().await; - if peer_guard.supports_headers2() { - tracing::warn!("⚠️ Peer {} supports headers2 but sent regular headers - possible protocol issue", addr); - } - drop(peer_guard); - // Forward to client - } - NetworkMessage::Headers2(headers2) => { - // Decompress headers in network layer and forward as regular Headers - tracing::info!( - "Received Headers2 from {} with {} compressed headers - decompressing", - addr, - headers2.headers.len() - ); - - match headers2_state.process_headers(&headers2.headers) { - Ok(headers) => { - tracing::info!( - "Decompressed {} headers from {} - forwarding as regular Headers", - headers.len(), - addr - ); - // Forward as regular Headers message - let headers_msg = NetworkMessage::Headers(headers); - let message = Message::new(msg.peer_address(), headers_msg); - message_dispatcher.lock().await.dispatch(&message); - continue; // Already sent, don't forward the original Headers2 - } - Err(e) => { - tracing::error!( - "Headers2 decompression failed from {}: {} - disabling headers2", - addr, - e - ); - headers2_disabled.lock().await.insert(addr); - // Apply reputation penalty - reputation_manager - .update_reputation( - addr, - ChangeReason::Headers2DecompressionFailed, - ) - .await; - continue; // Don't forward corrupted message - } - } - } - NetworkMessage::GetHeaders(_) => { - // SPV clients don't serve headers to peers - tracing::debug!( - "Received GetHeaders from {} - ignoring (SPV client)", - addr - ); - continue; // Don't forward to client - } - NetworkMessage::GetHeaders2(_) => { - // SPV clients don't serve compressed headers to peers - tracing::debug!( - "Received GetHeaders2 from {} - ignoring (SPV client)", - addr - ); - continue; // Don't forward to client - } - NetworkMessage::Unknown { - command, - payload, - } => { - // Log unknown messages with more detail - tracing::warn!("Received unknown message from {}: command='{}', payload_len={}", - addr, command, payload.len()); - // Still forward to client - } - _ => { - // Forward other messages to client - tracing::trace!( - "Forwarding {:?} from {} to client", - msg.cmd(), - addr - ); + // Downlink throughput (bytes read off the sockets — our downlink only). + // This is the saturation signal: unlike per-request service time — which + // balloons with cfilter payload size and out-of-order batch completion, + // reading "saturated" forever and pinning the cap at its floor — bytes/s + // directly reflects whether we are using the pipe. We grow the in-flight + // budget while throughput keeps climbing with it and stop at the plateau + // (the bandwidth-delay product: past it, more in-flight only grows queues, + // not bytes/s), backing off if it collapses (peers over-committed). + let now_bytes = bytes.load(Ordering::Relaxed); + let dl_rate = now_bytes.saturating_sub(last_bytes) as f64 / window_s; + last_bytes = now_bytes; + rate_ema = if rate_ema == 0.0 { + dl_rate + } else { + EMA_ALPHA * dl_rate + (1.0 - EMA_ALPHA) * rate_ema + }; + + let npeers = connected.lock().await.len(); + if npeers == 0 { + continue; + } + let floor = (npeers * FLOOR_PER_PEER).max(FLOOR_PER_PEER); + let ceiling = (npeers * PEER_CEIL).max(floor + 1); + let step = npeers.max(4); // ~one extra slot per peer per window + let cur = cap.load(Ordering::Relaxed); + + // Gradient hill-climb on smoothed throughput. + let (new, action) = if rate_ema < IDLE_BPS { + // Nothing meaningful downloading (e.g. the commit tail): hold the + // budget steady so it is ready when the download resumes. + (cur, "idle") + } else if cur <= floor || rate_ema >= last_rate * RISE { + // Still gaining (or at the floor): push the budget up. + hold = 0; + ((cur + step).min(ceiling), "grow") + } else if rate_ema < last_rate * DROP { + // Throughput collapsed — the peers are over-committed. Back off. + hold = 0; + (((cur as f64 * 0.8) as usize).max(floor), "backoff") + } else { + // Plateau: we are at the knee. Hold, re-probing up occasionally to + // catch a capacity increase (a faster peer, less congestion). + hold += 1; + if hold >= REPROBE { + hold = 0; + ((cur + step).min(ceiling), "reprobe") + } else { + (cur, "hold") + } + }; + // Anchor RISE/DROP to the rate at each real adjustment (skip idle/hold + // windows) so the next comparison is like-for-like. + if matches!(action, "grow" | "backoff" | "reprobe") { + last_rate = rate_ema; + } + cap.store(new, Ordering::Relaxed); + + // Size each peer's cap by AIMED on its OWN service time (lag): grow while + // its lag stays at the uncongested baseline (it has headroom), back off + // the moment its lag inflates (it is backing up). No hard per-peer + // request limit — a fast peer keeps growing so we fill the peers we have + // instead of forcing new connections while there is still room for work. + // The global cap above stays the host ceiling (route_tick binds by it). + let (mut cap_min, mut cap_max, mut cap_sum) = (usize::MAX, 0usize, 0usize); + let mut inflight_sum = 0usize; // total requests on the wire right now + { + let g = connected.lock().await; + let live: HashSet = g.iter().map(|(p, _)| p.addr()).collect(); + for (p, _) in g.iter() { + let addr = p.addr(); + let (count, total_ns) = p.latency_totals(); + let now_bytes = p.bytes_read(); + let st = peer_caps.entry(addr).or_default(); + let dc = count.saturating_sub(st.last_count); + let dt = total_ns.saturating_sub(st.last_total_ns); + let d_bytes = now_bytes.saturating_sub(st.last_bytes); + st.last_count = count; + st.last_total_ns = total_ns; + st.last_bytes = now_bytes; + let rate = d_bytes as f64 / window_s; // this peer's downlink (bytes/s) + + let (lambda, w) = if dc == 0 { + // No completions this window: either idle (no work queued to + // it) or stalled — the timeout monitor kicks a stalled peer at + // REQUEST_TIMEOUT. Keep at least the floor so the router can + // hand it work to bootstrap/keep measuring, but don't grow blind. + st.cap_ema = st.cap_ema.max(FLOOR_PER_PEER as f64); + (0.0, 0.0) + } else { + let lambda = dc as f64 / window_s; // completions/sec + let w = (dt as f64 / dc as f64) / 1e9; // avg service time (s) + // Windowed min-W baseline (BBR-style min filter): take a new low + // immediately, otherwise let the baseline go stale and re-track + // the current cost after MIN_W_WINDOW windows. Without the reset, + // one low sample from a cheap phase (fast headers) pins the + // baseline forever and every later heavier request (cfilters) + // reads as inflated => the cap decays to the floor and never + // recovers, throttling the very phase we want parallel. + if st.min_w == 0.0 || w < st.min_w { + st.min_w = w; + st.min_w_age = 0; + } else { + st.min_w_age += 1; + if st.min_w_age >= MIN_W_WINDOW { + st.min_w = w; + st.min_w_age = 0; } } + // AIMED on this peer's own lag: additive-increase while its + // service time sits at the uncongested baseline (headroom), + // multiplicative-decrease the moment it inflates (backing up). + if st.cap_ema == 0.0 { + st.cap_ema = FLOOR_PER_PEER as f64; + } else if w > st.min_w * W_INFLATE { + st.cap_ema = (st.cap_ema * CAP_BACKOFF).max(FLOOR_PER_PEER as f64); + } else { + st.cap_ema = (st.cap_ema + CAP_GROW).min(PEER_CEIL as f64); + } + (lambda, w) + }; - message_dispatcher.lock().await.dispatch(&msg); - } - Ok(None) => { - // No message available, continue immediately - // The socket read timeout already provides necessary delay - continue; - } - Err(e) => { - match e { - NetworkError::PeerDisconnected => { - tracing::info!("Peer {} disconnected", addr); - break; - } - NetworkError::Timeout => { - tracing::debug!("Timeout reading from {}, continuing...", addr); - // Minor reputation penalty for timeout - reputation_manager - .update_reputation(addr, ChangeReason::ReadTimeout) - .await; - continue; - } - _ => { - tracing::error!("Fatal error reading from {}: {}", addr, e); - - // Check if this is a serialization error that might have context - if let NetworkError::Serialization(ref decode_error) = e { - let error_msg = decode_error.to_string(); - if error_msg.contains("unknown special transaction type") { - tracing::warn!("Peer {} sent block with unsupported transaction type: {}", addr, decode_error); - tracing::error!( - "BLOCK DECODE FAILURE - Error details: {}", - error_msg - ); - // Reputation penalty for invalid data - reputation_manager - .update_reputation( - addr, - ChangeReason::InvalidTransactionInBlock, - ) - .await; - } else if error_msg - .contains("Failed to decode transactions for block") - { - // The error now includes the block hash - tracing::error!("Peer {} sent block that failed transaction decoding: {}", addr, decode_error); - // Try to extract the block hash from the error message - if let Some(hash_start) = error_msg.find("block ") { - if let Some(hash_end) = - error_msg[hash_start + 6..].find(':') - { - let block_hash = &error_msg - [hash_start + 6..hash_start + 6 + hash_end]; - tracing::error!( - "FAILING BLOCK HASH: {}", - block_hash - ); - } - } - } else if error_msg.contains("IO error") { - // This might be our wrapped error - log it prominently - tracing::error!("BLOCK DECODE FAILURE - IO error (possibly unknown transaction type) from peer {}", addr); - tracing::error!( - "Serialization error from {}: {}", - addr, - decode_error - ); - } else { - tracing::error!( - "Serialization error from {}: {}", - addr, - decode_error - ); - } - } + let cap_peer = (st.cap_ema.round() as usize).clamp(FLOOR_PER_PEER, PEER_CEIL); + p.set_cap(cap_peer); - break; - } - } - } + tracing::debug!( + target: "peer_speed", + "peer {}: cap={} in_flight={} lag={}ms lambda={:.1}/s W={:.0}ms rate={:.2} MB/s total={:.1} MB", + addr, + cap_peer, + p.in_flight(), + p.lag_ms(), + lambda, + w * 1e3, + rate / 1e6, + p.bytes_read() as f64 / 1e6, + ); + + cap_min = cap_min.min(cap_peer); + cap_max = cap_max.max(cap_peer); + cap_sum += cap_peer; + inflight_sum += p.in_flight(); } + // Drop state for peers that have disconnected. + peer_caps.retain(|addr, _| live.contains(addr)); + } + if cap_min == usize::MAX { + cap_min = 0; } - // Remove from pool and notify consumers - tracing::warn!("Disconnecting from {} (peer reader loop ended)", addr); - Self::remove_peer_and_notify( - &pool, - &addr, - &connected_peer_count, - &network_event_sender, - ) - .await; + // `bind` names the constraint the router is hitting: `global` if the + // host budget is the smaller room, `peers` if the summed per-peer caps + // are, `work` if neither is full (queue-limited or peers just slow). If + // the downlink is saturated the global cap should hold at the knee and + // `bind=global`. + let bind = if inflight_sum >= new.min(cap_sum) { + if new <= cap_sum { + "global" + } else { + "peers" + } + } else { + "work" + }; + tracing::debug!( + target: "peer_speed", + "bandwidth: {:.1} MB/s (ema {:.1}) | {} bind={} | global_cap={} sum_peer_cap={} in_flight={} | peers={} per-peer cap min/avg/max={}/{}/{}", + dl_rate / 1e6, + rate_ema / 1e6, + action, + bind, + new, + cap_sum, + inflight_sum, + npeers, + cap_min, + cap_sum / npeers.max(1), + cap_max, + ); + } + }) +} - headers2_disabled.lock().await.remove(&addr); +/// Keep the peer set topped up. +/// +/// Peers are connected once, in `start`; nothing put them back afterwards, so a client +/// whose peers all dropped — while idle or mid-sync — would simply sit there with zero +/// peers forever. This watches the count and refills it back to `max_peers`, pulling +/// fresh candidates from the discoverer when the backup list runs dry. +#[allow(clippy::too_many_arguments)] +/// Sort key for a connected peer's handshake ping: lower is better, and an +/// unmeasured lag (0) sorts as worst. +fn lag_key(peer: &ConnectedPeer) -> u32 { + match peer.lag_ms() { + 0 => u32::MAX, + ms => ms, + } +} + +/// The peer supervisor: the single task that owns the connected-peer set. +/// +/// It keeps `connected` filled toward `max_peers` with the lowest-latency peers it +/// can find, without ever blocking the caller: +/// +/// - Below cap it probes candidates in parallel and accepts the decent ones +/// (handshake ping under [`BAD_LAG_MS`]), emitting `PeersUpdated` as they connect +/// so sync starts on the first one. A "very bad" peer is taken only as a last +/// resort — when the set would otherwise be empty and nothing better connected. +/// - At cap it probes a few candidates every [`IMPROVE_TICK`] and, if one is clearly +/// better (ping ≤ worst / [`SWAP_IMPROVEMENT`]) than the slowest connected peer, +/// swaps it in. The displaced peer is handed to [`retire_drained`] so its in-flight +/// requests finish (or time out) before its socket closes. +/// - A peer kicked by the timeout monitor just drops the set below cap, so the next +/// fill round refills it — the same path as any other deficit. +/// +/// Probed-but-unused peers are closed and kept as ranked backups (`others`, carrying +/// their measured ping) so a later round reconnects the best of them without probing +/// blindly. +struct Supervisor { + discoverer: Arc>, + connected: Arc>>, + others: Arc>>, + inbound: UnboundedSender, + shutdown: CancellationToken, + bytes: Arc, + events: broadcast::Sender, + best_tip: Arc, + max_peers: usize, + required_services: ServiceFlags, +} - // Give small positive reputation if peer maintained long connection - let conn_duration = Duration::from_secs(60 * loop_iteration); // Rough estimate - if conn_duration > Duration::from_secs(3600) { - // 1 hour - reputation_manager.update_reputation(addr, ChangeReason::LongUptime).await; +impl Supervisor { + async fn run(self) { + loop { + let at_cap = self.connected.lock().await.len() >= self.max_peers; + let dur = if at_cap { + self.improve_round().await; + IMPROVE_TICK + } else { + self.fill_round().await; + FILL_TICK + }; + tokio::select! { + _ = self.shutdown.cancelled() => break, + _ = tokio::time::sleep(dur) => {} } - }); + } } - /// Start the request processing task for outgoing messages from managers via RequestSender. - async fn start_request_processor(&self) { - // Take the receiver (only one task can own it) - let request_rx = { - let mut rx_guard = self.request_rx.lock().await; - rx_guard.take() + /// Up to `want` candidate addresses to probe: best-ranked backups first, then + /// fresh discovery, de-duplicated against each other and the live set. + async fn next_candidates(&self, want: usize) -> Vec { + let mut out: Vec = { + let mut o = self.others.lock().await; + o.sort_by_key(|p| p.lag_ms().unwrap_or(u32::MAX)); + let take = want.min(o.len()); + o.drain(..take).collect() }; + if out.len() < want { + out.extend(self.discoverer.lock().await.get(want - out.len()).await); + } + let live: HashSet = + self.connected.lock().await.iter().map(|(p, _)| p.addr()).collect(); + let mut seen = HashSet::new(); + out.retain(|p| !live.contains(&p.addr()) && seen.insert(p.addr())); + out + } - let Some(mut request_rx) = request_rx else { - tracing::warn!("Request processor already started or receiver unavailable"); - return; - }; + /// Connect a batch in parallel, keeping the successful handshakes and advancing + /// `best_tip` from whatever chain height they advertise. + async fn connect_chunk(&self, batch: Vec) -> Vec { + let results = join_all(batch.into_iter().map(|c| { + c.connect( + self.inbound.clone(), + self.shutdown.clone(), + self.bytes.clone(), + self.required_services, + ) + })) + .await; + let peers: Vec = results.into_iter().filter_map(Result::ok).collect(); + for p in &peers { + self.best_tip.fetch_max(p.version().start_height.max(0) as u32, Ordering::Relaxed); + } + peers + } - let this = self.clone(); - let shutdown_token = self.shutdown_token.clone(); + /// Close probed-but-unused peers and remember them as ranked backups, keeping the + /// list de-duplicated (best ping per address) and bounded. + async fn stash_backups(&self, peers: impl IntoIterator) { + let mut o = self.others.lock().await; + for peer in peers { + peer.close(); + o.push(peer.disconnect()); + } + // Dedup by address keeping the lowest ping, then rank and cap the list. + o.sort_by(|a, b| { + a.addr() + .cmp(&b.addr()) + .then(a.lag_ms().unwrap_or(u32::MAX).cmp(&b.lag_ms().unwrap_or(u32::MAX))) + }); + o.dedup_by_key(|p| p.addr()); + o.sort_by_key(|p| p.lag_ms().unwrap_or(u32::MAX)); + o.truncate(self.max_peers * BACKUP_MULTIPLE); + } - let mut tasks = self.tasks.lock().await; - tasks.spawn(async move { - tracing::info!("Starting request processor task"); - loop { - tokio::select! { - request = request_rx.recv() => { - match request { - Some(NetworkRequest::SendMessage(msg)) => { - tracing::trace!("Request processor: sending {}", msg.cmd()); - // Spawn each send concurrently to allow parallel requests across peers. - let this = this.clone(); - tokio::spawn(async move { - let result = match &msg { - // Distribute across peers for parallel sync - NetworkMessage::GetCFHeaders(_) - | NetworkMessage::GetCFilters(_) - | NetworkMessage::GetData(_) - | NetworkMessage::GetMnListD(_) - | NetworkMessage::GetQRInfo(_) - | NetworkMessage::GetHeaders(_) - | NetworkMessage::GetHeaders2(_) => { - this.send_distributed(msg).await - } - _ => { - this.send_to_single_peer(msg).await - } - }; - if let Err(e) = result { - tracing::error!("Request processor: failed to send message: {}", e); - } - }); - } - Some(NetworkRequest::SendMessageToPeer(msg, peer_address)) => { - tracing::trace!("Request processor: sending {} to peer {}", msg.cmd(), peer_address); - let this = this.clone(); - tokio::spawn(async move { - let fallback_msg = msg.clone(); - let result = match this.pool.get_peer(&peer_address).await { - Some(peer) => match this.send_message_to_peer(&peer_address, &peer, msg).await { - Ok(()) => Ok(()), - Err(err) => { - tracing::warn!( - "Target peer {} send failed ({}), falling back to distributed send", - peer_address, - err - ); - this.send_distributed(fallback_msg).await - } - }, - None => { - tracing::warn!( - "Target peer {} disconnected, falling back to distributed send", - peer_address - ); - this.send_distributed(fallback_msg).await - } - }; - if let Err(e) = result { - tracing::error!("Request processor: failed to send message to peer {}: {}", peer_address, e); - } - }); - } - Some(NetworkRequest::BroadcastMessage(msg)) => { - tracing::debug!("Request processor: broadcasting {}", msg.cmd()); - let this = this.clone(); - tokio::spawn(async move { - let results = this.broadcast(msg).await; - let failures = results.iter().filter(|r| r.is_err()).count(); - if failures > 0 { - tracing::warn!( - "Request processor: broadcast had {} failures out of {} peers", - failures, - results.len() - ); - } - }); - } - None => { - tracing::info!("Request processor: channel closed"); - break; - } - } - } - _ = shutdown_token.cancelled() => { - tracing::info!("Request processor: shutting down"); - break; - } - } - } + async fn announce_update(&self) { + let count = self.connected.lock().await.len() as u32; + let _ = self.events.send(NetworkEvent::PeersUpdated { + connected_count: count, + best_height: self.best_tip.load(Ordering::Relaxed), }); } - pub(crate) async fn evict_mismatched_peers(&self) { - if self.required_services == ServiceFlags::NONE { + /// Below cap: probe and accept decent peers, emitting `PeersUpdated` as they land. + async fn fill_round(&self) { + if self.max_peers.saturating_sub(self.connected.lock().await.len()) == 0 { return; } - let all_peers = self.pool.get_all_peers().await; - let connected_count = all_peers.len(); - if connected_count <= 1 { + let batch = self.next_candidates(CONNECT_CHUNK).await; + if batch.is_empty() { return; } - let mut matched_count = 0; - let mut mismatched = Vec::new(); - for (addr, peer) in &all_peers { - let peer_guard = peer.read().await; - if peer_guard.services_known() && peer_guard.has_service(self.required_services) { - matched_count += 1; - } else if peer_guard.services_known() { - mismatched.push(*addr); + let mut probed = self.connect_chunk(batch).await; + probed.sort_by_key(lag_key); // best first + + let mut accepted = 0usize; + let mut leftover: Vec = Vec::new(); + for peer in probed { + let lag = peer.lag_ms(); + let acceptable = lag > 0 && lag < BAD_LAG_MS; + if acceptable && self.connected.lock().await.len() < self.max_peers { + let addr = peer.addr(); + self.connected.lock().await.push((peer, State {})); + let _ = self.events.send(NetworkEvent::PeerConnected(addr)); + accepted += 1; + } else { + leftover.push(peer); } } - if mismatched.is_empty() { - return; - } - let drop_count = if matched_count > 0 { - mismatched.len() - } else { - mismatched.len().min(connected_count - 1) - }; - if drop_count == 0 { - return; - } - tracing::info!( - "Capability churn: dropping {} of {} peers lacking required services", - drop_count, - connected_count, - ); - for addr in mismatched.into_iter().take(drop_count) { - self.record_capability_rejection(addr).await; - let _ = self - .disconnect_peer( - &addr, - &format!("missing required services ({})", self.required_services), - ) - .await; - } - } - async fn maintenance_tick(&self) { - // Remove peers that the reader loop failed to clean up. - // This should not trigger under normal operation. - let unhealthy = self.pool.remove_unhealthy().await; - for addr in &unhealthy { - tracing::warn!("Maintenance removed stale peer {} - reader loop missed cleanup", addr); - Self::notify_peer_removed( - &self.pool, + // Last resort: never sit at zero peers. If nothing decent connected and the + // set is empty, take the least-bad handshake we got so sync can start; the + // improve loop upgrades it once a decent peer appears. + if accepted == 0 && self.connected.lock().await.is_empty() && !leftover.is_empty() { + leftover.sort_by_key(lag_key); + let peer = leftover.remove(0); + let addr = peer.addr(); + tracing::warn!( + target: "dash_spv::network", + "no decent peer available; accepting {} (ping {}ms) as last resort", addr, - &self.connected_peer_count, - &self.network_event_sender, - ) - .await; + peer.lag_ms(), + ); + self.connected.lock().await.push((peer, State {})); + let _ = self.events.send(NetworkEvent::PeerConnected(addr)); + accepted += 1; } - let count = self.pool.peer_count().await; - tracing::debug!("Connected peers: {}", count); - // Keep the cached counter in sync with actual pool count - self.connected_peer_count.store(count, Ordering::Relaxed); - if self.exclusive_mode { - // In exclusive mode, only reconnect to originally specified peers - for addr in self.initial_peers.iter() { - if !self.pool.is_connected(addr).await && !self.pool.is_connecting(addr).await { - tracing::info!("Reconnecting to exclusive peer: {}", addr); - self.connect_to_peer(*addr).await; - } - } - } else { - // Evict peers that lack required services before top-up so replacements - // can be pulled in during the same tick. - self.evict_mismatched_peers().await; - // Re-read count after potential churn so top-up sees the current pool size. - let count = self.pool.peer_count().await; - if count < self.max_peers { - // Try known addresses first, sorted by reputation - let known = self.addrv2_handler.get_known_addresses().await; - let needed = self.max_peers.saturating_sub(count); - // Select best peers based on reputation - let best_peers = self.reputation_manager.select_best_peers(known, needed * 2).await; - let mut attempted = 0; - - for addr in best_peers { - if self.is_capability_rejected(&addr).await { - continue; - } - if !self.pool.is_connected(&addr).await && !self.pool.is_connecting(&addr).await - { - self.connect_to_peer(addr).await; - attempted += 1; - if attempted >= needed { - break; - } - } - } - } + self.stash_backups(leftover).await; + + if accepted > 0 { + self.announce_update().await; + tracing::info!( + target: "dash_spv::network", + "peer supervisor: +{} peers -> {}", + accepted, + self.connected.lock().await.len(), + ); } + } - if self.shutdown_token.is_cancelled() { + /// At cap: probe a few candidates and swap the slowest connected peer for a + /// clearly-faster one, retiring the displaced peer so its in-flight work drains. + async fn improve_round(&self) { + let batch = self.next_candidates(IMPROVE_PROBE).await; + if batch.is_empty() { return; } - - // Send ping to all peers if needed and disconnect unresponsive ones - for (addr, peer) in self.pool.get_all_peers().await { - let mut peer_guard = peer.write().await; - if peer_guard.should_ping() { - if let Err(e) = peer_guard.send_ping().await { - tracing::error!("Failed to ping {}: {}", addr, e); - // Update reputation for ping failure - self.reputation_manager.update_reputation(addr, ChangeReason::PingFailed).await; + let mut probed = self.connect_chunk(batch).await; + probed.sort_by_key(lag_key); // best first + + let mut swapped: Option<(SocketAddr, ConnectedPeer)> = None; + if let Some(cand_lag) = probed.first().map(ConnectedPeer::lag_ms) { + if cand_lag > 0 && cand_lag < DECENT_LAG_MS { + let mut peers = self.connected.lock().await; + if let Some((pos, worst_lag)) = peers + .iter() + .enumerate() + .map(|(i, (p, _))| (i, lag_key(p))) + .max_by_key(|&(_, l)| l) + { + let clearly_better = worst_lag > DECENT_LAG_MS + && cand_lag.saturating_mul(SWAP_IMPROVEMENT) <= worst_lag; + if clearly_better && peers.len() >= self.max_peers { + let candidate = probed.remove(0); + let new_addr = candidate.addr(); + let (old, _) = peers.swap_remove(pos); + peers.push((candidate, State {})); + swapped = Some((new_addr, old)); + } } } - let has_expired = peer_guard.remove_expired_pings(); - drop(peer_guard); - if has_expired { - let _ = self.disconnect_peer(&addr, "ping timeout").await; - } } - // Only save known peers if not in exclusive mode - if !self.exclusive_mode { - let addresses = self.addrv2_handler.get_known_addresses().await; - if !addresses.is_empty() { - if let Err(e) = self.peer_store.save_peers(&addresses).await { - tracing::warn!("Failed to save peers: {}", e); - } - } - - // Save reputation data periodically - if let Err(e) = self.reputation_manager.save_to_storage(&*self.peer_store).await { - tracing::warn!("Failed to save reputation data: {}", e); - } + if let Some((new_addr, old)) = swapped { + let old_addr = old.addr(); + tracing::info!( + target: "dash_spv::network", + "peer supervisor: swapped out slow {} for faster {}", + old_addr, + new_addr, + ); + // Keep the displaced peer alive until its in-flight requests drain or time + // out, then close it — don't strand work already routed to it. + retire_drained(old, self.shutdown.clone()); + let _ = self.events.send(NetworkEvent::PeerConnected(new_addr)); + let _ = self.events.send(NetworkEvent::PeerDisconnected(old_addr)); + self.announce_update().await; } + + self.stash_backups(probed).await; } +} - async fn dns_fallback_tick(&self) { - let count = self.pool.peer_count().await; - if count >= self.max_peers { - return; +#[allow(clippy::too_many_arguments)] +fn spawn_peer_supervisor( + discoverer: Arc>, + connected: Arc>>, + others: Arc>>, + inbound: UnboundedSender, + shutdown: CancellationToken, + bytes: Arc, + events: broadcast::Sender, + best_tip: Arc, + max_peers: usize, + required_services: ServiceFlags, +) -> JoinHandle<()> { + tokio::spawn( + Supervisor { + discoverer, + connected, + others, + inbound, + shutdown, + bytes, + events, + best_tip, + max_peers, + required_services, } - let dns_peers = tokio::select! { - peers = self.discovery.discover_peers(self.network) => peers, - _ = self.shutdown_token.cancelled() => { - tracing::info!("Maintenance loop shutting down during DNS discovery"); - return - } - }; - let needed = self.max_peers.saturating_sub(count); - tracing::debug!("DNS fallback tick found {} addresses. Needed {}", dns_peers.len(), needed); - let mut dns_attempted = 0; - for addr in dns_peers.iter() { - if self.is_capability_rejected(addr).await { - continue; - } - if !self.pool.is_connected(addr).await && !self.pool.is_connecting(addr).await { - self.connect_to_peer(*addr).await; - dns_attempted += 1; - if dns_attempted >= needed { - break; - } + .run(), + ) +} + +/// Time requests out and evict the peers that stalled on them. +/// +/// A peer is a culprit when it has in-flight work (`in_flight > 0`) AND no bytes +/// have arrived from it for a full [`REQUEST_TIMEOUT`] — i.e. it has gone silent +/// while owing us responses. Kicking is on peer liveness, not on any single +/// request's age: a healthy peer draining a deep backlog (the per-peer cap grows +/// with headroom, so we may have many requests queued at it) keeps sending bytes, +/// so its liveness clock keeps resetting even while its OLDEST request sits waiting +/// its turn. Only a peer that has stopped sending anything is worth dropping. +/// +/// Judging on the peer's own counters (not the broker registry) is deliberate: the +/// registry can desync from a peer's real `in_flight` — correlating a response +/// frees the registry key, while streaming in-flight units are freed on the peer +/// separately — so a wedged peer can pin `in_flight == cap` (starving the router, +/// which only sends to peers under cap) while showing zero registry entries. A +/// registry-based check would miss exactly that peer. +/// +/// When one is found we kick it immediately: every request routed to it is now +/// dead, so we pull ALL its on-wire entries, re-inject their messages (retry to a +/// fresh peer), and drop the connection. Its in-flight slots die with it — no +/// separate reclaim — and the reconnector refills the peer set. Requests the +/// registry lost track of are re-driven by their pipeline's wanted set. Immediate +/// kick is deliberate: leaving a dead peer connected would just route the freed +/// work straight back to it (the router fills the emptiest peer first). +fn spawn_timeout_monitor( + requests: Registry, + connected: Arc>>, + queue: Arc, + shutdown: CancellationToken, +) -> JoinHandle<()> { + tokio::spawn(async move { + let mut ticker = tokio::time::interval(TIMEOUT_CHECK); + // Per-peer liveness: (bytes read at last progress, when it last rose). + let mut progress: HashMap = HashMap::new(); + loop { + tokio::select! { + _ = shutdown.cancelled() => break, + _ = ticker.tick() => {} } - } - } - /// Start peer connection maintenance loop - async fn start_maintenance_loop(&self) { - let this = self.clone(); - let mut tasks = self.tasks.lock().await; - tasks.spawn(async move { - // Periodic DNS discovery check (only active in non-exclusive mode) - let mut dns_interval = - time::interval_at(Instant::now() + DNS_DISCOVERY_DELAY, DNS_DISCOVERY_DELAY); - // Periodic reconnection check (active in both modes) - let mut maintenance_interval = time::interval(MAINTENANCE_INTERVAL); - let mut network_events = this.network_event_sender.subscribe(); - while !this.shutdown_token.is_cancelled() { - tokio::select! { - _ = maintenance_interval.tick() => { - tracing::trace!("Maintenance interval elapsed"); - this.maintenance_tick().await; + let now = Instant::now(); + + // Liveness is judged on the peer itself, not on the broker registry. + // The registry can desync from a peer's real in-flight count (a request + // whose response we correlate frees the registry key, while streaming + // in-flight units are freed on the peer separately), so a wedged peer can + // hold `in_flight == cap` — blocking the router from sending it anything — + // while showing zero registry entries. Watching `in_flight` (has pending + // work) and `bytes_read` (is data still arriving — covers both completed + // requests and mid-stream progress) catches that: a peer with work + // outstanding whose byte counter is frozen for a full REQUEST_TIMEOUT is + // stuck and must be dropped, whatever the registry thinks. + let mut culprits: HashSet = { + let peers = connected.lock().await; + let live: HashSet = peers.iter().map(|(p, _)| p.addr()).collect(); + progress.retain(|addr, _| live.contains(addr)); + let mut culprits = HashSet::new(); + for (peer, _) in peers.iter() { + let addr = peer.addr(); + let bytes = peer.bytes_read(); + let entry = progress.entry(addr).or_insert((bytes, now)); + if bytes > entry.0 { + *entry = (bytes, now); } - _ = dns_interval.tick(), if !this.exclusive_mode => { - this.dns_fallback_tick().await; + if peer.in_flight() > 0 && now.duration_since(entry.1) > REQUEST_TIMEOUT { + culprits.insert(addr); } - event = network_events.recv() => { - match event { - Ok(event) => { - tracing::debug!("Network event in maintenance loop: {}", event); - dns_interval.reset(); - this.maintenance_tick().await; - } - Err(error) => { - tracing::error!("Network event error: {}", error); - break; - } + } + culprits + }; + // A request that stopped making progress also condemns its peer: the + // per-peer byte check above sees unrelated traffic and judges it healthy. + { + let reqs = requests.lock().await; + for state in reqs.values() { + if let ReqState::OnWire(o) = state { + if now.duration_since(o.last_progress) > REQUEST_TIMEOUT { + culprits.insert(o.peer); } } - _ = this.shutdown_token.cancelled() => { - tracing::info!("Maintenance loop shutting down"); - break; - } } } - }); - } - - /// Send a message to a single peer selected by message type requirements. - async fn send_to_single_peer(&self, message: NetworkMessage) -> NetworkResult<()> { - let peers = self.pool.get_all_peers().await; - if peers.is_empty() { - return Err(NetworkError::ConnectionFailed("No connected peers".to_string())); - } - - let preferred_service = match &message { - NetworkMessage::FilterLoad(_) - | NetworkMessage::FilterClear - | NetworkMessage::MemPool => Some((ServiceFlags::BLOOM, true)), - NetworkMessage::GetCFHeaders(_) | NetworkMessage::GetCFilters(_) => { - Some((ServiceFlags::COMPACT_FILTERS, true)) - } - NetworkMessage::GetHeaders(_) | NetworkMessage::GetHeaders2(_) => { - Some((ServiceFlags::NODE_HEADERS_COMPRESSED, false)) + if culprits.is_empty() { + continue; } - _ => None, - }; - let (addr, peer) = if let Some((flags, required)) = preferred_service { - match self.pool.peer_with_service(flags).await { - Some((address, peer)) => { - tracing::debug!( - "Selected peer {} with {} for {}", - address, - flags, - message.cmd() - ); - (address, peer) - } - None if required => { - tracing::warn!("No peers support {}, cannot send {}", flags, message.cmd()); - return Err(NetworkError::ProtocolError(format!("No peers support {}", flags))); + // Pull every on-wire request routed to a culprit (fresh ones included — + // the connection is going away, so they are dead too) and re-inject its + // message (key kept as Queued so de-dup still holds). + let mut reinject: Vec = Vec::new(); + { + let mut reqs = requests.lock().await; + let keys: Vec = reqs + .iter() + .filter_map(|(k, s)| match s { + ReqState::OnWire(o) if culprits.contains(&o.peer) => Some(k.clone()), + _ => None, + }) + .collect(); + for key in keys { + if let Some(ReqState::OnWire(o)) = reqs.insert(key, ReqState::Queued) { + reinject.push(o.msg); + } } - None => self.next_peer(&peers), } - } else { - self.next_peer(&peers) - }; - self.send_message_to_peer(&addr, &peer, message).await - } + // Drop the culprits still in the active set (some may already be gone + // — a retired-drained peer isn't here — which is fine). + let dropped = { + let mut peers = connected.lock().await; + let before = peers.len(); + peers.retain(|(p, _)| { + if culprits.contains(&p.addr()) { + p.close(); + false + } else { + true + } + }); + before - peers.len() + }; - /// Send a message distributed across connected peers using round-robin selection. - /// - /// Peer selection and message handling based on message type: - /// - Filters (GetCFHeaders/GetCFilters): requires peers that support compact filters - /// - Headers (GetHeaders/GetHeaders2): prefers headers2 peers, upgrades GetHeaders if supported - /// - Other (blocks, masternode data, etc.): uses all connected peers - async fn send_distributed(&self, message: NetworkMessage) -> NetworkResult<()> { - let peers = self.pool.get_all_peers().await; - - if peers.is_empty() { - return Err(NetworkError::ConnectionFailed("No connected peers".to_string())); - } + tracing::warn!( + target: "dash_spv::network", + "request timeout: kicked {} peer(s) {:?}, retried {} request(s)", + dropped, + culprits, + reinject.len(), + ); - // Select eligible peers based on message type - let (selected_peers, require_capability) = match &message { - NetworkMessage::GetCFHeaders(_) | NetworkMessage::GetCFilters(_) => { - let filter_peers = - self.pool.peers_with_service(ServiceFlags::COMPACT_FILTERS).await; - (filter_peers, true) - } - NetworkMessage::GetHeaders(_) | NetworkMessage::GetHeaders2(_) => { - // Prefer headers2 peers (excluding disabled), fall back to all - let disabled = self.headers2_disabled.lock().await; - let mut headers2_peers = - self.pool.peers_with_service(ServiceFlags::NODE_HEADERS_COMPRESSED).await; - headers2_peers.retain(|(addr, _)| !disabled.contains(addr)); - drop(disabled); - if headers2_peers.is_empty() { - (peers.clone(), false) - } else { - (headers2_peers, false) - } - } - _ => { - // All other messages use all connected peers - (peers.clone(), false) + for msg in reinject { + queue.push(msg).await; } - }; + // Freed capacity (dropped peers) — wake the router to re-evaluate. + queue.notify.notify_one(); + } + }) +} - if selected_peers.is_empty() { - return if require_capability { - Err(NetworkError::ProtocolError("No peers support required capability".to_string())) - } else { - Err(NetworkError::ConnectionFailed("No connected peers".to_string())) - }; +/// Retire a peer we are dropping from the active set WITHOUT stranding requests +/// already on the wire to it. During startup the reconnector connects peers and +/// the sync sends them pipeline requests before the probe has settled the final +/// peer set; replacing the set would then drop those peers mid-request, and the +/// stranded requests only recover on the pipelines' own (slow) timeout — +/// occasionally stalling whole header segments for a run. Instead keep the +/// connection alive in the background: its reader keeps delivering responses and +/// decrementing `in_flight`. Close it once it has drained, or after +/// `RETIRE_DRAIN_CAP` (a peer that never drains is dead), whichever comes first. +fn retire_drained(peer: ConnectedPeer, shutdown: CancellationToken) { + if peer.in_flight() == 0 { + peer.close(); + return; + } + tokio::spawn(async move { + let mut waited = Duration::ZERO; + while peer.in_flight() > 0 && waited < RETIRE_DRAIN_CAP { + tokio::select! { + _ = shutdown.cancelled() => return, + _ = tokio::time::sleep(DRAIN_POLL) => waited += DRAIN_POLL, + } } + peer.close(); + }); +} - let (addr, peer) = self.next_peer(&selected_peers); +fn spawn_pump( + mut inbound: UnboundedReceiver, + subscribers: Subscribers, + connected: Arc>>, + events: broadcast::Sender, + queue: Arc, + requests: Registry, + shutdown: CancellationToken, +) -> JoinHandle<()> { + tokio::spawn(async move { + // Per-peer received-message counter to check load balance across peers. + let mut recv_by_peer: HashMap = HashMap::new(); + let mut total_recv: u64 = 0; + loop { + let event = tokio::select! { + _ = shutdown.cancelled() => break, + ev = inbound.recv() => match ev { + Some(ev) => ev, + None => break, + }, + }; + match event { + PeerEvent::Message(addr, msg) => { + *recv_by_peer.entry(addr).or_insert(0) += 1; + total_recv += 1; + + // A `cfilter` is one piece of a `getcfilters` batch, which is only + // marked answered once the whole batch lands. Refresh that request's + // deadline so the timeout means "the pieces stopped coming". + if matches!(msg, NetworkMessage::CFilter(_)) { + let mut reqs = requests.lock().await; + for state in reqs.values_mut() { + if let ReqState::OnWire(o) = state { + if o.peer == addr + && matches!( + request_keys(&o.msg).first(), + Some(RequestKey::CFilters(_)) + ) + { + o.last_progress = Instant::now(); + } + } + } + } + if total_recv.is_multiple_of(250_000) { + let mut dist: Vec<(SocketAddr, u64)> = + recv_by_peer.iter().map(|(a, c)| (*a, *c)).collect(); + dist.sort_by_key(|(_, c)| std::cmp::Reverse(*c)); + tracing::info!( + target: "filt_depth", + "recv balance ({} peers, {} total): {:?}", + dist.len(), + total_recv, + dist, + ); + // Per-peer request->response latency (avg / worst). + let lat: Vec<(SocketAddr, u64, String)> = connected + .lock() + .await + .iter() + .map(|(p, _)| { + let (n, avg, max) = p.latency_stats(); + (p.addr(), n, format!("avg={avg:.1}ms max={max:.1}ms")) + }) + .collect(); + tracing::info!( + target: "filt_depth", + "peer latency: {:?}", + lat, + ); + } + let mt = MessageType::from_cmd(msg.cmd()); + // Single-message responses free a slot in the peer's reader; + // wake the router so it can use the freed capacity. `cfilter` + // is skipped — its batch frees a slot via `request_completed`, + // which notifies once per ~1000 messages instead of each. + if mt != Some(MessageType::CFilter) { + queue.notify.notify_one(); + } - tracing::trace!("Distributing {} request to peer {}", message.cmd(), addr); + { + let mut subscribers = subscribers.lock().await; + if let Some(list) = mt.and_then(|t| subscribers.get_mut(&t)) { + let last = list.len().saturating_sub(1); + let mut msg = Some(Arc::new(msg)); + let mut idx = 0; + + list.retain(|tx| { + // Hand the *last* subscriber our own reference instead of a + // clone. Every heavy message type (block, cfilter, cfheaders, + // headers) has exactly one subscriber, so it arrives with a + // refcount of 1 and the manager can take the payload without + // copying it. Only `inv`, which is tiny, is really shared. + let shared = if idx == last { + msg.take().expect("taken once, on the final subscriber") + } else { + Arc::clone( + msg.as_ref().expect("held until the final subscriber"), + ) + }; + idx += 1; + + tx.send((addr, shared)).is_ok() + }); + } + } + } + PeerEvent::Disconnected(addr) => { + let remaining = { + let mut guard = connected.lock().await; + guard.retain(|(peer, _)| peer.addr() != addr); + guard.len() + }; + + // A peer vanishing takes its in-flight requests with it: every + // request routed to it is now dead and no one else is tracking it. + // Pull them back to Queued and re-inject the messages so the router + // retries them on a live peer — the timeout monitor only watches + // connected peers, so a request whose peer is already gone would + // otherwise leak in the registry forever. + let mut reinject: Vec = Vec::new(); + { + let mut reqs = requests.lock().await; + let keys: Vec = reqs + .iter() + .filter_map(|(k, s)| match s { + ReqState::OnWire(o) if o.peer == addr => Some(k.clone()), + _ => None, + }) + .collect(); + for key in keys { + if let Some(ReqState::OnWire(o)) = reqs.insert(key, ReqState::Queued) { + reinject.push(o.msg); + } + } + } - self.send_message_to_peer(&addr, &peer, message).await - } + tracing::info!( + target: "dash_spv::network", + "peer disconnected: {} | {} peers remaining | {} request(s) re-queued", + addr, + remaining, + reinject.len(), + ); - /// Pick the next peer from `peers` using round-robin rotation. - fn next_peer( - &self, - peers: &[(SocketAddr, Arc>)], - ) -> (SocketAddr, Arc>) { - let idx = self.round_robin_counter.fetch_add(1, Ordering::Relaxed) % peers.len(); - (peers[idx].0, peers[idx].1.clone()) - } + for msg in reinject { + queue.push(msg).await; + } + queue.notify.notify_one(); - /// Send a message to the given peer. - /// For GetHeaders messages upgrade to GetHeaders2 if the peer supports it. - async fn send_message_to_peer( - &self, - addr: &SocketAddr, - peer: &Arc>, - message: NetworkMessage, - ) -> NetworkResult<()> { - let message = match message { - NetworkMessage::GetHeaders(get_headers) => { - let supports_headers2 = peer.read().await.can_request_headers2(); - if supports_headers2 && !self.headers2_disabled.lock().await.contains(addr) { - tracing::debug!("Upgrading GetHeaders to GetHeaders2 for peer {}", addr); - NetworkMessage::GetHeaders2(get_headers) - } else { - NetworkMessage::GetHeaders(get_headers) + let _ = events.send(NetworkEvent::PeerDisconnected(addr)); } } - other => other, - }; + } + }) +} - let mut peer_guard = peer.write().await; - peer_guard - .send_message(message) - .await - .map_err(|e| NetworkError::ProtocolError(format!("Failed to send to {}: {}", addr, e))) +impl MsgQueue { + fn new() -> Self { + Self { + other: Mutex::new(VecDeque::with_capacity(30)), + blocks: Mutex::new(VecDeque::with_capacity(30)), + cfilters: Mutex::new(VecDeque::with_capacity(30)), + cfheaders: Mutex::new(VecDeque::with_capacity(30)), + headers: Mutex::new(VecDeque::with_capacity(30)), + len: AtomicUsize::new(0), + notify: Notify::new(), + } } - /// Broadcast a message to all connected peers - pub async fn broadcast(&self, message: NetworkMessage) -> Vec> { - let peers = self.pool.get_all_peers().await; - let mut handles = Vec::new(); - - // Spawn tasks for concurrent sending - for (addr, peer) in peers { - // Reduce verbosity for common sync messages - match &message { - NetworkMessage::GetHeaders(_) | NetworkMessage::GetCFilters(_) => { - tracing::debug!("Broadcasting {} to {}", message.cmd(), addr); - } - _ => { - tracing::trace!("Broadcasting {:?} to {}", message.cmd(), addr); - } - } - let msg = message.clone(); + fn len(&self) -> usize { + self.len.load(Ordering::SeqCst) + } - let handle = tokio::spawn(async move { - let mut peer_guard = peer.write().await; - peer_guard.send_message(msg).await.map_err(Error::Network) - }); - handles.push(handle); + fn queue(&self, class: MsgClass) -> &Mutex> { + match class { + MsgClass::Other => &self.other, + MsgClass::Blocks => &self.blocks, + MsgClass::CFilters => &self.cfilters, + MsgClass::CfHeaders => &self.cfheaders, + MsgClass::Headers => &self.headers, } + } - // Wait for all sends to complete - let mut results = Vec::new(); - for handle in handles { - match handle.await { - Ok(result) => results.push(result), - Err(_) => results.push(Err(Error::Network(NetworkError::ConnectionFailed( - "Task panicked during broadcast".to_string(), - )))), + /// Take up to `n` messages in strict priority order ([`DRAIN_PRIORITY`]): + /// control traffic, then blocks, filters, filter headers, block headers. A + /// class is fully drained (up to the remaining budget) before the next is + /// touched, so a large backlog of one type never blocks another behind it. + async fn pop_n(&self, n: usize) -> Vec { + if n == 0 { + return Vec::new(); + } + let mut out: Vec = Vec::with_capacity(n); + for class in DRAIN_PRIORITY { + if out.len() >= n { + break; } + let mut q = self.queue(class).lock().await; + let take = (n - out.len()).min(q.len()); + out.extend(q.drain(..take)); } + self.len.fetch_sub(out.len(), Ordering::SeqCst); + out + } - results + async fn push(&self, msg: NetworkMessage) { + self.queue(classify(&msg)).lock().await.push_back(msg); + self.len.fetch_add(1, Ordering::SeqCst); + self.notify.notify_one(); } - /// Disconnect a specific peer - pub async fn disconnect_peer(&self, addr: &SocketAddr, reason: &str) -> Result<(), Error> { - tracing::info!("Disconnecting peer {} - reason: {}", addr, reason); + /// Return messages that were popped but could not be sent to the FRONT of + /// their class queue, preserving order. A popped message must never be dropped: + /// the owning pipeline has already recorded it as handed to the network and + /// only starts its response timeout once the router reports it on the wire, so + /// a dropped message is one the pipeline waits on forever — a permanent sync + /// stall (observed as: queue backed up, 0 MB/s, no sends, no timeouts). + async fn push_front_all(&self, msgs: Vec) { + if msgs.is_empty() { + return; + } + let n = msgs.len(); + for msg in msgs.into_iter().rev() { + self.queue(classify(&msg)).lock().await.push_front(msg); + } + self.len.fetch_add(n, Ordering::SeqCst); + self.notify.notify_one(); + } +} - Self::remove_peer_and_notify( - &self.pool, - addr, - &self.connected_peer_count, - &self.network_event_sender, - ) - .await; +#[cfg(test)] +mod tests { + use super::*; + use crate::client::ClientConfig; + use dashcore::network::message_blockdata::GetHeadersMessage; + use dashcore::network::message_filter::{GetCFHeaders, GetCFilters}; + use dashcore::network::message_sml::GetMnListDiff; + use dashcore::{BlockHash, Txid}; + use dashcore_hashes::Hash; + + fn get_headers(locator: u32) -> NetworkMessage { + NetworkMessage::GetHeaders(GetHeadersMessage::new( + vec![BlockHash::dummy(locator)], + BlockHash::dummy(0), + )) + } - Ok(()) + fn get_cfilters(start: u32) -> NetworkMessage { + NetworkMessage::GetCFilters(GetCFilters { + filter_type: 0, + start_height: start, + stop_hash: BlockHash::dummy(9), + }) } - /// Get reputation information for all peers - pub async fn get_peer_reputations(&self) -> HashMap { - let reputations = self.reputation_manager.get_all_reputations().await; - reputations.into_iter().map(|(addr, rep)| (addr, (rep.score, rep.is_banned()))).collect() + fn get_cfheaders(stop: u32) -> NetworkMessage { + NetworkMessage::GetCFHeaders(GetCFHeaders { + filter_type: 0, + start_height: 0, + stop_hash: BlockHash::dummy(stop), + }) } - /// Ban a specific peer manually - pub async fn ban_peer(&self, addr: &SocketAddr, reason: &str) -> Result<(), Error> { - tracing::info!("Manually banning peer {} - reason: {}", addr, reason); + /// A broker with no peers: `new` spawns the pump, router, timeout monitor and + /// bandwidth controller but never touches the network — the peer supervisor is + /// only started by `start()`, which these tests deliberately do not call. + async fn broker() -> PeerNetworkManager { + let config = ClientConfig::regtest().with_restrict_to_configured_peers(true); + PeerNetworkManager::new(&config).await + } - // Disconnect the peer first - self.disconnect_peer(addr, reason).await?; + // ---- request key derivation (the dedup identity) ---- - // Update reputation to trigger ban - self.reputation_manager.update_reputation(*addr, ChangeReason::ManuallyBanned).await; + #[test] + fn request_keys_derives_one_key_per_request_type() { + assert_eq!(request_keys(&get_headers(1)), vec![RequestKey::Headers(BlockHash::dummy(1))]); + assert_eq!( + request_keys(&get_cfheaders(2)), + vec![RequestKey::CfHeaders(BlockHash::dummy(2))] + ); + assert_eq!(request_keys(&get_cfilters(100)), vec![RequestKey::CFilters(100)]); + assert_eq!( + request_keys(&NetworkMessage::GetMnListD(GetMnListDiff { + base_block_hash: BlockHash::dummy(3), + block_hash: BlockHash::dummy(4), + })), + vec![RequestKey::MnListDiff(BlockHash::dummy(4))] + ); + } - Ok(()) + /// One `getdata` may name several blocks; each is tracked as its own request + /// so a single missing block can time out and retry on its own. + #[test] + fn request_keys_splits_getdata_per_block_and_ignores_other_inventory() { + let keys = request_keys(&NetworkMessage::GetData(vec![ + Inventory::Block(BlockHash::dummy(1)), + Inventory::Transaction(Txid::from_byte_array([7; 32])), + Inventory::Block(BlockHash::dummy(2)), + ])); + assert_eq!( + keys, + vec![RequestKey::Block(BlockHash::dummy(1)), RequestKey::Block(BlockHash::dummy(2))] + ); } - /// Unban a specific peer - pub async fn unban_peer(&self, addr: &SocketAddr) { - self.reputation_manager.unban_peer(addr).await; + /// Traffic the broker does not track for timeout/retry has no key, so it is + /// never de-duplicated — two `mempool` messages must both go out. + #[test] + fn request_keys_is_empty_for_untracked_traffic() { + assert!(request_keys(&NetworkMessage::MemPool).is_empty()); + assert!(request_keys(&NetworkMessage::Ping(1)).is_empty()); + assert!(request_keys(&NetworkMessage::GetData(vec![Inventory::Transaction( + Txid::from_byte_array([1; 32]) + )])) + .is_empty()); } - /// Shutdown the network manager - pub async fn shutdown(&self) { - tracing::info!("Shutting down peer network manager"); - self.shutdown_token.cancel(); + // ---- de-duplication ---- - // Save known peers before shutdown - let addresses = self.addrv2_handler.get_addresses_for_peer(MAX_ADDR_TO_STORE).await; - if !addresses.is_empty() { - if let Err(e) = self.peer_store.save_peers(&addresses).await { - tracing::warn!("Failed to save peers on shutdown: {}", e); - } - } + #[tokio::test] + async fn send_dedups_a_request_already_in_play() { + let net = broker().await; - // Save reputation data before shutdown - if let Err(e) = self.reputation_manager.save_to_storage(&*self.peer_store).await { - tracing::warn!("Failed to save reputation data on shutdown: {}", e); - } + net.send(get_cfilters(0)).await; + assert_eq!(net.msg_queue.len(), 1); - // Drain tasks while holding the lock. connect_to_peer() already uses - // `select!` with the cancellation token when acquiring this lock, so no - // deadlock can occur once the shutdown token is cancelled above. - let mut tasks = self.tasks.lock().await; - while let Some(result) = tasks.join_next().await { - if let Err(e) = result { - tracing::error!("Task join error: {}", e); - } - } + // Same key while the first is still tracked: dropped before the queue. + net.send(get_cfilters(0)).await; + assert_eq!(net.msg_queue.len(), 1, "a re-declared request must not be queued twice"); - // Disconnect all peers - for addr in self.pool.get_connected_addresses().await { - self.pool.remove_peer(&addr).await; - } + // A different key is a different request. + net.send(get_cfilters(1000)).await; + assert_eq!(net.msg_queue.len(), 2); } - async fn record_capability_rejection(&self, addr: SocketAddr) { - Self::record_capability_rejection_in(&self.capability_rejected, addr).await; - } + /// The key is held for the whole lifecycle, so re-declaring is a no-op until + /// the requester reports the response — that is what makes it safe for a + /// pipeline to re-declare its whole wanted set on every tick. + #[tokio::test] + async fn request_answered_releases_the_key_for_re_declaration() { + let net = broker().await; - async fn is_capability_rejected(&self, addr: &SocketAddr) -> bool { - let mut rejected = self.capability_rejected.write().await; - let now = Instant::now(); - rejected.retain(|_, rejected_at| { - now.saturating_duration_since(*rejected_at) < CAPABILITY_REJECTED_TTL - }); - rejected.contains_key(addr) - } + net.send(get_headers(1)).await; + net.send(get_headers(1)).await; + assert_eq!(net.msg_queue.len(), 1); - async fn record_capability_rejection_in( - capability_rejected: &RwLock>, - addr: SocketAddr, - ) { - capability_rejected.write().await.insert(addr, Instant::now()); - } + net.request_answered(RequestKey::Headers(BlockHash::dummy(1))).await; + assert!(net.requests.lock().await.is_empty(), "an answered request stops being tracked"); - async fn should_reject_after_handshake( - pool: &PeerPool, - peer: &Peer, - required_services: ServiceFlags, - ) -> bool { - required_services != ServiceFlags::NONE - && pool.has_peers_with_service(required_services).await - && peer.services_known() - && !peer.has_service(required_services) + // Now the same request is accepted again. + net.send(get_headers(1)).await; + assert_eq!(net.msg_queue.len(), 2); } -} -// Implement Clone for use in async closures -impl Clone for PeerNetworkManager { - fn clone(&self) -> Self { - Self { - pool: self.pool.clone(), - max_peers: self.max_peers, - discovery: self.discovery.clone(), - addrv2_handler: self.addrv2_handler.clone(), - peer_store: self.peer_store.clone(), - reputation_manager: self.reputation_manager.clone(), - network: self.network, - shutdown_token: self.shutdown_token.clone(), - tasks: self.tasks.clone(), - initial_peers: self.initial_peers.clone(), - data_dir: self.data_dir.clone(), - user_agent: self.user_agent.clone(), - exclusive_mode: self.exclusive_mode, - required_services: self.required_services, - capability_rejected: self.capability_rejected.clone(), - connected_peer_count: self.connected_peer_count.clone(), - headers2_disabled: self.headers2_disabled.clone(), - message_dispatcher: self.message_dispatcher.clone(), - request_tx: self.request_tx.clone(), - request_rx: self.request_rx.clone(), - round_robin_counter: self.round_robin_counter.clone(), - network_event_sender: self.network_event_sender.clone(), - } - } -} + #[tokio::test] + async fn keyless_messages_are_never_deduplicated() { + let net = broker().await; -// Implement NetworkManager trait -#[async_trait] -impl NetworkManager for PeerNetworkManager { - async fn message_receiver(&mut self, types: &[MessageType]) -> UnboundedReceiver { - self.message_dispatcher.lock().await.message_receiver(types) - } + net.send(NetworkMessage::MemPool).await; + net.send(NetworkMessage::MemPool).await; - fn request_sender(&self) -> RequestSender { - PeerNetworkManager::request_sender(self) + assert_eq!(net.msg_queue.len(), 2); + assert!( + net.requests.lock().await.is_empty(), + "untracked traffic must not enter the registry" + ); } - async fn connect(&mut self) -> NetworkResult<()> { - self.start().await.map_err(|e| NetworkError::ConnectionFailed(e.to_string())) + // ---- strict-priority scheduling ---- + + #[test] + fn classify_maps_each_request_to_its_scheduling_class() { + assert_eq!(classify(&get_headers(1)), MsgClass::Headers); + assert_eq!(classify(&get_cfheaders(1)), MsgClass::CfHeaders); + assert_eq!(classify(&get_cfilters(0)), MsgClass::CFilters); + assert_eq!( + classify(&NetworkMessage::GetData(vec![Inventory::Block(BlockHash::dummy(1))])), + MsgClass::Blocks + ); + assert_eq!(classify(&NetworkMessage::MemPool), MsgClass::Other); + // A `getdata` that is not purely blocks is control traffic, not a block download. + assert_eq!( + classify(&NetworkMessage::GetData(vec![Inventory::Transaction( + Txid::from_byte_array([1; 32]) + )])), + MsgClass::Other + ); } - async fn disconnect(&mut self) -> NetworkResult<()> { - self.shutdown().await; - Ok(()) + /// A backlog of one class must never block another behind it: the router + /// drains control first, then blocks, filters, filter headers, block headers. + #[tokio::test] + async fn queue_drains_in_strict_priority_order() { + let q = MsgQueue::new(); + + // Pushed in reverse priority on purpose. + q.push(get_headers(1)).await; + q.push(get_cfheaders(1)).await; + q.push(get_cfilters(0)).await; + q.push(NetworkMessage::GetData(vec![Inventory::Block(BlockHash::dummy(1))])).await; + q.push(NetworkMessage::MemPool).await; + assert_eq!(q.len(), 5); + + let drained: Vec = q.pop_n(5).await.iter().map(classify).collect(); + assert_eq!( + drained, + vec![ + MsgClass::Other, + MsgClass::Blocks, + MsgClass::CFilters, + MsgClass::CfHeaders, + MsgClass::Headers + ] + ); + assert_eq!(q.len(), 0); } - async fn send_message(&mut self, message: NetworkMessage) -> NetworkResult<()> { - // For sync messages that require consistent responses, send to only one peer - match &message { - NetworkMessage::GetHeaders(_) - | NetworkMessage::GetHeaders2(_) - | NetworkMessage::GetCFHeaders(_) - | NetworkMessage::GetCFilters(_) - | NetworkMessage::GetData(_) - | NetworkMessage::GetMnListD(_) => self.send_to_single_peer(message).await, - _ => { - // For other messages, broadcast to all peers - let results = self.broadcast(message).await; - - // Return error if all sends failed - if results.is_empty() { - return Err(NetworkError::ConnectionFailed("No connected peers".to_string())); - } + /// A message popped but not sent goes back to the FRONT of its class, so it + /// keeps its place: dropping it would strand the pipeline waiting forever. + #[tokio::test] + async fn unsent_messages_go_back_to_the_front_of_their_class() { + let q = MsgQueue::new(); + q.push(get_cfilters(0)).await; + q.push(get_cfilters(1000)).await; - let successes = results.iter().filter(|r| r.is_ok()).count(); - if successes == 0 { - return Err(NetworkError::ProtocolError( - "Failed to send to any peer".to_string(), - )); - } + let popped = q.pop_n(1).await; + assert_eq!(q.len(), 1); - Ok(()) - } - } // end match - } // end send_message + q.push_front_all(popped).await; + assert_eq!(q.len(), 2); - fn peer_count(&self) -> usize { - // Use cached counter to avoid blocking in async context - self.connected_peer_count.load(Ordering::Relaxed) + // The returned message is handed out first again, ahead of the one behind it. + let order: Vec = q + .pop_n(2) + .await + .iter() + .map(|m| match m { + NetworkMessage::GetCFilters(g) => g.start_height, + other => panic!("unexpected {other:?}"), + }) + .collect(); + assert_eq!(order, vec![0, 1000]); } - async fn broadcast(&self, message: NetworkMessage) -> NetworkResult<()> { - let results = PeerNetworkManager::broadcast(self, message).await; - - if results.is_empty() { - return Err(NetworkError::ConnectionFailed("No connected peers".to_string())); - } - - let successes = results.iter().filter(|r| r.is_ok()).count(); - if successes == 0 { - return Err(NetworkError::ConnectionFailed("All broadcast sends failed".to_string())); + #[tokio::test] + async fn pop_n_respects_its_budget() { + let q = MsgQueue::new(); + for start in [0, 1000, 2000] { + q.push(get_cfilters(start)).await; } - Ok(()) + assert_eq!(q.pop_n(2).await.len(), 2); + assert_eq!(q.len(), 1); + assert!(q.pop_n(0).await.is_empty()); } - async fn dispatch_local(&self, message: NetworkMessage) { - let local_addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0)); - let msg = Message::new(local_addr, message); - self.message_dispatcher.lock().await.dispatch(&msg); - } + // ---- inbound routing ---- + + /// `dispatch_local` feeds a message through the same pump as a real peer, + /// tagged with the `0.0.0.0:0` sentinel that managers read as "self-originated". + #[tokio::test] + async fn dispatch_local_reaches_subscribers_with_the_local_sentinel() { + let net = broker().await; + let mut rx = net.subscribe(&[MessageType::Tx]).await; + + let tx = dashcore::Transaction { + version: 1, + lock_time: 0, + input: vec![], + output: vec![], + special_transaction_payload: None, + }; + net.dispatch_local(NetworkMessage::Tx(tx)).await; - async fn disconnect_peer(&self, addr: &SocketAddr, reason: &str) -> NetworkResult<()> { - PeerNetworkManager::disconnect_peer(self, addr, reason) + let (peer, msg) = tokio::time::timeout(Duration::from_secs(2), rx.recv()) .await - .map_err(|e| NetworkError::ConnectionFailed(e.to_string())) + .expect("the pump must deliver the injected message") + .expect("channel open"); + assert!(peer.ip().is_unspecified(), "locally injected messages carry the sentinel address"); + assert_eq!(peer.port(), 0); + assert!(matches!(*msg, NetworkMessage::Tx(_))); } - fn subscribe_network_events(&self) -> broadcast::Receiver { - self.network_event_sender.subscribe() - } -} + /// Subscriptions are per message type: a manager only wakes for what it asked for. + #[tokio::test] + async fn subscribers_only_receive_the_types_they_asked_for() { + let net = broker().await; + let mut headers_rx = net.subscribe(&[MessageType::Headers]).await; + + net.dispatch_local(NetworkMessage::Tx(dashcore::Transaction { + version: 1, + lock_time: 0, + input: vec![], + output: vec![], + special_transaction_payload: None, + })) + .await; -#[cfg(test)] -impl PeerNetworkManager { - pub(crate) async fn new_for_test(required_services: ServiceFlags) -> Self { - let test_dir = tempfile::tempdir().expect("test dir creation failed").keep(); - let peer_store = - PersistentPeerStorage::open(&test_dir).await.expect("test peer store init failed"); - let discovery = DnsDiscovery::new(); - let (request_tx, request_rx) = unbounded_channel(); - Self { - pool: Arc::new(PeerPool::new(8)), - max_peers: 8, - discovery: Arc::new(discovery), - addrv2_handler: Arc::new(AddrV2Handler::new()), - peer_store: Arc::new(peer_store), - reputation_manager: Arc::new(PeerReputationManager::new()), - network: Network::Testnet, - shutdown_token: CancellationToken::new(), - tasks: Arc::new(Mutex::new(JoinSet::new())), - initial_peers: vec![], - data_dir: test_dir, - user_agent: None, - exclusive_mode: false, - required_services, - capability_rejected: Arc::new(RwLock::new(HashMap::new())), - connected_peer_count: Arc::new(AtomicUsize::new(0)), - headers2_disabled: Arc::new(Mutex::new(HashSet::new())), - message_dispatcher: Arc::new(Mutex::new(MessageDispatcher::default())), - request_tx, - request_rx: Arc::new(Mutex::new(Some(request_rx))), - round_robin_counter: Arc::new(AtomicUsize::new(0)), - network_event_sender: broadcast::Sender::new(DEFAULT_NETWORK_EVENT_CAPACITY), - } + // Give the pump a chance to run before asserting nothing arrived. + tokio::time::sleep(Duration::from_millis(50)).await; + assert!(headers_rx.try_recv().is_err(), "a Headers subscriber must not see a Tx"); } - pub(crate) async fn insert_test_peer(&self, addr: SocketAddr, flags: ServiceFlags) { - self.pool.insert_peer_with_services(addr, flags).await; - self.connected_peer_count.fetch_add(1, Ordering::Relaxed); - } + // ---- peer-set accessors with no peers ---- - pub(crate) async fn test_peer_count(&self) -> usize { - self.pool.peer_count().await + #[tokio::test] + async fn reports_no_peers_and_no_tip_before_start() { + let net = broker().await; + assert_eq!(net.connected_count().await, 0); + assert_eq!(net.tip(), 0); + // Broadcasting with no peers is a no-op, not a panic. + net.broadcast(NetworkMessage::MemPool); } - pub(crate) async fn test_is_connected(&self, addr: &SocketAddr) -> bool { - self.pool.is_connected(addr).await + /// `inv` goes to several managers at once, so the pump must fan the same + /// message out to every subscriber of that type. + #[tokio::test] + async fn multiple_subscribers_of_the_same_type_all_receive_it() { + let net = broker().await; + let mut first = net.subscribe(&[MessageType::Inv]).await; + let mut second = net.subscribe(&[MessageType::Inv]).await; + + net.dispatch_local(NetworkMessage::Inv(vec![Inventory::Block(BlockHash::dummy(1))])).await; + + for rx in [&mut first, &mut second] { + let (_, msg) = tokio::time::timeout(Duration::from_secs(2), rx.recv()) + .await + .expect("both subscribers must be served") + .expect("channel open"); + assert!(matches!(*msg, NetworkMessage::Inv(_))); + } } - pub(crate) async fn insert_test_capability_rejected(&self, addr: SocketAddr) { - self.record_capability_rejection(addr).await; - } + /// Asking for the same type more than once must still deliver it once: a + /// manager that repeats a type in `wanted_message_types` would otherwise + /// process every message of that type as many times as it listed it. + #[tokio::test] + async fn duplicate_requested_types_deliver_once() { + let net = broker().await; + let mut rx = net.subscribe(&[MessageType::Inv, MessageType::Inv, MessageType::Inv]).await; - pub(crate) async fn test_capability_rejected_count(&self) -> usize { - self.capability_rejected.read().await.len() - } + net.dispatch_local(NetworkMessage::Inv(vec![Inventory::Block(BlockHash::dummy(1))])).await; - pub(crate) async fn test_is_capability_rejected(&self, addr: &SocketAddr) -> bool { - self.is_capability_rejected(addr).await + tokio::time::timeout(Duration::from_secs(2), rx.recv()) + .await + .expect("the message must arrive") + .expect("channel open"); + tokio::time::sleep(Duration::from_millis(50)).await; + assert!(rx.try_recv().is_err(), "a repeated type must not duplicate delivery"); } - pub(crate) async fn test_has_capable_peer(&self) -> bool { - self.required_services != ServiceFlags::NONE - && self.pool.has_peers_with_service(self.required_services).await + /// A manager task that ended drops its receiver; the pump prunes it and the + /// surviving subscribers keep working. + #[tokio::test] + async fn a_dropped_subscriber_does_not_break_the_others() { + let net = broker().await; + let dead = net.subscribe(&[MessageType::Inv]).await; + let mut alive = net.subscribe(&[MessageType::Inv]).await; + drop(dead); + + net.dispatch_local(NetworkMessage::Inv(vec![Inventory::Block(BlockHash::dummy(1))])).await; + + let (_, msg) = tokio::time::timeout(Duration::from_secs(2), alive.recv()) + .await + .expect("the live subscriber must still be served") + .expect("channel open"); + assert!(matches!(*msg, NetworkMessage::Inv(_))); } - pub(crate) async fn test_should_reject_after_handshake(&self, peer: &Peer) -> bool { - Self::should_reject_after_handshake(&self.pool, peer, self.required_services).await + /// Routing keys off the wire command: an unknown command has no type and is + /// therefore delivered to nobody. + #[test] + fn message_type_from_cmd_maps_known_commands_and_rejects_the_rest() { + assert_eq!(MessageType::from_cmd("headers"), Some(MessageType::Headers)); + assert_eq!(MessageType::from_cmd("inv"), Some(MessageType::Inv)); + assert_eq!(MessageType::from_cmd("cfilter"), Some(MessageType::CFilter)); + assert_eq!(MessageType::from_cmd("cfheaders"), Some(MessageType::CfHeaders)); + assert_eq!(MessageType::from_cmd("block"), Some(MessageType::Block)); + assert_eq!(MessageType::from_cmd("mnlistdiff"), Some(MessageType::MnListDiff)); + assert_eq!(MessageType::from_cmd("qrinfo"), Some(MessageType::QrInfo)); + assert_eq!(MessageType::from_cmd("tx"), Some(MessageType::Tx)); + assert_eq!(MessageType::from_cmd("isdlock"), Some(MessageType::IsDLock)); + assert_eq!(MessageType::from_cmd("clsig"), Some(MessageType::ChainLock)); + + assert_eq!(MessageType::from_cmd("ping"), None); + assert_eq!(MessageType::from_cmd(""), None); + assert_eq!(MessageType::from_cmd("Headers"), None, "the match is case-sensitive"); } } diff --git a/dash-spv/src/network/message_dispatcher.rs b/dash-spv/src/network/message_dispatcher.rs deleted file mode 100644 index d88540d7d..000000000 --- a/dash-spv/src/network/message_dispatcher.rs +++ /dev/null @@ -1,227 +0,0 @@ -//! Message dispatcher for network message distribution. -//! -//! This module filters incoming network messages by type and forwards -//! them to registered receivers. -//! -//! - [`Message`]: Wraps a `NetworkMessage` with the originating peer address -//! - [`MessageDispatcher`]: Manages channels and dispatches messages to interested parties - -use std::collections::{HashMap, HashSet}; -use std::net::SocketAddr; - -use dashcore::network::message::NetworkMessage; -use tokio::sync::mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender}; - -use crate::network::MessageType; - -/// A network message tagged with the peer that sent it. -#[derive(Clone, Debug, PartialEq)] -pub struct Message { - peer_address: SocketAddr, - inner: NetworkMessage, -} - -impl Message { - /// Creates a new Message from a peer address and network message. - pub fn new(peer_address: SocketAddr, inner: NetworkMessage) -> Self { - Self { - peer_address, - inner, - } - } - - /// Forwards the cmd() of the underlying NetworkMessage. - pub fn cmd(&self) -> &'static str { - self.inner.cmd() - } - - /// Returns the SocketAddr of the peer that sent this message. - pub fn peer_address(&self) -> SocketAddr { - self.peer_address - } - - /// Returns a reference to the underlying network message. - pub fn inner(&self) -> &NetworkMessage { - &self.inner - } -} - -/// Routes incoming network messages to subscribers based on message type. -/// -/// Subscribers call [`message_receiver`](Self::message_receiver) with the message types they -/// want, receiving an unbounded channel. When [`dispatch`](Self::dispatch) is called, the -/// message is sent to all subscribers registered for that type. Dead channels are pruned -/// automatically on dispatch. -#[derive(Debug, Default)] -pub struct MessageDispatcher { - senders: HashMap>>, -} - -impl MessageDispatcher { - /// Creates and returns a receiver that yields only messages matching the provided message types. - pub fn message_receiver( - &mut self, - message_types: &[MessageType], - ) -> UnboundedReceiver { - let (sender, receiver) = unbounded_channel(); - let unique_types: HashSet = message_types.iter().copied().collect(); - for message_type in unique_types { - self.senders.entry(message_type).or_default().push(sender.clone()); - } - receiver - } - - /// Distributes a message to all subscribers interested in its type. Prunes dead senders automatically. - pub fn dispatch(&mut self, message: &Message) { - let message_type = MessageType::from(message); - if let Some(senders) = self.senders.get_mut(&message_type) { - senders.retain(|sender| sender.send(message.clone()).is_ok()); - }; - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::test_utils::test_socket_address; - - #[test] - fn test_message_creation() { - let peer_address = test_socket_address(1); - let inner = NetworkMessage::Headers(vec![]); - let msg = Message::new(peer_address, inner.clone()); - - assert_eq!(msg.cmd(), inner.cmd()); - assert_eq!(msg.peer_address(), peer_address); - assert_eq!(*msg.inner(), inner); - } - - #[tokio::test] - async fn test_dispatch_to_interested_receiver() { - let mut message_dispatcher = MessageDispatcher::default(); - let mut receiver = message_dispatcher.message_receiver(&[MessageType::Inv]); - - let msg = Message::new(test_socket_address(1), NetworkMessage::Inv(vec![])); - message_dispatcher.dispatch(&msg); - - assert_eq!(receiver.recv().await.unwrap(), msg); - } - - #[tokio::test] - async fn test_dispatch_skips_uninterested() { - let mut message_dispatcher = MessageDispatcher::default(); - let mut receiver = message_dispatcher.message_receiver(&[MessageType::Headers]); - - let msg = Message::new(test_socket_address(1), NetworkMessage::Inv(vec![])); - message_dispatcher.dispatch(&msg); - - assert!(receiver.try_recv().is_err()); - } - - #[test] - fn test_dispatch_no_subscribers() { - let mut message_dispatcher = MessageDispatcher::default(); - let msg = Message::new(test_socket_address(1), NetworkMessage::Headers(vec![])); - - // Should not panic with no subscribers - message_dispatcher.dispatch(&msg); - } - - #[tokio::test] - async fn test_message_receiver_multiple_types() { - let mut message_dispatcher = MessageDispatcher::default(); - let mut receiver1 = message_dispatcher.message_receiver(&[MessageType::Headers]); - let mut receiver2 = message_dispatcher.message_receiver(&[MessageType::Inv]); - - let headers_msg = Message::new(test_socket_address(1), NetworkMessage::Headers(vec![])); - let inv_msg = Message::new(test_socket_address(2), NetworkMessage::Inv(vec![])); - - message_dispatcher.dispatch(&headers_msg); - message_dispatcher.dispatch(&inv_msg); - - assert_eq!(receiver1.recv().await.unwrap(), headers_msg); - assert_eq!(receiver2.recv().await.unwrap(), inv_msg); - } - - #[tokio::test] - async fn test_dispatch_multiple_subscribers_same_type() { - let mut message_dispatcher = MessageDispatcher::default(); - let mut receiver1 = message_dispatcher.message_receiver(&[MessageType::Headers]); - let mut receiver2 = message_dispatcher.message_receiver(&[MessageType::Headers]); - - let msg = Message::new(test_socket_address(1), NetworkMessage::Headers(vec![])); - message_dispatcher.dispatch(&msg); - - assert_eq!(receiver1.recv().await.unwrap(), msg); - assert_eq!(receiver2.recv().await.unwrap(), msg); - } - - #[tokio::test] - async fn test_dropped_receiver_does_not_affect_others() { - let mut message_dispatcher = MessageDispatcher::default(); - let receiver1 = message_dispatcher.message_receiver(&[MessageType::Headers]); - let mut receiver2 = message_dispatcher.message_receiver(&[MessageType::Headers]); - - drop(receiver1); - - let msg = Message::new(test_socket_address(1), NetworkMessage::Headers(vec![])); - message_dispatcher.dispatch(&msg); - - assert_eq!(receiver2.recv().await.unwrap(), msg); - } - - #[tokio::test] - async fn test_messages_received_in_order() { - let mut message_dispatcher = MessageDispatcher::default(); - let mut receiver = - message_dispatcher.message_receiver(&[MessageType::Headers, MessageType::Inv]); - - let peer_1 = test_socket_address(1); - let peer_2 = test_socket_address(2); - let peer_3 = test_socket_address(3); - - let msg1 = Message::new(peer_1, NetworkMessage::Headers(vec![])); - let msg2 = Message::new(peer_2, NetworkMessage::Inv(vec![])); - let msg3 = Message::new(peer_3, NetworkMessage::Headers(vec![])); - - message_dispatcher.dispatch(&msg1); - message_dispatcher.dispatch(&msg2); - message_dispatcher.dispatch(&msg3); - - assert_eq!(receiver.recv().await.unwrap().peer_address(), peer_1); - assert_eq!(receiver.recv().await.unwrap().peer_address(), peer_2); - assert_eq!(receiver.recv().await.unwrap().peer_address(), peer_3); - } - - #[tokio::test] - async fn test_message_receiver_receives_multiple_types() { - let mut message_dispatcher = MessageDispatcher::default(); - let mut receiver = - message_dispatcher.message_receiver(&[MessageType::Headers, MessageType::Inv]); - - let headers_msg = Message::new(test_socket_address(1), NetworkMessage::Headers(vec![])); - let inv_msg = Message::new(test_socket_address(2), NetworkMessage::Inv(vec![])); - - message_dispatcher.dispatch(&headers_msg); - message_dispatcher.dispatch(&inv_msg); - - assert_eq!(receiver.recv().await.unwrap(), headers_msg); - assert_eq!(receiver.recv().await.unwrap(), inv_msg); - } - - #[tokio::test] - async fn test_duplicate_message_types_no_duplicate_delivery() { - let mut message_dispatcher = MessageDispatcher::default(); - let mut receiver = message_dispatcher.message_receiver(&[ - MessageType::Headers, - MessageType::Headers, - MessageType::Headers, - ]); - - let msg = Message::new(test_socket_address(1), NetworkMessage::Headers(vec![])); - message_dispatcher.dispatch(&msg); - - assert_eq!(receiver.recv().await.unwrap(), msg); - assert!(receiver.try_recv().is_err()); - } -} diff --git a/dash-spv/src/network/message_type.rs b/dash-spv/src/network/message_type.rs deleted file mode 100644 index 9d2dfb86d..000000000 --- a/dash-spv/src/network/message_type.rs +++ /dev/null @@ -1,170 +0,0 @@ -//! Message type enum for easier message mapping to NetworkMessage variants. -//! -//! Uses a macro to keep MessageType in sync with NetworkMessage. -//! If NetworkMessage adds a new variant, compilation will fail until -//! the variant is added here. - -use crate::network::Message; -use dashcore::network::message::NetworkMessage; - -/// Generates the `MessageType` enum -/// -/// Implements: -/// - `From<&Message>` -/// -/// Each `NetworkMessage` variant maps to a corresponding `MessageType` variant -/// (e.g., `NetworkMessage::Headers(_)` -> `MessageType::Headers`). -/// -/// Syntax for entries: -/// - `Name` for unit variants (e.g., `Verack`) -/// - `Name (..)` for tuple variants with data (e.g., `Headers (..)`) -/// - `Name { .. }` for struct variants (e.g., `Unknown { .. }`) -macro_rules! define_message_types { - ($($(#[$meta:meta])* $variant:ident $( ( $($tuple:tt)* ) )? $( { $($field:tt)* } )?),* $(,)?) => { - /// Message types that subscribers can subscribe to. - #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] - pub enum MessageType { - $($(#[$meta])* $variant,)* - } - - impl From<&Message> for MessageType { - fn from(value: &Message) -> Self { - match value.inner() { - $(NetworkMessage::$variant $( ( $($tuple)* ) )? $( { $($field)* } )? => MessageType::$variant,)* - } - } - } - }; -} - -define_message_types! { - /// `version` - Version (..), - /// `verack` - Verack, - /// `addr` - Addr (..), - /// `inv` - Inv (..), - /// `getdata` - GetData (..), - /// `notfound` - NotFound (..), - /// `getblocks` - GetBlocks (..), - /// `getheaders` - GetHeaders (..), - /// `mempool` - MemPool, - /// `tx` - Tx (..), - /// `block` - Block (..), - /// `headers` - Headers (..), - /// `sendheaders` - SendHeaders, - /// `getheaders2` - GetHeaders2 (..), - /// `sendheaders2` - SendHeaders2, - /// `headers2` - Headers2 (..), - /// `getaddr` - GetAddr, - /// `ping` - Ping (..), - /// `pong` - Pong (..), - /// `merkleblock` - MerkleBlock (..), - /// `filterload` - FilterLoad (..), - /// `filteradd` - FilterAdd (..), - /// `filterclear` - FilterClear, - /// `getcfilters` - GetCFilters (..), - /// `cfilter` - CFilter (..), - /// `getcfheaders` - GetCFHeaders (..), - /// `cfheaders` - CFHeaders (..), - /// `getcfcheckpt` - GetCFCheckpt (..), - /// `cfcheckpt` - CFCheckpt (..), - /// `sendcmpct` - SendCmpct (..), - /// `cmpctblock` - CmpctBlock (..), - /// `getblocktxn` - GetBlockTxn (..), - /// `blocktxn` - BlockTxn (..), - /// `alert` - Alert (..), - /// `reject` - Reject (..), - /// `feefilter` - FeeFilter (..), - /// `wtxidrelay` - WtxidRelay, - /// `addrv2` - AddrV2 (..), - /// `sendaddrv2` - SendAddrV2, - /// `getmnlistd` - GetMnListD (..), - /// `mnlistdiff` - MnListDiff (..), - /// `getqrinfo` - GetQRInfo (..), - /// `qrinfo` - QRInfo (..), - /// `clsig` - CLSig (..), - /// `isdlock` - ISLock (..), - /// `senddsq` - SendDsq (..), - /// Unknown message type - Unknown { .. }, -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::test_utils::test_socket_address; - - #[test] - fn from_message_unit_variant() { - let addr = test_socket_address(1); - - let msg = Message::new(addr, NetworkMessage::SendHeaders); - assert_eq!(MessageType::from(&msg), MessageType::SendHeaders); - } - - #[test] - fn from_message_tuple_variant() { - let addr = test_socket_address(1); - - let msg = Message::new(addr, NetworkMessage::Alert(vec![])); - assert_eq!(MessageType::from(&msg), MessageType::Alert); - } - - #[test] - fn from_message_unknown_variant() { - use dashcore::network::message::CommandString; - - let addr = test_socket_address(1); - let unknown_msg = NetworkMessage::Unknown { - command: CommandString::try_from_static("test").unwrap(), - payload: vec![], - }; - let msg = Message::new(addr, unknown_msg); - assert_eq!(MessageType::from(&msg), MessageType::Unknown); - } -} diff --git a/dash-spv/src/network/mod.rs b/dash-spv/src/network/mod.rs index 20c4266c7..f263a3768 100644 --- a/dash-spv/src/network/mod.rs +++ b/dash-spv/src/network/mod.rs @@ -1,272 +1,121 @@ -//! Network layer for the Dash SPV client. - -pub mod addrv2; -pub mod constants; -pub mod discovery; -mod event; -pub mod handshake; -pub mod manager; -mod message_dispatcher; -pub mod peer; -pub mod pool; -mod reputation; - -mod message_type; -#[cfg(test)] -mod tests; - -pub use event::NetworkEvent; +mod discovery; +mod manager; +mod peer; use async_trait::async_trait; -use tokio::sync::{broadcast, mpsc}; - -use crate::error::NetworkResult; -use crate::NetworkError; use dashcore::network::message::NetworkMessage; -use dashcore::network::message_blockdata::{GetHeadersMessage, Inventory}; -use dashcore::network::message_bloom::FilterLoad; -use dashcore::network::message_filter::{GetCFHeaders, GetCFilters}; -use dashcore::network::message_qrinfo::GetQRInfo; -use dashcore::network::message_sml::GetMnListDiff; -use dashcore::BlockHash; -use dashcore_hashes::Hash; -pub use handshake::{HandshakeManager, HandshakeState}; -pub use manager::PeerNetworkManager; -pub use message_dispatcher::{Message, MessageDispatcher}; -pub use message_type::MessageType; -pub use peer::Peer; -pub(crate) use reputation::PeerReputation; use std::net::SocketAddr; +use tokio::sync::broadcast; use tokio::sync::mpsc::UnboundedReceiver; -const FILTER_TYPE_DEFAULT: u8 = 0; +// `NetworkEvent` is part of the public `EventHandler` API (delivered to +// `on_network_event`), so it stays exported. The peer-to-peer manager and its +// request/message plumbing are internal: the client builds and owns the manager +// itself (see `DashSpvClient::new`), so none of these are part of the public API. +pub use manager::NetworkEvent; +// These form the `NetworkManager` trait's interface, which appears in the public +// `DashSpvClient` bound, so they are part of the public API. +pub use manager::{Inbound, MessageType, PeerNetworkManager, RequestKey}; + +/// Abstraction over the peer-to-peer network manager. +/// +/// The sync managers and pipelines depend on this trait rather than the concrete +/// [`PeerNetworkManager`], so a lightweight mock can drive them in unit tests and +/// the client stays generic over the network implementation. It is the *minimum* +/// surface the sync layer needs: declare requests, subscribe to inbound messages +/// and peer events, correlate answered requests, and lifecycle. +/// +/// Requests are fire-and-forget: the implementation de-duplicates by request key, +/// paces, times out and retries. Once a response is correlated to a request, the +/// caller reports it via [`request_answered`](Self::request_answered) (and +/// [`request_completed`](Self::request_completed) for streaming batches) so the +/// implementation stops tracking it. +#[async_trait] +pub trait NetworkManager: Send + Sync + 'static { + /// Begin peer discovery/connection. Call *after* every consumer has + /// subscribed, so the initial `PeersUpdated`/`PeerConnected` events are seen. + fn start(&self); -/// Request to send to network. -#[derive(Debug)] -pub enum NetworkRequest { - /// Send a message to the network. - SendMessage(NetworkMessage), - /// Send a message to a specific peer. - SendMessageToPeer(NetworkMessage, SocketAddr), - /// Broadcast a message to all connected peers. - BroadcastMessage(NetworkMessage), -} + /// Tear down all peer connections and background tasks. + fn stop(&self); -/// Handle for managers to queue outgoing network requests. -#[derive(Clone)] -pub struct RequestSender { - tx: mpsc::UnboundedSender, -} + /// Declare a request/message. Keyed requests are de-duplicated, paced, + /// timed out and retried by the implementation. + async fn send(&self, msg: NetworkMessage); -impl RequestSender { - /// Create a new RequestSender. - pub fn new(tx: mpsc::UnboundedSender) -> Self { - Self { - tx, - } - } + /// Send a message to one specific peer. Returns `false` if the peer is not + /// connected or the write failed. + async fn send_to(&self, addr: SocketAddr, msg: NetworkMessage) -> bool; - /// Queue a message to be sent to the network. - fn send_message(&self, msg: NetworkMessage) -> NetworkResult<()> { - self.tx - .send(NetworkRequest::SendMessage(msg)) - .map_err(|e| NetworkError::ProtocolError(e.to_string())) - } + /// Fire-and-forget send to every connected peer. + fn broadcast(&self, msg: NetworkMessage); - /// Queue a message to be sent to a specific peer. - fn send_message_to_peer( - &self, - msg: NetworkMessage, - peer_address: SocketAddr, - ) -> NetworkResult<()> { - self.tx - .send(NetworkRequest::SendMessageToPeer(msg, peer_address)) - .map_err(|e| NetworkError::ProtocolError(e.to_string())) - } + /// Inject a message into the local pump as if received from a peer. + async fn dispatch_local(&self, msg: NetworkMessage); - /// Queue a message to be broadcast to all connected peers. - pub(crate) fn broadcast(&self, msg: NetworkMessage) -> NetworkResult<()> { - self.tx - .send(NetworkRequest::BroadcastMessage(msg)) - .map_err(|e| NetworkError::ProtocolError(e.to_string())) - } + /// Report that a request key has been answered so it stops being tracked + /// for timeout/retry. + async fn request_answered(&self, key: RequestKey); - /// Send a transaction to a specific peer. - pub(crate) fn send_transaction( - &self, - tx: dashcore::Transaction, - peer_address: SocketAddr, - ) -> NetworkResult<()> { - self.send_message_to_peer(NetworkMessage::Tx(tx), peer_address) - } + /// Report that `n` streaming requests served by `peer` fully completed, + /// freeing that peer's in-flight units. + async fn request_completed(&self, peer: SocketAddr, n: usize); - /// Request inventory from a specific peer. - pub fn request_inventory( - &self, - inventory: Vec, - peer_address: SocketAddr, - ) -> NetworkResult<()> { - self.send_message_to_peer(NetworkMessage::GetData(inventory), peer_address) - } + /// Subscribe to inbound messages of the given types. Each inbound item is a + /// `(peer, message)` pair. + async fn subscribe(&self, kinds: &[MessageType]) -> UnboundedReceiver; - pub fn request_block_headers(&self, start_hash: BlockHash) -> NetworkResult<()> { - self.send_message(NetworkMessage::GetHeaders(GetHeadersMessage::new( - vec![start_hash], - BlockHash::all_zeros(), - ))) - } + /// Subscribe to peer-set lifecycle events. + fn events(&self) -> broadcast::Receiver; - pub fn request_block_headers_from_peer( - &self, - start_hash: BlockHash, - address: SocketAddr, - ) -> NetworkResult<()> { - self.send_message_to_peer( - NetworkMessage::GetHeaders(GetHeadersMessage::new( - vec![start_hash], - BlockHash::all_zeros(), - )), - address, - ) - } + /// Best tip height advertised across connected peers. + fn tip(&self) -> u32; - pub fn request_filter_headers( - &self, - start_height: u32, - stop_hash: BlockHash, - ) -> NetworkResult<()> { - self.send_message(NetworkMessage::GetCFHeaders(GetCFHeaders { - filter_type: FILTER_TYPE_DEFAULT, - start_height, - stop_hash, - })) - } + /// Number of currently connected peers. + async fn connected_count(&self) -> u32; +} - pub fn request_filters(&self, start_height: u32, stop_hash: BlockHash) -> NetworkResult<()> { - self.send_message(NetworkMessage::GetCFilters(GetCFilters { - filter_type: FILTER_TYPE_DEFAULT, - start_height, - stop_hash, - })) +// TODO: Inline the methods +// Thin delegation to the inherent methods of `PeerNetworkManager`. Method-call +// syntax (`self.method(..)`) resolves to the inherent method (inherent methods +// take precedence over trait methods), so this does not recurse and adds no +// behaviour of its own. +#[async_trait] +impl NetworkManager for PeerNetworkManager { + fn start(&self) { + self.start() } - - pub fn request_mnlist_diff( - &self, - base_block_hash: BlockHash, - block_hash: BlockHash, - ) -> NetworkResult<()> { - self.send_message(NetworkMessage::GetMnListD(GetMnListDiff { - base_block_hash, - block_hash, - })) + fn stop(&self) { + self.stop() } - - pub fn request_qr_info( - &self, - known_block_hashes: Vec, - target_block_hash: BlockHash, - extra_share: bool, - ) -> NetworkResult<()> { - self.send_message(NetworkMessage::GetQRInfo(GetQRInfo { - base_block_hashes: known_block_hashes, - block_request_hash: target_block_hash, - extra_share, - })) + async fn send(&self, msg: NetworkMessage) { + self.send(msg).await } - - pub fn request_blocks(&self, hashes: Vec) -> NetworkResult<()> { - self.send_message(NetworkMessage::GetData( - hashes.into_iter().map(Inventory::Block).collect(), - )) + async fn send_to(&self, addr: SocketAddr, msg: NetworkMessage) -> bool { + self.send_to(addr, msg).await } - - /// Send a filterload message to a specific peer. - pub fn send_filter_load(&self, filter_load: FilterLoad, peer: SocketAddr) -> NetworkResult<()> { - self.send_message_to_peer(NetworkMessage::FilterLoad(filter_load), peer) + fn broadcast(&self, msg: NetworkMessage) { + self.broadcast(msg) } - - /// Send a filterclear message to a specific peer. - pub fn send_filter_clear(&self, peer: SocketAddr) -> NetworkResult<()> { - self.send_message_to_peer(NetworkMessage::FilterClear, peer) + async fn dispatch_local(&self, msg: NetworkMessage) { + self.dispatch_local(msg).await } - - /// Send a mempool message to request inventory from a specific peer. - pub fn request_mempool(&self, peer: SocketAddr) -> NetworkResult<()> { - self.send_message_to_peer(NetworkMessage::MemPool, peer) + async fn request_answered(&self, key: RequestKey) { + self.request_answered(key).await } -} - -/// Network manager trait for abstracting network operations. -#[async_trait] -pub trait NetworkManager: Send + Sync + 'static { - /// Creates and returns a receiver that yields only messages of the matching the provided message types. - async fn message_receiver(&mut self, types: &[MessageType]) -> UnboundedReceiver; - - /// Get a sender for queuing outgoing network requests. - /// - /// Messages sent via this sender are delivered to the network asynchronously. - fn request_sender(&self) -> RequestSender; - - /// Connect to the network. - async fn connect(&mut self) -> NetworkResult<()>; - - /// Disconnect from the network. - async fn disconnect(&mut self) -> NetworkResult<()>; - - /// Send a message to a peer. - async fn send_message(&mut self, message: NetworkMessage) -> NetworkResult<()>; - - /// Get the number of connected peers. - fn peer_count(&self) -> usize; - - /// Request QRInfo from the network. - /// - /// # Arguments - /// * `base_block_hashes` - Array of base block hashes for the masternode lists the light client already knows - /// * `block_request_hash` - Hash of the block for which the masternode list diff is requested - /// * `extra_share` - Optional flag to indicate if an extra share is requested - async fn request_qr_info( - &mut self, - base_block_hashes: Vec, - block_request_hash: BlockHash, - extra_share: bool, - ) -> NetworkResult<()> { - use dashcore::network::message_qrinfo::GetQRInfo; - - let get_qr_info = GetQRInfo { - base_block_hashes: base_block_hashes.clone(), - block_request_hash, - extra_share, - }; - - let base_hashes_count = get_qr_info.base_block_hashes.len(); - - self.send_message(NetworkMessage::GetQRInfo(get_qr_info)).await?; - - tracing::debug!( - "Requested QRInfo with {} base hashes for block {}, extra_share={}", - base_hashes_count, - block_request_hash, - extra_share - ); - - Ok(()) + async fn request_completed(&self, peer: SocketAddr, n: usize) { + self.request_completed(peer, n).await + } + async fn subscribe(&self, kinds: &[MessageType]) -> UnboundedReceiver { + self.subscribe(kinds).await + } + fn events(&self) -> broadcast::Receiver { + self.events() + } + fn tip(&self) -> u32 { + self.tip() + } + async fn connected_count(&self) -> u32 { + self.connected_count().await } - - /// Broadcast a message to all connected peers. - async fn broadcast(&self, _message: NetworkMessage) -> NetworkResult<()>; - - /// Inject a message into the local message dispatcher as if received from a peer. - /// - /// Used for locally-originated messages (e.g., self-broadcast transactions) that - /// should be processed through the same pipeline as peer-received messages. - async fn dispatch_local(&self, message: NetworkMessage); - - /// Disconnect a specific peer by address. - async fn disconnect_peer(&self, _addr: &SocketAddr, _reason: &str) -> NetworkResult<()>; - - /// Subscribe to network events (peer connections, disconnections). - /// - /// Returns a broadcast receiver for network events. - fn subscribe_network_events(&self) -> broadcast::Receiver; } diff --git a/dash-spv/src/network/peer.rs b/dash-spv/src/network/peer.rs index 48e2f631c..34ffce4d3 100644 --- a/dash-spv/src/network/peer.rs +++ b/dash-spv/src/network/peer.rs @@ -1,850 +1,568 @@ -//! Dash peer connection management. - -use dashcore::network::constants::ServiceFlags; -use std::collections::HashMap; +use dashcore::{ + consensus::encode, + network::{ + address::Address, + constants::ServiceFlags, + message::{NetworkMessage, RawNetworkMessage, RawNetworkMessageCodec}, + message_network::VersionMessage, + }, + Network, +}; +use futures::lock::Mutex; +use std::collections::VecDeque; use std::net::SocketAddr; +use std::sync::atomic::{AtomicU32, AtomicU64, AtomicUsize, Ordering}; use std::sync::Arc; -use std::time::{Duration, SystemTime}; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; -use tokio::net::TcpStream; -use tokio::sync::Mutex; - -use dashcore::consensus::{encode, Decodable}; -use dashcore::network::message::{NetworkMessage, RawNetworkMessage}; -use dashcore::Network; - -use crate::error::{NetworkError, NetworkResult}; -use crate::network::constants::PING_INTERVAL; -use crate::network::Message; - -/// Internal state for the TCP connection -struct ConnectionState { - stream: TcpStream, - // Stateful message framing buffer to ensure full frames before decoding - framing_buffer: Vec, -} - -/// Dash P2P peer -pub struct Peer { - address: SocketAddr, - // Use a single mutex to protect both the write stream and read buffer - // This ensures no concurrent access to the underlying socket - state: Option>>, - timeout: Duration, - connected_at: Option, - bytes_sent: u64, - network: Network, - // Ping/pong state - last_ping_sent: Option, - last_pong_received: Option, - pending_pings: HashMap, // nonce -> sent_time - // Peer information from Version message - version: Option, - services: Option, - user_agent: Option, - best_height: Option, - relay: Option, - prefers_headers2: bool, - sent_sendheaders2: bool, +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; +use tokio::sync::mpsc::UnboundedSender; +use tokio::{ + io::{AsyncRead, AsyncWriteExt, ReadBuf}, + net::{ + tcp::{OwnedReadHalf, OwnedWriteHalf}, + TcpStream, + }, +}; +use tokio_stream::StreamExt; +use tokio_util::codec::FramedRead; +use tokio_util::sync::CancellationToken; + +use crate::{error::NetworkResult, NetworkError}; + +const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(5); +const USER_AGENT: &str = concat!("/dash-spv:", env!("CARGO_PKG_VERSION"), "/"); + +/// Per-peer response-latency tracker: the time each pipeline request spends in +/// flight, from send to the response that completes it. Send times are queued +/// FIFO; each completing response pops the oldest and folds the elapsed time +/// into running total/count/max (so we can report avg and worst per peer). +#[derive(Default)] +pub(crate) struct Latency { + pending: Mutex>, + total_ns: AtomicU64, + count: AtomicU64, + max_ns: AtomicU64, } -impl Peer { - /// Get the remote peer socket address. - pub fn address(&self) -> SocketAddr { - self.address - } - /// Create a new peer. - pub fn new(address: SocketAddr, timeout: Duration, network: Network) -> Self { - Self { - address, - state: None, - timeout, - connected_at: None, - bytes_sent: 0, - network, - last_ping_sent: None, - last_pong_received: None, - pending_pings: HashMap::new(), - version: None, - services: None, - user_agent: None, - best_height: None, - relay: None, - prefers_headers2: false, - sent_sendheaders2: false, - } +impl Latency { + async fn on_send(&self) { + self.pending.lock().await.push_back(Instant::now()); } - /// Connect to a peer and return a connected instance. - pub async fn connect( - address: SocketAddr, - timeout_secs: u64, - network: Network, - ) -> NetworkResult { - let timeout = Duration::from_secs(timeout_secs); - - let stream = tokio::time::timeout(timeout, TcpStream::connect(address)) - .await - .map_err(|_| { - NetworkError::ConnectionFailed(format!("Connection to {} timed out", address)) - })? - .map_err(|e| { - NetworkError::ConnectionFailed(format!("Failed to connect to {}: {}", address, e)) - })?; - - stream.set_nodelay(true).map_err(|e| { - NetworkError::ConnectionFailed(format!("Failed to set TCP_NODELAY: {}", e)) - })?; - - let state = ConnectionState { - stream, - framing_buffer: Vec::new(), - }; - - Ok(Self { - address, - state: Some(Arc::new(Mutex::new(state))), - timeout, - connected_at: Some(SystemTime::now()), - bytes_sent: 0, - network, - last_ping_sent: None, - last_pong_received: None, - pending_pings: HashMap::new(), - version: None, - services: None, - user_agent: None, - best_height: None, - relay: None, - prefers_headers2: false, - sent_sendheaders2: false, - }) - } - - pub fn version(&self) -> Option { - self.version + /// Undo an `on_send` whose write then failed, so the queue does not drift. + async fn cancel_one(&self) { + self.pending.lock().await.pop_back(); } - pub fn best_height(&self) -> Option { - self.best_height + /// Pop the oldest pending send and record its round-trip. + async fn complete_one(&self) { + let sent = self.pending.lock().await.pop_front(); + if let Some(sent) = sent { + let ns = sent.elapsed().as_nanos() as u64; + self.total_ns.fetch_add(ns, Ordering::Relaxed); + self.count.fetch_add(1, Ordering::Relaxed); + self.max_ns.fetch_max(ns, Ordering::Relaxed); + } } - /// Check if peer supports compact filters (BIP 157/158). - pub fn supports_compact_filters(&self) -> bool { - self.has_service(ServiceFlags::COMPACT_FILTERS) + /// (completed request count, average ms, worst ms). + fn snapshot(&self) -> (u64, f64, f64) { + let count = self.count.load(Ordering::Relaxed); + let total = self.total_ns.load(Ordering::Relaxed); + let max = self.max_ns.load(Ordering::Relaxed); + let avg_ms = if count > 0 { + total as f64 / count as f64 / 1e6 + } else { + 0.0 + }; + (count, avg_ms, max as f64 / 1e6) } - /// Check if peer supports headers2 compression (DIP-0025). - pub fn supports_headers2(&self) -> bool { - self.has_service(ServiceFlags::NODE_HEADERS_COMPRESSED) + /// Cumulative (completed request count, total service-time nanoseconds). + /// The bandwidth controller diffs these across a window to get THIS peer's + /// completion rate and service time, sizing its in-flight cap by Little's Law. + fn totals(&self) -> (u64, u64) { + (self.count.load(Ordering::Relaxed), self.total_ns.load(Ordering::Relaxed)) } +} - pub fn has_service(&self, flags: ServiceFlags) -> bool { - self.services.map(|s| ServiceFlags::from(s).has(flags)).unwrap_or(false) - } +/// Wraps a socket read half and adds every byte read into a shared counter, so +/// the network manager can estimate download throughput (and size the global +/// in-flight budget to ~90% of it). +struct CountingReader { + inner: R, + /// Host-wide download counter (all peers), the global bandwidth signal. + bytes: Arc, + /// This connection's own download counter, so the controller can measure + /// per-peer throughput. + peer_bytes: Arc, +} - pub(crate) fn services_known(&self) -> bool { - self.services.is_some() +impl AsyncRead for CountingReader { + fn poll_read( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> std::task::Poll> { + let before = buf.filled().len(); + let r = std::pin::Pin::new(&mut self.inner).poll_read(cx, buf); + if let std::task::Poll::Ready(Ok(())) = &r { + let n = (buf.filled().len() - before) as u64; + self.bytes.fetch_add(n, Ordering::Relaxed); + self.peer_bytes.fetch_add(n, Ordering::Relaxed); + } + r } +} - /// Connect to the peer (instance method for compatibility). - pub async fn connect_instance(&mut self) -> NetworkResult<()> { - let stream = tokio::time::timeout(self.timeout, TcpStream::connect(self.address)) - .await - .map_err(|_| { - NetworkError::ConnectionFailed(format!("Connection to {} timed out", self.address)) - })? - .map_err(|e| { - NetworkError::ConnectionFailed(format!( - "Failed to connect to {}: {}", - self.address, e - )) - })?; - - // Disable Nagle's algorithm for lower latency - stream.set_nodelay(true).map_err(|e| { - NetworkError::ConnectionFailed(format!("Failed to set TCP_NODELAY: {}", e)) - })?; +type PeerReader = FramedRead, RawNetworkMessageCodec>; - let state = ConnectionState { - stream, - framing_buffer: Vec::new(), - }; +// `NetworkMessage` is a large enum (unboxed to mirror upstream `dashcore`); wrapping it +// here by value trips `large_enum_variant`. The event lives briefly on the inbound channel +// and boxing every message would penalise the common small ones, so allow the size skew. +#[allow(clippy::large_enum_variant)] +pub enum PeerEvent { + Message(SocketAddr, NetworkMessage), + Disconnected(SocketAddr), +} - self.state = Some(Arc::new(Mutex::new(state))); - self.connected_at = Some(SystemTime::now()); +pub struct ConnectedPeer { + network: Network, + addr: SocketAddr, + version: VersionMessage, + lag_ms: AtomicU32, + in_flight: Arc, + latency: Arc, + writer: Arc>, + /// This peer's measured serving capacity — the number of requests it can have + /// in flight before ITS responses start queuing (its bandwidth-delay product). + /// Sized continuously by the bandwidth controller from the peer's own + /// completion rate and service time, mirroring how the global budget is sized + /// from our download rate. Fast peers earn a high cap, slow peers a low one. + cap: Arc, + /// Cumulative bytes downloaded from THIS peer, for per-peer throughput. + bytes: Arc, + /// Per-connection cancel token (child of the global shutdown). Cancelling it + /// stops this peer's reader and closes the socket. Used to drop peers we + /// probed but don't keep, so we only hold connections we actually use. + token: CancellationToken, +} - tracing::info!("Connected to peer {}", self.address); +pub struct DisconnectedPeer { + network: Network, + addr: SocketAddr, + /// Handshake ping measured the last time we were connected to this peer, if any. + /// Lets the supervisor rank backups by measured quality instead of treating every + /// disconnected address as an unknown. `None` for an address we have never probed. + lag_ms: Option, +} - Ok(()) +impl ConnectedPeer { + pub fn addr(&self) -> SocketAddr { + self.addr } - /// Disconnect from the peer. - pub async fn disconnect(&mut self) -> NetworkResult<()> { - if let Some(state_arc) = self.state.take() { - if let Ok(state_mutex) = Arc::try_unwrap(state_arc) { - let mut state = state_mutex.into_inner(); - let _ = state.stream.shutdown().await; - } - } - self.connected_at = None; - - tracing::info!("Disconnected from peer {}", self.address); - - Ok(()) + pub fn version(&self) -> &VersionMessage { + &self.version } - /// Update peer information from a received Version message - pub fn update_peer_info( - &mut self, - version_msg: &dashcore::network::message_network::VersionMessage, - ) { - // Define validation constants - const MIN_PROTOCOL_VERSION: u32 = 60001; // Minimum version that supports ping/pong - const MAX_PROTOCOL_VERSION: u32 = 100000; // Reasonable upper bound for protocol version - const MAX_USER_AGENT_LENGTH: usize = 256; // Maximum reasonable user agent length - const MAX_START_HEIGHT: i32 = 10_000_000; // Reasonable upper bound for block height - - // Validate protocol version - if version_msg.version < MIN_PROTOCOL_VERSION { - tracing::warn!( - "Peer {} reported protocol version {} below minimum {}, skipping update", - self.address, - version_msg.version, - MIN_PROTOCOL_VERSION - ); - return; - } - - if version_msg.version > MAX_PROTOCOL_VERSION { - tracing::warn!( - "Peer {} reported suspiciously high protocol version {}, skipping update", - self.address, - version_msg.version - ); - return; - } - - // Validate start height - if version_msg.start_height < 0 { - tracing::warn!( - "Peer {} reported negative start height {}, skipping update", - self.address, - version_msg.start_height - ); - return; - } - - if version_msg.start_height > MAX_START_HEIGHT { - tracing::warn!( - "Peer {} reported suspiciously high start height {}, skipping update", - self.address, - version_msg.start_height - ); - return; - } - - // Validate user agent - if version_msg.user_agent.is_empty() { - tracing::warn!("Peer {} provided empty user agent, skipping update", self.address); - return; - } - - if version_msg.user_agent.len() > MAX_USER_AGENT_LENGTH { - tracing::warn!( - "Peer {} provided excessively long user agent ({} bytes), skipping update", - self.address, - version_msg.user_agent.len() - ); - return; - } - - // Validate services - ensure they contain expected flags - let services = version_msg.services.as_u64(); - const KNOWN_SERVICE_FLAGS: u64 = 0x0000_0000_0000_1FFF; // All known service flags up to bit 12 - if services & !KNOWN_SERVICE_FLAGS != 0 { - tracing::warn!( - "Peer {} reported unknown service flags: 0x{:016x}, proceeding with caution", - self.address, - services - ); - // Note: We don't return here as unknown flags might be from newer versions - } - - // All validations passed, update peer info - self.version = Some(version_msg.version); - self.services = Some(version_msg.services.as_u64()); - self.user_agent = Some(version_msg.user_agent.clone()); - self.best_height = Some(version_msg.start_height as u32); - self.relay = Some(version_msg.relay); - - tracing::info!( - "Updated peer info for {}: height={}, version={}, services={:?}", - self.address, - version_msg.start_height, - version_msg.version, - version_msg.services - ); - - // Also log with standard logging for debugging - tracing::info!( - "PEER_INFO_DEBUG: Updated peer {} with height={}, version={}", - self.address, - version_msg.start_height, - version_msg.version - ); + /// Net in-flight to this peer: `+1` per message we send it, `-1` per message + /// we read from it (managed internally by `send` and the reader task). The + /// router reads this to send to the least-loaded peer. + pub fn in_flight(&self) -> usize { + self.in_flight.load(Ordering::Relaxed) } - /// Helper function to read some bytes into the framing buffer. - async fn read_some(state: &mut ConnectionState) -> std::io::Result { - let mut tmp = [0u8; 8192]; - match state.stream.read(&mut tmp).await { - Ok(0) => Ok(0), - Ok(n) => { - state.framing_buffer.extend_from_slice(&tmp[..n]); - Ok(n) - } - Err(e) => Err(e), + pub fn disconnect(self) -> DisconnectedPeer { + DisconnectedPeer { + network: self.network, + addr: self.addr, + // Carry the measured handshake ping (0 means unmeasured) so the supervisor + // can rank this address against others without re-probing it. + lag_ms: (self.lag_ms() > 0).then(|| self.lag_ms()), } } - /// Send a message to the peer. - pub async fn send_message(&mut self, message: NetworkMessage) -> NetworkResult<()> { - let state_arc = self - .state - .as_ref() - .ok_or_else(|| NetworkError::ConnectionFailed("Not connected".to_string()))?; - - let raw_message = RawNetworkMessage { + pub async fn send(&self, msg: &NetworkMessage) -> NetworkResult<()> { + // TODO: Take a reference to msg instead of cloning it + let raw = RawNetworkMessage { magic: self.network.magic(), - payload: message, + payload: msg.clone(), }; - - let serialized = encode::serialize(&raw_message); - - // Log details for debugging headers2 issues - if matches!( - raw_message.payload, - NetworkMessage::GetHeaders2(_) | NetworkMessage::GetHeaders(_) - ) { - let msg_type = match raw_message.payload { - NetworkMessage::GetHeaders2(_) => "GetHeaders2", - NetworkMessage::GetHeaders(_) => "GetHeaders", - _ => "Unknown", - }; - tracing::debug!( - "Sending {} raw bytes (len={}): {:02x?}", - msg_type, - serialized.len(), - &serialized[..std::cmp::min(100, serialized.len())] - ); + let serialized = encode::serialize(&raw); + + // A pipeline request counts as one unit of in-flight work for this peer, and + // it must be counted BEFORE the write: the peer can reply before this task is + // scheduled again, and the reader's decrement saturates at zero. Counting + // afterwards loses that decrement and leaks the slot for the connection's + // lifetime — with a starting cap of 2, one leak halves the peer's capacity. + let accounted = is_pipeline_request(msg); + if accounted { + self.in_flight.fetch_add(1, Ordering::Relaxed); + self.latency.on_send().await; } - // Lock the state for the entire write operation - let mut state = state_arc.lock().await; - - // Write with error handling - match state.stream.write_all(&serialized).await { - Ok(_) => { - // Flush to ensure data is sent immediately - if let Err(e) = state.stream.flush().await { - tracing::warn!("Failed to flush socket {}: {}", self.address, e); - } - self.bytes_sent += serialized.len() as u64; - tracing::trace!("Sent message to {}: {:?}", self.address, raw_message.payload); - Ok(()) - } - Err(e) => { - tracing::warn!("Disconnecting {} due to write error: {}", self.address, e); - // Drop the lock before clearing connection state - drop(state); - // Clear connection state on write error - self.state = None; - self.connected_at = None; - Err(NetworkError::ConnectionFailed(format!("Write failed: {}", e))) + if let Err(e) = self.writer.lock().await.write_all(&serialized).await { + // Nothing reached the peer, so no response will free this unit. + if accounted { + self.in_flight.fetch_sub(1, Ordering::Relaxed); + self.latency.cancel_one().await; } + tracing::warn!("Disconnecting {} due to write error: {}", self.addr, e); + return Err(NetworkError::ConnectionFailed(format!("Write failed: {}", e))); } + Ok(()) } - /// Receive a message from the peer. - pub async fn receive_message(&mut self) -> NetworkResult> { - // If the state was cleared e.g. by a write-path broken pipe, treat as disconnected - // so the reader loop handles it identically to a read-path EOF. - let state_arc = self.state.as_ref().ok_or(NetworkError::PeerDisconnected)?; - - // Lock the state for the entire read operation - // This ensures no concurrent access to the socket - let mut state = state_arc.lock().await; - - // Buffered, stateful framing - const HEADER_LEN: usize = 24; // magic[4] + cmd[12] + length[4] + checksum[4] - const MAX_RESYNC_STEPS_PER_CALL: usize = 64; - - let result = async { - let magic_bytes = self.network.magic().to_le_bytes(); - let mut resync_steps = 0usize; - - loop { - // Ensure header availability - if state.framing_buffer.len() < HEADER_LEN { - match Self::read_some(&mut state).await { - Ok(0) => { - tracing::info!("Peer {} closed connection (EOF)", self.address); - return Err(NetworkError::PeerDisconnected); - } - Ok(_) => {} - Err(ref e) - if e.kind() == std::io::ErrorKind::ConnectionAborted - || e.kind() == std::io::ErrorKind::ConnectionReset => - { - tracing::info!("Peer {} connection reset/aborted", self.address); - return Err(NetworkError::PeerDisconnected); - } - Err(e) => { - return Err(NetworkError::ConnectionFailed(format!( - "Read failed: {}", - e - ))); - } - } - } - - // Align to magic - if state.framing_buffer.len() >= 4 && state.framing_buffer[..4] != magic_bytes { - if let Some(pos) = - state.framing_buffer.windows(4).position(|w| w == magic_bytes) - { - if pos > 0 { - tracing::warn!( - "{}: stream desync: skipping {} stray bytes before magic", - self.address, - pos - ); - state.framing_buffer.drain(0..pos); - resync_steps += 1; - if resync_steps >= MAX_RESYNC_STEPS_PER_CALL { - return Ok(None); - } - continue; - } - } else { - // Keep last 3 bytes of potential magic prefix - if state.framing_buffer.len() > 3 { - let dropped = state.framing_buffer.len() - 3; - tracing::warn!( - "{}: stream desync: dropping {} bytes (no magic found)", - self.address, - dropped - ); - state.framing_buffer.drain(0..dropped); - resync_steps += 1; - if resync_steps >= MAX_RESYNC_STEPS_PER_CALL { - return Ok(None); - } - } - // Need more data - match Self::read_some(&mut state).await { - Ok(0) => { - tracing::info!("Peer {} closed connection (EOF)", self.address); - return Err(NetworkError::PeerDisconnected); - } - Ok(_) => {} - Err(e) => { - return Err(NetworkError::ConnectionFailed(format!( - "Read failed: {}", - e - ))); - } - } - continue; - } - } - - // Ensure full header - if state.framing_buffer.len() < HEADER_LEN { - match Self::read_some(&mut state).await { - Ok(0) => { - tracing::info!("Peer {} closed connection (EOF)", self.address); - return Err(NetworkError::PeerDisconnected); - } - Ok(_) => {} - Err(e) => { - return Err(NetworkError::ConnectionFailed(format!( - "Read failed: {}", - e - ))); - } - } - continue; - } - - // Parse header fields - let length_le = u32::from_le_bytes([ - state.framing_buffer[16], - state.framing_buffer[17], - state.framing_buffer[18], - state.framing_buffer[19], - ]) as usize; - let header_checksum = [ - state.framing_buffer[20], - state.framing_buffer[21], - state.framing_buffer[22], - state.framing_buffer[23], - ]; - // Validate announced length to prevent unbounded accumulation or overflow - if length_le > dashcore::network::message::MAX_MSG_SIZE { - return Err(NetworkError::ProtocolError(format!( - "Declared payload length {} exceeds MAX_MSG_SIZE {}", - length_le, - dashcore::network::message::MAX_MSG_SIZE - ))); - } - let total_len = match HEADER_LEN.checked_add(length_le) { - Some(v) => v, - None => { - return Err(NetworkError::ProtocolError( - "Message length overflow".to_string(), - )); - } - }; - - // Ensure full frame available - if state.framing_buffer.len() < total_len { - match Self::read_some(&mut state).await { - Ok(0) => { - tracing::info!("Peer {} closed connection (EOF)", self.address); - return Err(NetworkError::PeerDisconnected); - } - Ok(_) => {} - Err(e) => { - return Err(NetworkError::ConnectionFailed(format!( - "Read failed: {}", - e - ))); - } - } - continue; - } - - // Verify checksum - let payload_slice = &state.framing_buffer[HEADER_LEN..total_len]; - let expected = { - let checksum = ::hash( - payload_slice, - ); - [checksum[0], checksum[1], checksum[2], checksum[3]] - }; - if expected != header_checksum { - tracing::warn!( - "Skipping message with invalid checksum from {}: expected {:02x?}, actual {:02x?}", - self.address, - expected, - header_checksum - ); - if header_checksum == [0, 0, 0, 0] { - tracing::warn!( - "All-zeros checksum detected from {}, likely corrupted stream - resyncing", - self.address - ); - } - // Resync by dropping a byte and retrying - state.framing_buffer.drain(0..1); - resync_steps += 1; - if resync_steps >= MAX_RESYNC_STEPS_PER_CALL { - return Ok(None); - } - continue; - } - - // Decode full RawNetworkMessage from the frame using existing decoder - let mut cursor = std::io::Cursor::new(&state.framing_buffer[..total_len]); - match RawNetworkMessage::consensus_decode(&mut cursor) { - Ok(raw_message) => { - // Consume bytes - state.framing_buffer.drain(0..total_len); - - // Validate magic matches our network - if raw_message.magic != self.network.magic() { - tracing::warn!( - "Received message with wrong magic bytes: expected {:#x}, got {:#x}", - self.network.magic(), - raw_message.magic - ); - return Err(NetworkError::ProtocolError(format!( - "Wrong magic bytes: expected {:#x}, got {:#x}", - self.network.magic(), - raw_message.magic - ))); - } - - tracing::trace!( - "Successfully decoded message from {}: {:?}", - self.address, - raw_message.payload.cmd() - ); - - return Ok(Some(Message::new(self.address, raw_message.payload))); - } - Err(e) => { - tracing::warn!( - "{}: decode error after framing ({}), attempting resync", - self.address, - e - ); - state.framing_buffer.drain(0..1); - resync_steps += 1; - if resync_steps >= MAX_RESYNC_STEPS_PER_CALL { - return Ok(None); - } - continue; - } - } - } - } - .await; - - // Drop the lock before disconnecting - drop(state); - - // Handle disconnection if needed - if let Err(NetworkError::PeerDisconnected) = &result { - self.state = None; - self.connected_at = None; + /// Note that `n` earlier pipeline requests have fully completed, freeing that + /// much in-flight work. Used for streaming responses (`getcfilters` -> many + /// `cfilter`s) that the reader can't attribute to a finished request on its + /// own; single-response requests are decremented directly in the reader. + pub(crate) async fn response_completed(&self, n: usize) { + let _ = self + .in_flight + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |v| Some(v.saturating_sub(n))); + for _ in 0..n { + self.latency.complete_one().await; } - - result } - /// Check if the connection is active. - pub fn is_connected(&self) -> bool { - self.state.is_some() + /// Per-peer response latency: (completed requests, average ms, worst ms). + pub(crate) fn latency_stats(&self) -> (u64, f64, f64) { + self.latency.snapshot() } - /// Check if connection appears healthy (not just connected). - pub fn is_healthy(&self) -> bool { - if !self.is_connected() { - tracing::debug!("Connection to {} marked unhealthy: not connected", self.address); - return false; - } - - let now = SystemTime::now(); - - // If we have exchanged pings/pongs, check the last activity - if let Some(last_pong) = self.last_pong_received { - if let Ok(duration) = now.duration_since(last_pong) { - // If no pong in 10 minutes, consider unhealthy - if duration > Duration::from_secs(600) { - tracing::warn!("Connection to {} marked unhealthy: no pong received for {} seconds (limit: 600)", - self.address, duration.as_secs()); - return false; - } - } - } else if let Some(connected_at) = self.connected_at { - // If we haven't received any pongs yet, check how long we've been connected - if let Ok(duration) = now.duration_since(connected_at) { - // Give new connections 5 minutes before considering them unhealthy - if duration > Duration::from_secs(300) { - tracing::warn!("Connection to {} marked unhealthy: no pong activity after {} seconds (limit: 300, last_ping_sent: {:?})", - self.address, duration.as_secs(), self.last_ping_sent.is_some()); - return false; - } - } - } - - // Connection is healthy - true + /// Cumulative (completed request count, total service-time ns) for this peer. + /// The bandwidth controller diffs these across a window to size this peer's + /// in-flight cap independently by Little's Law. + pub(crate) fn latency_totals(&self) -> (u64, u64) { + self.latency.totals() } - /// Get connection statistics. - pub fn stats(&self) -> (u64, u64) { - (self.bytes_sent, 0) // TODO: Track bytes received + /// Handshake round-trip latency in ms (0 if unmeasured). + pub(crate) fn lag_ms(&self) -> u32 { + self.lag_ms.load(Ordering::Relaxed) } - /// Send a ping message with a random nonce. - pub async fn send_ping(&mut self) -> NetworkResult { - let nonce = rand::random::(); - let ping_message = NetworkMessage::Ping(nonce); + /// Close this connection: cancel its reader so the socket shuts down. Used to + /// drop probed-but-unselected peers instead of leaking their readers. + pub(crate) fn close(&self) { + self.token.cancel(); + } - self.send_message(ping_message).await?; + /// This peer's current measured in-flight capacity (its serving BDP). + pub(crate) fn cap(&self) -> usize { + self.cap.load(Ordering::Relaxed) + } - let now = SystemTime::now(); - self.last_ping_sent = Some(now); - self.pending_pings.insert(nonce, now); + /// Cumulative bytes downloaded from this peer. The controller diffs it across + /// a window for this peer's throughput. + pub(crate) fn bytes_read(&self) -> u64 { + self.bytes.load(Ordering::Relaxed) + } - tracing::trace!("Sent ping to {} with nonce {}", self.address, nonce); + /// Update this peer's measured in-flight capacity (called by the controller). + pub(crate) fn set_cap(&self, n: usize) { + self.cap.store(n, Ordering::Relaxed); + } +} - Ok(nonce) +/// Pipeline requests we send and expect a response for (each adds one in-flight). +fn is_pipeline_request(msg: &NetworkMessage) -> bool { + match msg { + NetworkMessage::GetHeaders(_) + | NetworkMessage::GetHeaders2(_) + | NetworkMessage::GetCFHeaders(_) + | NetworkMessage::GetCFilters(_) => true, + // Block `getdata` is paced like any other request: the blocks pipeline + // sends ONE block per message, so a request is exactly one in-flight unit + // and one `block` in reply. Other `getdata` (chainlocks, islocks, txs) is + // control traffic — it carries several inventory items whose replies the + // reader cannot attribute one-for-one, so counting it would leak slots. + NetworkMessage::GetData(inv) => { + !inv.is_empty() + && inv + .iter() + .all(|i| matches!(i, dashcore::network::message_blockdata::Inventory::Block(_))) + } + _ => false, } +} - /// Handle a received ping message by sending a pong response. - pub async fn handle_ping(&mut self, nonce: u64) -> NetworkResult<()> { - let pong_message = NetworkMessage::Pong(nonce); - self.send_message(pong_message).await?; +/// Single-message responses: the reader decrements one in-flight per message. +/// `cfilter` is excluded — one `getcfilters` yields up to 1000 `cfilter`s, so its +/// unit is freed once per batch by the filters pipeline via `response_completed`. +fn is_single_response(msg: &NetworkMessage) -> bool { + matches!( + msg, + NetworkMessage::Headers(_) + | NetworkMessage::Headers2(_) + | NetworkMessage::CFHeaders(_) + | NetworkMessage::Block(_) + ) +} - tracing::trace!("Responded to ping from {} with pong nonce {}", self.address, nonce); +impl DisconnectedPeer { + pub fn new(addr: SocketAddr, network: Network) -> Self { + DisconnectedPeer { + network, + addr, + lag_ms: None, + } + } - Ok(()) + pub fn addr(&self) -> SocketAddr { + self.addr } - /// Handle a received pong message by validating the nonce. - pub fn handle_pong(&mut self, nonce: u64) -> NetworkResult<()> { - if let Some(sent_time) = self.pending_pings.remove(&nonce) { - let now = SystemTime::now(); - let rtt = now.duration_since(sent_time).unwrap_or(Duration::from_secs(0)); + /// Handshake ping measured on a prior connection, if this address has been probed + /// before. `None` sorts as worst (unknown) when ranking backups. + pub(crate) fn lag_ms(&self) -> Option { + self.lag_ms + } - self.last_pong_received = Some(now); + #[allow(clippy::too_many_arguments)] + pub async fn connect( + self, + inbound: UnboundedSender, + shutdown: CancellationToken, + bytes: Arc, + required_services: ServiceFlags, + ) -> NetworkResult { + let stream = TcpStream::connect(&self.addr).await.map_err(|e| { + NetworkError::ConnectionFailed(format!("Failed to connect to {}: {}", self.addr, e)) + })?; - tracing::debug!( - "Received valid pong from {} with nonce {} (RTT: {:?})", - self.address, - nonce, - rtt - ); + let peer_bytes = Arc::new(AtomicU64::new(0)); + let (read_half, mut writer) = stream.into_split(); + let mut reader = FramedRead::new( + CountingReader { + inner: read_half, + bytes, + peer_bytes: peer_bytes.clone(), + }, + RawNetworkMessageCodec, + ); + let magic = self.network.magic(); - Ok(()) - } else { - tracing::warn!("Received unexpected pong from {} with nonce {}", self.address, nonce); - Err(NetworkError::ProtocolError(format!( - "Unexpected pong nonce {} from {}", - nonce, self.address - ))) - } - } + handshake_send(&mut writer, magic, NetworkMessage::Version(build_version(self.addr))) + .await?; - /// Check if we need to send a ping (no ping/pong activity for 2 minutes). - pub fn should_ping(&self) -> bool { - let now = SystemTime::now(); + let deadline = tokio::time::Instant::now() + HANDSHAKE_TIMEOUT; + let mut peer_version: Option = None; + let mut got_verack = false; - // Check if we've sent a ping recently - if let Some(last_ping) = self.last_ping_sent { - if now.duration_since(last_ping).unwrap_or(Duration::MAX) < PING_INTERVAL { - return false; + while !(peer_version.is_some() && got_verack) { + let raw = match tokio::time::timeout_at(deadline, reader.next()).await { + Err(_) => return Err(NetworkError::Timeout), + Ok(None) => return Err(NetworkError::PeerDisconnected), + Ok(Some(Err(e))) => return Err(e.into()), + Ok(Some(Ok(raw))) => raw, + }; + if raw.magic != magic { + return Err(NetworkError::ProtocolError("wrong network magic".into())); } - } - - // Check if we've received a pong recently - if let Some(last_pong) = self.last_pong_received { - if now.duration_since(last_pong).unwrap_or(Duration::MAX) < PING_INTERVAL { - return false; + match raw.payload { + NetworkMessage::Version(v) => { + // BIP155: sendaddrv2 must be sent BEFORE verack. + handshake_send(&mut writer, magic, NetworkMessage::SendAddrV2).await?; + handshake_send(&mut writer, magic, NetworkMessage::Verack).await?; + peer_version = Some(v); + } + NetworkMessage::Verack => got_verack = true, + NetworkMessage::Ping(n) => { + handshake_send(&mut writer, magic, NetworkMessage::Pong(n)).await? + } + _ => {} } } - // If we haven't sent a ping or received a pong in 2 minutes, we should ping - true - } - - /// Remove pending pings that have timed out. - /// Returns `true` if any pings were removed. - pub fn remove_expired_pings(&mut self) -> bool { - const PING_TIMEOUT: Duration = Duration::from_secs(60); // 1 minute timeout for pings - - let now = SystemTime::now(); - let mut expired_nonces = Vec::new(); + let version = peer_version.ok_or(NetworkError::PeerDisconnected)?; - for (&nonce, &sent_time) in &self.pending_pings { - if now.duration_since(sent_time).unwrap_or(Duration::ZERO) > PING_TIMEOUT { - expired_nonces.push(nonce); - } + // Only keep peers that advertise the services we need (see `new` in `manager` + // for how the set is composed). A peer missing one of them doesn't fail loudly: + // it stays connected and simply ignores the requests it can't serve, so its + // slots stall until the request times out. Drop it now rather than waste it. + if !version.services.has(required_services) { + tracing::debug!( + target: "dash_spv::network", + "dropping {}: lacks required services {:?} (advertises {:?})", + self.addr, required_services, version.services + ); + return Err(NetworkError::ConnectionFailed(format!( + "peer {} lacks required services {:?}", + self.addr, required_services + ))); } - let has_expired = !expired_nonces.is_empty(); - for nonce in expired_nonces { - self.pending_pings.remove(&nonce); - tracing::warn!("Ping timeout for {} with nonce {}", self.address, nonce); + // Announce sendheaders only after the handshake is fully complete. + handshake_send(&mut writer, magic, NetworkMessage::SendHeaders).await?; + + // Measure round-trip lag with a post-handshake ping/pong. Sending a ping + // before the handshake completes makes some peers drop us, so we do it here. + let mut lag_ms: u32 = 0; + let ping_nonce: u64 = rand::random(); + let ping_sent = tokio::time::Instant::now(); + if handshake_send(&mut writer, magic, NetworkMessage::Ping(ping_nonce)).await.is_ok() { + let deadline = ping_sent + HANDSHAKE_TIMEOUT; + while let Ok(Some(Ok(raw))) = tokio::time::timeout_at(deadline, reader.next()).await { + if raw.magic != magic { + continue; + } + match raw.payload { + NetworkMessage::Pong(n) if n == ping_nonce => { + lag_ms = ping_sent.elapsed().as_millis().clamp(1, u32::MAX as u128) as u32; + break; + } + NetworkMessage::Ping(n) => { + let _ = handshake_send(&mut writer, magic, NetworkMessage::Pong(n)).await; + } + _ => {} + } + } } - has_expired - } - - /// Get ping/pong statistics. - pub fn ping_stats(&self) -> (Option, Option, usize) { - (self.last_ping_sent, self.last_pong_received, self.pending_pings.len()) - } + let writer = Arc::new(Mutex::new(writer)); + let in_flight = Arc::new(AtomicUsize::new(0)); + let latency = Arc::new(Latency::default()); + // Start with a tiny in-flight capacity; the controller grows it from this + // peer's measured serving rate. + let cap = Arc::new(AtomicUsize::new(2)); + // Per-connection token: cancelled by the global shutdown (parent) OR by + // `close()` to drop just this peer. + let token = shutdown.child_token(); + spawn_reader( + self.addr, + magic, + reader, + writer.clone(), + inbound, + in_flight.clone(), + latency.clone(), + token.clone(), + ); - /// Set that peer prefers headers2. - pub fn set_prefers_headers2(&mut self, prefers: bool) { - self.prefers_headers2 = prefers; - if prefers { - tracing::info!("Peer {} prefers headers2 compression", self.address); - } - } + tracing::debug!( + target: "dash_spv::network", + "peer connected: {} | lag={}ms height={}", + self.addr, + lag_ms, + version.start_height, + ); - /// Check if peer prefers headers2. - pub fn prefers_headers2(&self) -> bool { - self.prefers_headers2 + Ok(ConnectedPeer { + network: self.network, + addr: self.addr, + version, + lag_ms: AtomicU32::new(lag_ms), + in_flight, + latency, + writer, + cap, + bytes: peer_bytes, + token, + }) } +} - /// Set that peer sent us SendHeaders2. - pub fn set_peer_sent_sendheaders2(&mut self, sent: bool) { - self.sent_sendheaders2 = sent; - if sent { - tracing::info!( - "Peer {} sent SendHeaders2 - they will send compressed headers", - self.address - ); +#[allow(clippy::too_many_arguments)] +fn spawn_reader( + addr: SocketAddr, + magic: u32, + mut reader: PeerReader, + writer: Arc>, + inbound: UnboundedSender, + in_flight: Arc, + latency: Arc, + shutdown: CancellationToken, +) { + tokio::spawn(async move { + loop { + let next = tokio::select! { + _ = shutdown.cancelled() => break, + next = reader.next() => next, + }; + match next { + None => break, + Some(Err(e)) => { + tracing::error!("NETWORK: reader {} stopped: {}", addr, e); + break; + } + Some(Ok(raw)) => { + if raw.magic != magic { + continue; + } + // A single-message response completes one unit of in-flight + // work. Streaming responses (cfilter) are freed per batch by + // the filters pipeline instead. + if is_single_response(&raw.payload) { + let _ = in_flight.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |v| { + Some(v.saturating_sub(1)) + }); + latency.complete_one().await; + } + match raw.payload { + NetworkMessage::Ping(nonce) => { + let pong = RawNetworkMessage { + magic, + payload: NetworkMessage::Pong(nonce), + }; + if writer + .lock() + .await + .write_all(&encode::serialize(&pong)) + .await + .is_err() + { + break; + } + } + payload => { + if inbound.send(PeerEvent::Message(addr, payload)).is_err() { + break; + } + } + } + } + } } - } - - /// Check if peer sent us SendHeaders2. - pub fn peer_sent_sendheaders2(&self) -> bool { - self.sent_sendheaders2 - } - /// Check if we can request headers2 from this peer. - pub fn can_request_headers2(&self) -> bool { - // We can request headers2 if peer has the service flag for headers2 support - // Note: We don't wait for SendHeaders2 from peer as that creates a race condition - // during initial sync. The service flag is sufficient to know they support headers2. - if let Some(services) = self.services { - dashcore::network::constants::ServiceFlags::from(services) - .has(dashcore::network::constants::NODE_HEADERS_COMPRESSED) - } else { - false - } - } + tracing::info!("NETWORK: peer {} disconnected", addr); + let _ = inbound.send(PeerEvent::Disconnected(addr)); + }); } -#[cfg(test)] -impl Peer { - pub(crate) fn set_services(&mut self, flags: ServiceFlags) { - self.services = Some(flags.as_u64()); - } +fn build_version(peer: SocketAddr) -> VersionMessage { + let now = SystemTime::now().duration_since(UNIX_EPOCH).map(|d| d.as_secs() as i64).unwrap_or(0); + let unspecified: SocketAddr = ([0u8, 0, 0, 0], 0).into(); + VersionMessage::new( + ServiceFlags::NONE, + now, + Address::new(&peer, ServiceFlags::NETWORK), + Address::new(&unspecified, ServiceFlags::NONE), + rand::random(), + USER_AGENT.to_string(), + 0, + false, + [0u8; 32], + ) } -#[cfg(test)] -mod tests { - use std::net::SocketAddr; - use std::time::{Duration, SystemTime}; - - use super::Peer; - - #[test] - fn remove_expired_pings() { - let addr: SocketAddr = "127.0.0.1:9999".parse().unwrap(); - let mut peer = Peer::dummy(addr); - let now = SystemTime::now(); - let expired = now - Duration::from_secs(61); - - // No pings at all - assert!(!peer.remove_expired_pings()); - - // Only recent pings — nothing removed - peer.pending_pings.insert(1, now); - peer.pending_pings.insert(2, now); - assert!(!peer.remove_expired_pings()); - assert_eq!(peer.pending_pings.len(), 2); - - // Add an expired ping — only it gets removed - peer.pending_pings.insert(3, expired); - assert!(peer.remove_expired_pings()); - assert_eq!(peer.pending_pings.len(), 2); - assert!(!peer.pending_pings.contains_key(&3)); - - // All expired — map ends up empty - peer.pending_pings.clear(); - peer.pending_pings.insert(10, expired); - peer.pending_pings.insert(20, expired); - assert!(peer.remove_expired_pings()); - assert!(peer.pending_pings.is_empty()); - } +async fn handshake_send( + writer: &mut OwnedWriteHalf, + magic: u32, + payload: NetworkMessage, +) -> NetworkResult<()> { + let raw = RawNetworkMessage { + magic, + payload, + }; + writer + .write_all(&encode::serialize(&raw)) + .await + .map_err(|e| NetworkError::ConnectionFailed(format!("handshake write failed: {}", e)))?; + writer + .flush() + .await + .map_err(|e| NetworkError::ConnectionFailed(format!("handshake flush failed: {}", e)))?; + Ok(()) } diff --git a/dash-spv/src/network/pool.rs b/dash-spv/src/network/pool.rs deleted file mode 100644 index 1a2e68b21..000000000 --- a/dash-spv/src/network/pool.rs +++ /dev/null @@ -1,316 +0,0 @@ -//! Peer pool for managing multiple peer connections - -use crate::error::{NetworkError, SpvError as Error}; -use crate::network::peer::Peer; -use dashcore::network::constants::ServiceFlags; -use dashcore::prelude::CoreBlockHeight; -use std::collections::{HashMap, HashSet}; -use std::net::SocketAddr; -use std::sync::Arc; -use tokio::sync::RwLock; - -/// Pool for managing multiple peer instances -pub struct PeerPool { - /// Active peers mapped by address - peers: Arc>>>>, - /// Addresses currently being connected to - connecting: Arc>>, - /// Maximum number of simultaneous peer connections (from `ClientConfig::max_peers`). - max_peers: usize, -} - -impl PeerPool { - /// Create a new peer pool with a connection cap. - pub fn new(max_peers: usize) -> Self { - // Assert peers are greater than 0. We may change this - // so 0 means 'connect to as many peers as you can' - debug_assert!(max_peers > 0, "max_peers must be greater than 0 for the spv client to sync"); - - Self { - peers: Arc::new(RwLock::new(HashMap::new())), - connecting: Arc::new(RwLock::new(HashSet::new())), - max_peers, - } - } - - /// Mark an address as being connected to - pub async fn mark_connecting(&self, addr: SocketAddr) -> bool { - let mut connecting = self.connecting.write().await; - connecting.insert(addr) - } - - /// Add a peer to the pool - pub async fn add_peer(&self, addr: SocketAddr, peer: Peer) -> Result<(), Error> { - let mut peers = self.peers.write().await; - let mut connecting = self.connecting.write().await; - - // Remove from connecting set - connecting.remove(&addr); - - // Check if we're at capacity - if peers.len() >= self.max_peers { - return Err(Error::Network(NetworkError::ConnectionFailed(format!( - "Maximum peers ({}) reached", - self.max_peers - )))); - } - - // Check if already connected - if peers.contains_key(&addr) { - return Err(Error::Network(NetworkError::ConnectionFailed(format!( - "Already connected to {}", - addr - )))); - } - - peers.insert(addr, Arc::new(RwLock::new(peer))); - tracing::info!("Added peer {}, total peers: {}", addr, peers.len()); - Ok(()) - } - - /// Remove a peer from the pool and clear connecting state - pub async fn remove_peer(&self, addr: &SocketAddr) -> Option>> { - self.connecting.write().await.remove(addr); - let removed = self.peers.write().await.remove(addr); - if removed.is_some() { - tracing::info!("Removed peer {}", addr); - } - removed - } - - /// Get all active peers - pub async fn get_all_peers(&self) -> Vec<(SocketAddr, Arc>)> { - self.peers.read().await.iter().map(|(addr, peer)| (*addr, peer.clone())).collect() - } - - /// Get a specific peer - pub async fn get_peer(&self, addr: &SocketAddr) -> Option>> { - self.peers.read().await.get(addr).cloned() - } - - /// Get the number of active peers - pub async fn peer_count(&self) -> usize { - self.peers.read().await.len() - } - - /// Check if connected to a specific peer - pub async fn is_connected(&self, addr: &SocketAddr) -> bool { - self.peers.read().await.contains_key(addr) - } - - /// Check if currently connecting to a peer - pub async fn is_connecting(&self, addr: &SocketAddr) -> bool { - self.connecting.read().await.contains(addr) - } - - /// Get all connected peer addresses - pub async fn get_connected_addresses(&self) -> Vec { - self.peers.read().await.keys().copied().collect() - } - - pub async fn get_best_height(&self) -> Option { - let peers = self.get_all_peers().await; - - if peers.is_empty() { - tracing::debug!("get_best_height: No peers available"); - return None; - } - - let mut best_height = 0u32; - let mut peer_count = 0; - - for (addr, peer) in peers.iter() { - let peer_guard = peer.read().await; - peer_count += 1; - - tracing::debug!( - "get_best_height: Peer {} - best_height: {:?}, version: {:?}, connected: {}", - addr, - peer_guard.best_height(), - peer_guard.version(), - peer_guard.is_connected(), - ); - - if let Some(peer_height) = peer_guard.best_height() { - if peer_height > 0 { - best_height = best_height.max(peer_height); - tracing::debug!( - "get_best_height: Updated best_height to {} from peer {}", - best_height, - addr - ); - } - } - } - - tracing::debug!( - "get_best_height: Checked {} peers, best_height: {}", - peer_count, - best_height - ); - - if best_height > 0 { - Some(best_height) - } else { - None - } - } - - /// Find the first connected peer that advertises the given service flags. - pub(crate) async fn peer_with_service( - &self, - flags: ServiceFlags, - ) -> Option<(SocketAddr, Arc>)> { - let peers = self.peers.read().await; - for (addr, peer) in peers.iter() { - if peer.read().await.has_service(flags) { - return Some((*addr, Arc::clone(peer))); - } - } - None - } - - /// Collect all connected peers that advertise the given service flags. - pub(crate) async fn peers_with_service( - &self, - flags: ServiceFlags, - ) -> Vec<(SocketAddr, Arc>)> { - let peers = self.peers.read().await; - let mut result = Vec::new(); - for (addr, peer) in peers.iter() { - if peer.read().await.has_service(flags) { - result.push((*addr, peer.clone())); - } - } - result - } - - /// Check whether any connected peer advertises the given service flags. - pub(crate) async fn has_peers_with_service(&self, flags: ServiceFlags) -> bool { - let peers = self.peers.read().await; - for peer in peers.values() { - if peer.read().await.has_service(flags) { - return true; - } - } - false - } - - /// Check if we need more peers - pub async fn needs_more_peers(&self) -> bool { - self.peer_count().await < self.max_peers - } - - /// Check if we can accept more peers - pub async fn can_accept_peers(&self) -> bool { - self.peer_count().await < self.max_peers - } - - /// Remove unhealthy peers and return their addresses so the caller can - /// emit the appropriate network events. - pub async fn remove_unhealthy(&self) -> Vec { - let peers = self.peers.read().await; - let mut unhealthy = Vec::new(); - - // Check each peer's health - for (addr, peer) in peers.iter() { - // Use blocking read to properly check health - let peer_guard = peer.read().await; - if !peer_guard.is_healthy() { - unhealthy.push(*addr); - } - } - - // Release read lock before taking write lock - drop(peers); - - // Remove unhealthy connections - if !unhealthy.is_empty() { - let mut peers = self.peers.write().await; - unhealthy.retain(|addr| peers.remove(addr).is_some()); - } - - unhealthy - } -} - -impl Default for PeerPool { - fn default() -> Self { - Self::new(8) - } -} - -#[cfg(test)] -impl PeerPool { - pub(crate) async fn insert_peer_with_services(&self, addr: SocketAddr, flags: ServiceFlags) { - let mut peer = Peer::dummy(addr); - peer.set_services(flags); - self.peers.write().await.insert(addr, Arc::new(RwLock::new(peer))); - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[tokio::test] - async fn test_peer_pool_basic() { - let pool = PeerPool::new(8); - - // Initial state - assert_eq!(pool.peer_count().await, 0); - assert!(pool.needs_more_peers().await); - assert!(pool.can_accept_peers().await); - - // Test marking as connecting - let addr = "127.0.0.1:9999".parse().expect("Failed to parse test address"); - assert!(pool.mark_connecting(addr).await); - assert!(!pool.mark_connecting(addr).await); // Already marked - assert!(pool.is_connecting(&addr).await); - } - - #[tokio::test] - async fn test_service_lookup() { - let pool = PeerPool::new(8); - let compact_filters = ServiceFlags::COMPACT_FILTERS; - let combined = compact_filters | ServiceFlags::NODE_HEADERS_COMPRESSED; - - // No matches on empty pool - assert!(pool.peer_with_service(compact_filters).await.is_none()); - assert!(pool.peers_with_service(compact_filters).await.is_empty()); - - // No matches when peers lack the requested flag - let addr1: SocketAddr = "127.0.0.1:1001".parse().unwrap(); - pool.insert_peer_with_services(addr1, ServiceFlags::NETWORK).await; - assert!(pool.peer_with_service(compact_filters).await.is_none()); - assert!(pool.peers_with_service(compact_filters).await.is_empty()); - - // Single-flag lookup returns matching peers - let addr2: SocketAddr = "127.0.0.1:1002".parse().unwrap(); - let addr3: SocketAddr = "127.0.0.1:1003".parse().unwrap(); - pool.insert_peer_with_services(addr2, ServiceFlags::NETWORK | compact_filters).await; - pool.insert_peer_with_services(addr3, ServiceFlags::NETWORK | combined).await; - - let (found_addr, found_peer) = pool.peer_with_service(compact_filters).await.unwrap(); - assert!(found_addr == addr2 || found_addr == addr3); - assert!(found_peer.read().await.has_service(compact_filters)); - - let filter_peers: HashMap = - pool.peers_with_service(compact_filters).await.into_iter().collect(); - assert_eq!(filter_peers.len(), 2); - assert!(filter_peers.contains_key(&addr2)); - assert!(filter_peers.contains_key(&addr3)); - - // Combined flags require all bits present - let (found_addr, _) = pool.peer_with_service(combined).await.unwrap(); - assert_eq!(found_addr, addr3); - let combined_peers = pool.peers_with_service(combined).await; - assert_eq!(combined_peers.len(), 1); - assert_eq!(combined_peers[0].0, addr3); - - // NONE matches every peer in the pool - assert!(pool.peer_with_service(ServiceFlags::NONE).await.is_some()); - let all = pool.peers_with_service(ServiceFlags::NONE).await; - assert_eq!(all.len(), 3); - } -} diff --git a/dash-spv/src/network/reputation.rs b/dash-spv/src/network/reputation.rs deleted file mode 100644 index f90656584..000000000 --- a/dash-spv/src/network/reputation.rs +++ /dev/null @@ -1,439 +0,0 @@ -//! Peer reputation management system -//! -//! This module implements a reputation system to track peer behavior and protect -//! against malicious peers. It tracks both positive and negative behaviors, -//! implements automatic banning for excessive misbehavior, and provides reputation -//! decay over time for recovery. - -use crate::storage::PeerStorage; -use dashcore::network::address::AddrV2Message; -use serde::{Deserialize, Deserializer, Serialize}; -use std::collections::HashMap; -use std::net::SocketAddr; -use std::sync::Arc; -use std::time::{Duration, Instant}; -use tokio::sync::RwLock; - -/// Reason for a peer reputation change. Each reason owns its score delta -/// (positive = penalty, negative = reward) and a human-readable label. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum ChangeReason { - HandshakeFailed, - ConnectionFailed, - Headers2DecompressionFailed, - ReadTimeout, - PingFailed, - InvalidTransactionInBlock, - ManuallyBanned, - LongUptime, -} - -impl ChangeReason { - /// Score delta for this reason: positive for misbehavior (penalty), - /// negative for good behavior (reward). - pub fn score(&self) -> i32 { - match self { - ChangeReason::HandshakeFailed => 10, - ChangeReason::ConnectionFailed => 2, - ChangeReason::Headers2DecompressionFailed => 10, - ChangeReason::ReadTimeout => 5, - ChangeReason::PingFailed => 5, - ChangeReason::InvalidTransactionInBlock => 20, - ChangeReason::ManuallyBanned => 100, - ChangeReason::LongUptime => -5, - } - } -} - -impl std::fmt::Display for ChangeReason { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let label = match self { - ChangeReason::HandshakeFailed => "Handshake failed", - ChangeReason::ConnectionFailed => "Connection failed", - ChangeReason::Headers2DecompressionFailed => "Headers2 decompression failed", - ChangeReason::ReadTimeout => "Read timeout", - ChangeReason::PingFailed => "Ping failed", - ChangeReason::InvalidTransactionInBlock => "Invalid transaction type in block", - ChangeReason::ManuallyBanned => "Manually banned", - ChangeReason::LongUptime => "Long connection uptime", - }; - f.write_str(label) - } -} - -/// Ban duration for misbehaving peers -const BAN_DURATION: Duration = Duration::from_secs(24 * 60 * 60); // 24 hours - -/// Reputation decay interval -const DECAY_INTERVAL: Duration = Duration::from_secs(60 * 60); // 1 hour - -/// Amount to decay reputation score per interval -const DECAY_AMOUNT: i32 = 5; - -/// Maximum misbehavior score before a peer is banned -const MAX_MISBEHAVIOR_SCORE: i32 = 100; - -/// Minimum score (most positive reputation) -const MIN_MISBEHAVIOR_SCORE: i32 = -50; - -const MAX_BAN_COUNT: u32 = 1000; - -const MAX_ACTION_COUNT: u64 = 1_000_000; - -fn clamp_peer_score<'de, D>(deserializer: D) -> Result -where - D: Deserializer<'de>, -{ - let mut v = i32::deserialize(deserializer)?; - - if v < MIN_MISBEHAVIOR_SCORE { - tracing::warn!("Peer has invalid score {v}, clamping to min {MIN_MISBEHAVIOR_SCORE}"); - v = MIN_MISBEHAVIOR_SCORE - } else if v > MAX_MISBEHAVIOR_SCORE { - tracing::warn!("Peer has invalid score {v}, clamping to max {MAX_MISBEHAVIOR_SCORE}"); - v = MAX_MISBEHAVIOR_SCORE - } - - Ok(v) -} - -fn clamp_peer_ban_count<'de, D>(deserializer: D) -> Result -where - D: Deserializer<'de>, -{ - let mut v = u32::deserialize(deserializer)?; - - if v > MAX_BAN_COUNT { - tracing::warn!("Peer has excessive ban count {v}, clamping to {MAX_BAN_COUNT}"); - v = MAX_BAN_COUNT - } - - Ok(v) -} - -fn clamp_peer_connection_attempts<'de, D>(deserializer: D) -> Result -where - D: Deserializer<'de>, -{ - let mut v = u64::deserialize(deserializer)?; - - v = v.min(MAX_ACTION_COUNT); - - Ok(v) -} - -/// Peer reputation entry -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct PeerReputation { - /// Current misbehavior score - #[serde(deserialize_with = "clamp_peer_score")] - pub score: i32, - - /// Number of times this peer has been banned - #[serde(deserialize_with = "clamp_peer_ban_count")] - pub ban_count: u32, - - /// Time when the peer was banned (if currently banned) - #[serde(skip)] - pub banned_until: Option, - - /// Last time the reputation was updated - #[serde(skip, default = "Instant::now")] - pub last_update: Instant, - - /// Total number of positive actions - pub positive_actions: u64, - - /// Total number of negative actions - pub negative_actions: u64, - - /// Connection count - #[serde(deserialize_with = "clamp_peer_connection_attempts")] - pub connection_attempts: u64, - - /// Successful connection count - pub successful_connections: u64, - - /// Last connection time - #[serde(skip)] - pub last_connection: Option, -} - -impl Default for PeerReputation { - fn default() -> Self { - Self { - score: 0, - ban_count: 0, - banned_until: None, - last_update: Instant::now(), - positive_actions: 0, - negative_actions: 0, - connection_attempts: 0, - successful_connections: 0, - last_connection: None, - } - } -} - -impl PeerReputation { - /// Check if the peer is currently banned - pub fn is_banned(&self) -> bool { - self.banned_until.is_some_and(|until| Instant::now() < until) - } - - /// Get remaining ban time - pub fn ban_time_remaining(&self) -> Option { - self.banned_until.and_then(|until| { - let now = Instant::now(); - if now < until { - Some(until - now) - } else { - None - } - }) - } - - /// Apply reputation decay - pub fn apply_decay(&mut self) { - let now = Instant::now(); - let elapsed = now - self.last_update; - - // Apply decay for each interval that has passed - let intervals = elapsed.as_secs() / DECAY_INTERVAL.as_secs(); - if intervals > 0 { - // Use saturating conversion to prevent overflow - // Cap at a reasonable maximum to avoid excessive decay - let intervals_i32 = intervals.min(i32::MAX as u64) as i32; - let decay = intervals_i32.saturating_mul(DECAY_AMOUNT); - self.score = (self.score - decay).max(MIN_MISBEHAVIOR_SCORE); - self.last_update = now; - } - - // Check if ban has expired - if self.is_banned() && self.ban_time_remaining().is_none() { - self.banned_until = None; - } - } -} - -/// Peer reputation manager -pub struct PeerReputationManager { - /// Reputation data for each peer - reputations: Arc>>, -} - -impl Default for PeerReputationManager { - fn default() -> Self { - Self::new() - } -} - -impl PeerReputationManager { - /// Create a new reputation manager - pub fn new() -> Self { - Self { - reputations: Arc::new(RwLock::new(HashMap::new())), - } - } - - /// Update peer reputation by the score delta of `reason`. - pub async fn update_reputation(&self, peer: SocketAddr, reason: ChangeReason) -> bool { - let score_change = reason.score(); - - let mut reputations = self.reputations.write().await; - let reputation = reputations.entry(peer).or_default(); - - // Apply decay first - reputation.apply_decay(); - - // Update score - let old_score = reputation.score; - reputation.score = - (reputation.score + score_change).clamp(MIN_MISBEHAVIOR_SCORE, MAX_MISBEHAVIOR_SCORE); - - // Track positive/negative actions - if score_change > 0 { - reputation.negative_actions += 1; - } else if score_change < 0 { - reputation.positive_actions += 1; - } - - // Check if peer should be banned - let should_ban = reputation.score >= MAX_MISBEHAVIOR_SCORE && !reputation.is_banned(); - if should_ban { - reputation.banned_until = Some(Instant::now() + BAN_DURATION); - reputation.ban_count += 1; - tracing::warn!( - "Peer {} banned for misbehavior (score: {}, ban #{}, reason: {})", - peer, - reputation.score, - reputation.ban_count, - reason - ); - } - - // Log significant changes - if score_change.abs() >= 10 || should_ban { - tracing::info!( - "Peer {} reputation changed: {} -> {} (change: {}, reason: {})", - peer, - old_score, - reputation.score, - score_change, - reason - ); - } - - should_ban - } - - /// Check if a peer is banned - pub async fn is_banned(&self, peer: &SocketAddr) -> bool { - let mut reputations = self.reputations.write().await; - if let Some(reputation) = reputations.get_mut(peer) { - reputation.apply_decay(); - reputation.is_banned() - } else { - false - } - } - - /// Record a connection attempt - pub async fn record_connection_attempt(&self, peer: SocketAddr) { - let mut reputations = self.reputations.write().await; - let reputation = reputations.entry(peer).or_default(); - reputation.connection_attempts += 1; - reputation.last_connection = Some(Instant::now()); - } - - /// Record a successful connection - pub async fn record_successful_connection(&self, peer: SocketAddr) { - let mut reputations = self.reputations.write().await; - let reputation = reputations.entry(peer).or_default(); - reputation.successful_connections += 1; - } - - /// Get all peer reputations - pub async fn get_all_reputations(&self) -> HashMap { - let mut reputations = self.reputations.write().await; - - // Apply decay to all peers - for reputation in reputations.values_mut() { - reputation.apply_decay(); - } - - reputations.clone() - } - - /// Clear banned status for a peer (admin function) - pub async fn unban_peer(&self, peer: &SocketAddr) { - let mut reputations = self.reputations.write().await; - if let Some(reputation) = reputations.get_mut(peer) { - reputation.banned_until = None; - reputation.score = reputation.score.min(MAX_MISBEHAVIOR_SCORE - 10); - tracing::info!("Manually unbanned peer {}", peer); - } - } - - /// Save reputation data to persistent storage - pub async fn save_to_storage(&self, storage: &impl PeerStorage) -> std::io::Result<()> { - let reputations = self.reputations.read().await; - - storage.save_peers_reputation(&reputations).await.map_err(std::io::Error::other) - } - - /// Load reputation data from persistent storage - pub async fn load_from_storage(&self, storage: &impl PeerStorage) -> std::io::Result<()> { - let data = storage.load_peers_reputation().await.map_err(std::io::Error::other)?; - - let mut reputations = self.reputations.write().await; - let mut loaded_count = 0; - let mut skipped_count = 0; - - for (addr, mut reputation) in data { - // Validate successful connections don't exceed attempts - reputation.successful_connections = - reputation.successful_connections.min(reputation.connection_attempts); - - // Skip entry if data appears corrupted - if reputation.positive_actions > MAX_ACTION_COUNT - || reputation.negative_actions > MAX_ACTION_COUNT - { - tracing::warn!("Skipping peer {} with potentially corrupted action counts", addr); - skipped_count += 1; - continue; - } - - // Apply initial decay based on ban count - if reputation.ban_count > 0 { - reputation.score = reputation.score.max(50); // Start with higher score for previously banned peers - } - - reputations.insert(addr, reputation); - loaded_count += 1; - } - - tracing::info!( - "Loaded reputation data for {} peers (skipped {} corrupted entries)", - loaded_count, - skipped_count - ); - Ok(()) - } -} - -/// Helper trait for reputation-aware peer selection -pub trait ReputationAware { - /// Select best peers based on reputation - fn select_best_peers( - &self, - available_peers: Vec, - count: usize, - ) -> impl std::future::Future> + Send; - - /// Check if we should connect to a peer based on reputation - fn should_connect_to_peer( - &self, - peer: &SocketAddr, - ) -> impl std::future::Future + Send; -} - -impl ReputationAware for PeerReputationManager { - async fn select_best_peers( - &self, - available_peers: Vec, - count: usize, - ) -> Vec { - let mut peer_scores = Vec::new(); - let mut reputations = self.reputations.write().await; - - for peer in available_peers { - let Ok(socket_addr) = peer.socket_addr() else { - tracing::warn!("Skip invalid peer address: {:?}", peer); - continue; - }; - - let reputation = reputations.entry(socket_addr).or_default(); - reputation.apply_decay(); - - if !reputation.is_banned() { - peer_scores.push((socket_addr, reputation.score)); - } - } - - // Sort by score (lower is better) - peer_scores.sort_by_key(|(_, score)| *score); - - // Return the best peers - peer_scores.into_iter().take(count).map(|(peer, _)| peer).collect() - } - - async fn should_connect_to_peer(&self, peer: &SocketAddr) -> bool { - !self.is_banned(peer).await - } -} - -// Include tests module -#[cfg(test)] -#[path = "reputation_tests.rs"] -mod reputation_tests; diff --git a/dash-spv/src/network/reputation_tests.rs b/dash-spv/src/network/reputation_tests.rs deleted file mode 100644 index 68b74e13b..000000000 --- a/dash-spv/src/network/reputation_tests.rs +++ /dev/null @@ -1,109 +0,0 @@ -//! Unit tests for reputation system (in-module tests) - -#[cfg(test)] -mod tests { - use crate::storage::{PersistentPeerStorage, PersistentStorage}; - - use super::super::*; - use std::net::SocketAddr; - - async fn score(manager: &PeerReputationManager, peer: &SocketAddr) -> i32 { - manager.get_all_reputations().await.get(peer).map_or(0, |rep| rep.score) - } - - #[tokio::test] - async fn test_basic_reputation_operations() { - let manager = PeerReputationManager::new(); - let peer: SocketAddr = "127.0.0.1:8333".parse().unwrap(); - - assert_eq!(score(&manager, &peer).await, 0); - - manager.update_reputation(peer, ChangeReason::HandshakeFailed).await; - assert_eq!(score(&manager, &peer).await, 10); - - manager.update_reputation(peer, ChangeReason::LongUptime).await; - assert_eq!(score(&manager, &peer).await, 5); - } - - #[tokio::test] - async fn test_banning_mechanism() { - let manager = PeerReputationManager::new(); - let peer: SocketAddr = "192.168.1.1:8333".parse().unwrap(); - - // Banned on the 10th violation (10 * 10 = 100). - for i in 0..10 { - let banned = manager.update_reputation(peer, ChangeReason::HandshakeFailed).await; - if i == 9 { - assert!(banned); - } else { - assert!(!banned); - } - } - - assert!(manager.is_banned(&peer).await); - } - - #[tokio::test] - async fn test_reputation_persistence() { - let manager = PeerReputationManager::new(); - let peer1: SocketAddr = "10.0.0.1:8333".parse().unwrap(); - let peer2: SocketAddr = "10.0.0.2:8333".parse().unwrap(); - - manager.update_reputation(peer1, ChangeReason::LongUptime).await; - manager.update_reputation(peer1, ChangeReason::LongUptime).await; - manager.update_reputation(peer2, ChangeReason::InvalidTransactionInBlock).await; - - let temp_dir = tempfile::TempDir::new().unwrap(); - let peer_storage = PersistentPeerStorage::open(temp_dir.path()) - .await - .expect("Failed to open PersistentPeerStorage"); - manager.save_to_storage(&peer_storage).await.unwrap(); - - let new_manager = PeerReputationManager::new(); - new_manager.load_from_storage(&peer_storage).await.unwrap(); - - assert_eq!(score(&new_manager, &peer1).await, -10); - assert_eq!(score(&new_manager, &peer2).await, 20); - } - - #[tokio::test] - async fn test_peer_selection() { - let manager = PeerReputationManager::new(); - - let good_peer = AddrV2Message::dummy(0, "1.1.1.1".parse().unwrap(), 8333); - let neutral_peer = AddrV2Message::dummy(0, "2.2.2.2".parse().unwrap(), 8333); - let bad_peer = AddrV2Message::dummy(0, "3.3.3.3".parse().unwrap(), 8333); - - manager.update_reputation(good_peer.socket_addr().unwrap(), ChangeReason::LongUptime).await; - manager - .update_reputation( - bad_peer.socket_addr().unwrap(), - ChangeReason::InvalidTransactionInBlock, - ) - .await; - - let all_peers = vec![good_peer.clone(), neutral_peer.clone(), bad_peer.clone()]; - let selected = manager.select_best_peers(all_peers, 2).await; - - assert_eq!(selected.len(), 2); - assert_eq!(selected[0], good_peer.socket_addr().unwrap()); - assert_eq!(selected[1], neutral_peer.socket_addr().unwrap()); - } - - #[tokio::test] - async fn test_connection_tracking() { - let manager = PeerReputationManager::new(); - let peer: SocketAddr = "127.0.0.1:9999".parse().unwrap(); - - // Track connection attempts - manager.record_connection_attempt(peer).await; - manager.record_connection_attempt(peer).await; - manager.record_successful_connection(peer).await; - - let reputations = manager.get_all_reputations().await; - let rep = &reputations[&peer]; - - assert_eq!(rep.connection_attempts, 2); - assert_eq!(rep.successful_connections, 1); - } -} diff --git a/dash-spv/src/network/tests.rs b/dash-spv/src/network/tests.rs deleted file mode 100644 index 4bb3447e0..000000000 --- a/dash-spv/src/network/tests.rs +++ /dev/null @@ -1,119 +0,0 @@ -//! Unit tests for network module - -#[cfg(test)] -mod peer_tests { - use crate::network::peer::Peer; - use dashcore::Network; - use std::time::Duration; - - #[test] - fn test_peer_creation() { - let addr = "127.0.0.1:9999".parse().unwrap(); - let timeout = Duration::from_secs(30); - let peer = Peer::new(addr, timeout, Network::Mainnet); - - assert!(!peer.is_connected()); - assert_eq!(peer.address(), addr); - } -} - -#[cfg(test)] -mod pool_tests { - use crate::network::manager::PeerNetworkManager; - use crate::network::peer::Peer; - use crate::network::pool::PeerPool; - use crate::test_utils::test_socket_address; - use dashcore::network::constants::ServiceFlags; - use dashcore::Network; - use tokio::time::Duration; - - #[tokio::test] - async fn test_pool_limits() { - let pool = PeerPool::new(8); - - // Test needs_more_peers logic - assert!(pool.needs_more_peers().await); - - // Can accept up to 8 peers - assert!(pool.can_accept_peers().await); - - // Test peer count - assert_eq!(pool.peer_count().await, 0); - } - - #[tokio::test] - async fn test_capability_policy_for_handshake_and_eviction() { - let cf = ServiceFlags::COMPACT_FILTERS; - let mut incapable = - Peer::new(test_socket_address(9), Duration::from_secs(10), Network::Testnet); - incapable.set_services(ServiceFlags::NETWORK); - - // Handshake admission: keep fallback when no capable peer exists yet. - let manager = PeerNetworkManager::new_for_test(cf).await; - assert!(!manager.test_has_capable_peer().await); - assert!(!manager.test_should_reject_after_handshake(&incapable).await); - - // Handshake admission: reject incapable peers once a capable peer exists. - let manager = PeerNetworkManager::new_for_test(cf).await; - manager.insert_test_peer(test_socket_address(1), cf).await; - assert!(manager.test_has_capable_peer().await); - assert!(manager.test_should_reject_after_handshake(&incapable).await); - - // Healthy pool: all peers match, nothing evicted - let manager = PeerNetworkManager::new_for_test(cf).await; - manager.insert_test_peer(test_socket_address(1), cf).await; - manager.insert_test_peer(test_socket_address(2), cf).await; - manager.insert_test_peer(test_socket_address(3), cf).await; - manager.evict_mismatched_peers().await; - assert_eq!(manager.test_peer_count().await, 3); - - // Lone mismatched peer is preserved (never drop to zero) - let manager = PeerNetworkManager::new_for_test(cf).await; - manager.insert_test_peer(test_socket_address(1), ServiceFlags::NETWORK).await; - manager.evict_mismatched_peers().await; - assert_eq!(manager.test_peer_count().await, 1); - - // All peers lack service: tick 1 drops all but 1, tick 2 preserves the lone peer - let manager = PeerNetworkManager::new_for_test(cf).await; - manager.insert_test_peer(test_socket_address(1), ServiceFlags::NETWORK).await; - manager.insert_test_peer(test_socket_address(2), ServiceFlags::NETWORK).await; - manager.insert_test_peer(test_socket_address(3), ServiceFlags::NETWORK).await; - manager.evict_mismatched_peers().await; - assert_eq!(manager.test_peer_count().await, 1); - manager.evict_mismatched_peers().await; - assert_eq!(manager.test_peer_count().await, 1); - - // Mixed pool: only mismatched peers are dropped, matching peers survive - let manager = PeerNetworkManager::new_for_test(cf).await; - let p1 = test_socket_address(1); - let p2 = test_socket_address(2); - let p3 = test_socket_address(3); - let p4 = test_socket_address(4); - manager.insert_test_peer(p1, cf).await; - manager.insert_test_peer(p2, cf).await; - manager.insert_test_peer(p3, ServiceFlags::NETWORK).await; - manager.insert_test_peer(p4, ServiceFlags::NETWORK).await; - manager.evict_mismatched_peers().await; - assert_eq!(manager.test_peer_count().await, 2); - assert!(manager.test_is_connected(&p1).await); - assert!(manager.test_is_connected(&p2).await); - assert!(!manager.test_is_connected(&p3).await); - assert!(!manager.test_is_connected(&p4).await); - } - - #[tokio::test(start_paused = true)] - async fn test_capability_rejection_cache_expires() { - let manager = PeerNetworkManager::new_for_test(ServiceFlags::COMPACT_FILTERS).await; - let fresh = test_socket_address(42); - let expired = test_socket_address(43); - - manager.insert_test_capability_rejected(expired).await; - tokio::time::advance(Duration::from_secs(31 * 60)).await; - manager.insert_test_capability_rejected(fresh).await; - - assert!(manager.test_is_capability_rejected(&fresh).await); - assert!(!manager.test_is_capability_rejected(&expired).await); - - assert_eq!(manager.test_capability_rejected_count().await, 1); - } -} diff --git a/dash-spv/src/storage/mod.rs b/dash-spv/src/storage/mod.rs index 248ab10d0..af2419072 100644 --- a/dash-spv/src/storage/mod.rs +++ b/dash-spv/src/storage/mod.rs @@ -10,7 +10,6 @@ mod io; mod lockfile; mod masternode; mod metadata; -mod peers; mod segments; use crate::error::StorageResult; use crate::storage::lockfile::LockFile; @@ -33,7 +32,6 @@ pub use crate::storage::filter_headers::{FilterHeaderStorage, PersistentFilterHe pub use crate::storage::filters::{FilterStorage, PersistentFilterStorage}; pub use crate::storage::masternode::{MasternodeStateStorage, PersistentMasternodeStateStorage}; pub use crate::storage::metadata::{MetadataStorage, PersistentMetadataStorage}; -pub use crate::storage::peers::{PeerStorage, PersistentPeerStorage}; pub use types::*; diff --git a/dash-spv/src/storage/peers.rs b/dash-spv/src/storage/peers.rs deleted file mode 100644 index 360e83650..000000000 --- a/dash-spv/src/storage/peers.rs +++ /dev/null @@ -1,204 +0,0 @@ -use std::{collections::HashMap, fs::File, io::BufReader, net::SocketAddr, path::PathBuf}; - -use tokio::fs; - -use async_trait::async_trait; -use dashcore::{ - consensus::{encode, Decodable, Encodable}, - network::address::AddrV2Message, -}; - -use crate::{ - error::StorageResult, - network::PeerReputation, - storage::{io::atomic_write, PersistentStorage}, - StorageError, -}; - -#[async_trait] -pub trait PeerStorage { - async fn save_peers( - &self, - peers: &[dashcore::network::address::AddrV2Message], - ) -> StorageResult<()>; - - async fn load_peers(&self) -> StorageResult>; - - async fn save_peers_reputation( - &self, - reputations: &HashMap, - ) -> StorageResult<()>; - - async fn load_peers_reputation(&self) -> StorageResult>; -} - -pub struct PersistentPeerStorage { - storage_path: PathBuf, -} - -impl PersistentPeerStorage { - const FOLDER_NAME: &str = "peers"; - - fn peers_data_file(&self) -> PathBuf { - self.storage_path.join("peers.dat") - } - - fn peers_reputation_file(&self) -> PathBuf { - self.storage_path.join("reputations.json") - } -} - -#[async_trait] -impl PersistentStorage for PersistentPeerStorage { - async fn open(storage_path: impl Into + Send) -> StorageResult { - let storage_path = storage_path.into(); - - Ok(PersistentPeerStorage { - storage_path: storage_path.join(Self::FOLDER_NAME), - }) - } - - async fn persist(&mut self, _storage_path: impl Into + Send) -> StorageResult<()> { - // Current implementation persists data everytime data is stored - Ok(()) - } -} - -#[async_trait] -impl PeerStorage for PersistentPeerStorage { - async fn save_peers( - &self, - peers: &[dashcore::network::address::AddrV2Message], - ) -> StorageResult<()> { - let peers_file = self.peers_data_file(); - - let mut buffer = Vec::new(); - - for item in peers.iter() { - item.consensus_encode(&mut buffer) - .map_err(|e| StorageError::WriteFailed(format!("Failed to encode peer: {}", e)))?; - } - - let peers_file_parent = peers_file - .parent() - .ok_or(StorageError::NotFound("peers_file doesn't have a parent".to_string()))?; - - tokio::fs::create_dir_all(peers_file_parent).await?; - - atomic_write(&peers_file, &buffer).await?; - - Ok(()) - } - - async fn load_peers(&self) -> StorageResult> { - let peers_file = self.peers_data_file(); - - if !fs::try_exists(&peers_file).await? { - return Ok(Vec::new()); - }; - - let peers = tokio::task::spawn_blocking(move || { - let file = File::open(&peers_file)?; - let mut reader = BufReader::new(file); - - let mut peers = Vec::new(); - - loop { - match AddrV2Message::consensus_decode(&mut reader) { - Ok(peer) => peers.push(peer), - Err(encode::Error::Io(ref e)) - if e.kind() == std::io::ErrorKind::UnexpectedEof => - { - break - } - Err(e) => { - return Err(StorageError::ReadFailed(format!("Failed to decode peer: {e}"))) - } - } - } - Ok(peers) - }) - .await - .map_err(|e| StorageError::ReadFailed(format!("Failed to load peers: {e}")))??; - - Ok(peers) - } - - async fn save_peers_reputation( - &self, - reputations: &HashMap, - ) -> StorageResult<()> { - let reputation_file = self.peers_reputation_file(); - - let json = serde_json::to_string_pretty(reputations).map_err(|e| { - StorageError::Serialization(format!("Failed to serialize peers reputations: {e}")) - })?; - - let reputation_file_parent = reputation_file - .parent() - .ok_or(StorageError::NotFound("reputation_file doesn't have a parent".to_string()))?; - - fs::create_dir_all(reputation_file_parent).await?; - - atomic_write(&reputation_file, json.as_bytes()).await - } - - async fn load_peers_reputation(&self) -> StorageResult> { - let reputation_file = self.peers_reputation_file(); - - if !fs::try_exists(&reputation_file).await? { - return Ok(HashMap::new()); - } - - let json = fs::read_to_string(reputation_file).await?; - serde_json::from_str(&json).map_err(|e| { - StorageError::ReadFailed(format!("Failed to deserialize peers reputations: {e}")) - }) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use dashcore::network::address::{AddrV2, AddrV2Message}; - use dashcore::network::constants::ServiceFlags; - use tempfile::TempDir; - - #[tokio::test] - async fn test_persistent_peer_storage_save_load() { - let temp_dir = TempDir::new().expect("Failed to create temporary directory for test"); - let store = PersistentPeerStorage::open(temp_dir.path()) - .await - .expect("Failed to open persistent peer storage"); - - // Create test peer messages - let addr: std::net::SocketAddr = - "192.168.1.1:9999".parse().expect("Failed to parse test address"); - let msg = AddrV2Message { - time: 1234567890, - services: ServiceFlags::NETWORK, - addr: AddrV2::Ipv4( - addr.ip().to_string().parse().expect("Failed to parse IPv4 address"), - ), - port: addr.port(), - }; - - store.save_peers(&[msg]).await.expect("Failed to save peers in test"); - - let loaded = store.load_peers().await.expect("Failed to load peers in test"); - assert_eq!(loaded.len(), 1); - assert_eq!(loaded[0].socket_addr().unwrap(), addr); - } - - #[tokio::test] - async fn test_persistent_peer_storage_empty() { - let temp_dir = TempDir::new().expect("Failed to create temporary directory for test"); - let store = PersistentPeerStorage::open(temp_dir.path()) - .await - .expect("Failed to open persistent peer storage"); - - // Load from non-existent file - let loaded = store.load_peers().await.expect("Failed to load peers from empty store"); - assert!(loaded.is_empty()); - } -} diff --git a/dash-spv/src/sync/block_headers/manager.rs b/dash-spv/src/sync/block_headers/manager.rs index b3699e979..2e8bf90c1 100644 --- a/dash-spv/src/sync/block_headers/manager.rs +++ b/dash-spv/src/sync/block_headers/manager.rs @@ -13,7 +13,7 @@ use std::time::Instant; use crate::chain::CheckpointManager; use crate::error::{SyncError, SyncResult}; -use crate::network::RequestSender; +use crate::network::{NetworkManager, RequestKey}; use crate::storage::{BlockHeaderStorage, BlockHeaderTip, MetadataStorage}; use crate::sync::block_headers::HeadersPipeline; use crate::sync::{BlockHeadersProgress, ProgressPercentage, SyncEvent, SyncManager, SyncState}; @@ -123,7 +123,7 @@ impl BlockHeadersManager { pub(super) async fn handle_headers_pipeline( &mut self, headers: &[Header], - requests: &RequestSender, + network: &Arc, ) -> SyncResult> { if !self.pipeline.is_initialized() { // Pipeline not initialized (shouldn't happen in normal flow) @@ -134,10 +134,26 @@ impl BlockHeadersManager { let was_syncing = self.state() == SyncState::Syncing; let tip_was_complete = self.pipeline.is_tip_complete(); + // Capture the `getheaders` locator this response answers before routing — + // routing advances the matched segment's `current_tip_hash`. A non-empty + // response is keyed by its first header's `prev_blockhash` (the locator we + // sent from); an empty response carries nothing to key on, so it answers + // the active tip segment's locator. + let answered_locator = match headers.first() { + Some(first) => Some(first.prev_blockhash), + None => self.pipeline.active_tip_locator(), + }; + // Route headers to the pipeline, validates checkpoint match. let matched = self.pipeline.receive_headers(headers)?; - if matched.is_none() && !headers.is_empty() { + // Correlated to a segment: tell the network manager the request is + // answered so it stops timing it out and frees the key for re-request. + if matched.is_some() { + if let Some(locator) = answered_locator { + network.request_answered(RequestKey::Headers(locator)).await; + } + } else if !headers.is_empty() { tracing::debug!( "Headers not matched by pipeline (prev_hash: {}), may be post-sync update", headers[0].prev_blockhash @@ -147,7 +163,7 @@ impl BlockHeadersManager { // Send more requests during initial sync or active post-sync catch-up. // Skip for unsolicited headers. if was_syncing || !tip_was_complete { - let sent = self.pipeline.send_pending(requests)?; + let sent = self.pipeline.send_pending(network).await?; if sent > 0 { tracing::debug!("Pipeline sent {} more requests", sent); } @@ -199,7 +215,7 @@ impl BlockHeadersManager { self.pending_announcements.len() ); self.pipeline.reset_tip_segment(); - self.pipeline.send_pending(requests)?; + self.pipeline.send_pending(network).await?; } else { // Synced to the tip and no pending announcements, finalize and emit event let tip = self.tip().await?; @@ -229,7 +245,7 @@ impl BlockHeadersManager { pub(super) async fn handle_inventory( &mut self, inv: &[Inventory], - _requests: &RequestSender, + _network: &Arc, ) -> SyncResult<()> { for inv_item in inv { if let Inventory::Block(block_hash) = inv_item { @@ -258,13 +274,13 @@ impl BlockHeadersManager { mod tests { use super::*; use crate::chain::checkpoints::testnet_checkpoints; - use crate::network::{MessageType, NetworkEvent, NetworkRequest, RequestSender}; + use crate::network::{MessageType, NetworkEvent}; use crate::storage::{ DiskStorageManager, PersistentBlockHeaderStorage, PersistentMetadataStorage, StorageManager, }; use crate::sync::{ManagerIdentifier, SyncManager, SyncManagerProgress}; + use crate::test_utils::{test_socket_address, MockNetworkManager}; use dashcore::network::message::NetworkMessage; - use tokio::sync::mpsc::unbounded_channel; type TestBlockHeadersManager = BlockHeadersManager; @@ -303,7 +319,7 @@ mod tests { let manager = create_test_manager().await; assert_eq!(manager.identifier(), ManagerIdentifier::BlockHeader); assert_eq!(manager.state(), SyncState::WaitingForConnections); - assert_eq!(manager.wanted_message_types(), vec![MessageType::Headers, MessageType::Inv]); + assert_eq!(manager.wanted_message_types(), [MessageType::Headers, MessageType::Inv]); } #[tokio::test] @@ -332,12 +348,6 @@ mod tests { assert_eq!(manager.pipeline.segment_count(), 0); } - fn create_test_request_sender( - ) -> (RequestSender, tokio::sync::mpsc::UnboundedReceiver) { - let (tx, rx) = unbounded_channel(); - (RequestSender::new(tx), rx) - } - #[tokio::test] async fn test_unsolicited_post_sync_header_does_not_trigger_get_headers() { let mut manager = create_test_manager().await; @@ -349,11 +359,12 @@ mod tests { manager.pipeline.mark_tip_complete(); manager.progress.set_state(SyncState::Synced); - let (sender, mut rx) = create_test_request_sender(); + let mock = Arc::new(MockNetworkManager::new()); + let network: Arc = mock.clone(); let header = Header::dummy_chain(1, tip_hash).remove(0); - let events = manager.handle_headers_pipeline(&[header], &sender).await.unwrap(); + let events = manager.handle_headers_pipeline(&[header], &network).await.unwrap(); // Header should have been stored assert_eq!(events.len(), 1); @@ -365,7 +376,8 @@ mod tests { )); // No GetHeaders request should have been sent - assert!(rx.try_recv().is_err()); + assert!(mock.sent_messages().is_empty()); + assert!(mock.sent_to_messages().is_empty()); // Tip segment marked complete again for the next unsolicited header assert!(manager.pipeline.is_tip_complete()); @@ -374,111 +386,114 @@ mod tests { #[tokio::test] async fn test_peer_tip_announcement_lifecycle() { let mut manager = create_synced_manager().await; - let (requests, mut rx) = create_test_request_sender(); + // An idle synced tip has no in-flight catch-up request. + manager.pipeline.mark_tip_complete(); + let mock = Arc::new(MockNetworkManager::new()); + let network: Arc = mock.clone(); - let addr: SocketAddr = "1.2.3.4:9999".parse().unwrap(); - let connect = NetworkEvent::PeerConnected { - address: addr, - }; + let addr = test_socket_address(1); + let connect = NetworkEvent::PeerConnected(addr); // Connect sends a peer-targeted GetHeaders - let events = manager.handle_network_event(&connect, &requests).await.unwrap(); + let events = manager.handle_network_event(&connect, &network).await.unwrap(); assert!(events.is_empty()); assert!(manager.announced_peers.contains(&addr)); - match rx.try_recv().unwrap() { - NetworkRequest::SendMessageToPeer(_, target_addr) => { - assert_eq!(target_addr, addr); - } - other => panic!("Expected SendMessageToPeer, got {:?}", other), - } + let sent_to = mock.sent_to_messages(); + assert_eq!(sent_to.len(), 1); + assert_eq!(sent_to[0].0, addr); + assert!(matches!(sent_to[0].1, NetworkMessage::GetHeaders(_))); // Same peer again sends nothing (already announced) - manager.handle_network_event(&connect, &requests).await.unwrap(); - assert!(rx.try_recv().is_err()); + manager.handle_network_event(&connect, &network).await.unwrap(); + assert_eq!(mock.sent_to_messages().len(), 1); // Disconnect removes from announced set - let disconnect = NetworkEvent::PeerDisconnected { - address: addr, - }; - manager.handle_network_event(&disconnect, &requests).await.unwrap(); + let disconnect = NetworkEvent::PeerDisconnected(addr); + manager.handle_network_event(&disconnect, &network).await.unwrap(); assert!(!manager.announced_peers.contains(&addr)); // Reconnect sends GetHeaders again - manager.handle_network_event(&connect, &requests).await.unwrap(); + manager.handle_network_event(&connect, &network).await.unwrap(); assert!(manager.announced_peers.contains(&addr)); - assert!(rx.try_recv().is_ok()); + assert_eq!(mock.sent_to_messages().len(), 2); } #[tokio::test] async fn test_peer_tip_announcement_guards() { // Not synced: peer connect does nothing let mut manager = create_test_manager().await; - let (requests, mut rx) = create_test_request_sender(); - let addr: SocketAddr = "1.2.3.4:9999".parse().unwrap(); - let connect = NetworkEvent::PeerConnected { - address: addr, - }; + let mock = Arc::new(MockNetworkManager::new()); + let network: Arc = mock.clone(); + let addr = test_socket_address(1); + let connect = NetworkEvent::PeerConnected(addr); - manager.handle_network_event(&connect, &requests).await.unwrap(); + manager.handle_network_event(&connect, &network).await.unwrap(); assert!(!manager.announced_peers.contains(&addr)); - assert!(rx.try_recv().is_err()); + assert!(mock.sent_to_messages().is_empty()); // Active catch-up: peer connect skipped while pipeline has pending request let mut manager = create_synced_manager().await; + let mock = Arc::new(MockNetworkManager::new()); + let network: Arc = mock.clone(); manager.pipeline.reset_tip_segment(); - manager.pipeline.send_pending(&requests).unwrap(); - rx.try_recv().unwrap(); // drain the pipeline GetHeaders + manager.pipeline.send_pending(&network).await.unwrap(); + // The pipeline GetHeaders is declared via `send`, not `send_to`. + assert!(!mock.sent_messages().is_empty()); - manager.handle_network_event(&connect, &requests).await.unwrap(); + manager.handle_network_event(&connect, &network).await.unwrap(); assert!(!manager.announced_peers.contains(&addr)); - assert!(rx.try_recv().is_err()); + // No peer-targeted announcement while a catch-up request is in flight. + assert!(mock.sent_to_messages().is_empty()); } #[tokio::test] async fn test_disconnect_preserves_pipeline_and_resumes_from_advanced_tip() { let mut manager = create_test_manager().await; - let (requests, mut rx) = create_test_request_sender(); + let mock = Arc::new(MockNetworkManager::new()); + let network: Arc = mock.clone(); // Use a target below the first testnet checkpoint (50000) so the // pipeline produces a single open-ended tip segment. let initial_event = NetworkEvent::PeersUpdated { connected_count: 1, - best_height: Some(40_000), - addresses: vec![], + best_height: 40_000, }; - manager.handle_network_event(&initial_event, &requests).await.unwrap(); + manager.handle_network_event(&initial_event, &network).await.unwrap(); assert_eq!(manager.state(), SyncState::Syncing); assert!(manager.pipeline.is_initialized()); assert_eq!(manager.pipeline.segment_count(), 1); - let initial_locator = match rx.try_recv().expect("initial GetHeaders not sent") { - NetworkRequest::SendMessage(NetworkMessage::GetHeaders(msg)) => msg.locator_hashes[0], + let sent = mock.sent_messages(); + assert_eq!(sent.len(), 1, "initial GetHeaders not sent"); + let initial_locator = match &sent[0] { + NetworkMessage::GetHeaders(msg) => msg.locator_hashes[0], other => panic!("Expected GetHeaders, got {:?}", other), }; - assert!(rx.try_recv().is_err()); + mock.clear_sent(); // Simulate a peer response. The single tip segment drains its buffer // through take_ready_to_store, advancing the storage tip and the // segment's current_tip_hash to advanced_hash. let header = Header::dummy_chain(1, initial_locator).remove(0); let advanced_hash = header.block_hash(); - manager.handle_headers_pipeline(&[header], &requests).await.unwrap(); + manager.handle_headers_pipeline(&[header], &network).await.unwrap(); // Drain the follow-up GetHeaders that send_pending issued. - match rx.try_recv().expect("follow-up GetHeaders not sent") { - NetworkRequest::SendMessage(NetworkMessage::GetHeaders(msg)) => { + let sent = mock.sent_messages(); + assert_eq!(sent.len(), 1, "follow-up GetHeaders not sent"); + match &sent[0] { + NetworkMessage::GetHeaders(msg) => { assert_eq!(msg.locator_hashes[0], advanced_hash); } other => panic!("Expected GetHeaders, got {:?}", other), } - assert!(rx.try_recv().is_err()); + mock.clear_sent(); let disconnect_event = NetworkEvent::PeersUpdated { connected_count: 0, - best_height: Some(40_000), - addresses: vec![], + best_height: 40_000, }; - manager.handle_network_event(&disconnect_event, &requests).await.unwrap(); + manager.handle_network_event(&disconnect_event, &network).await.unwrap(); assert_eq!(manager.state(), SyncState::WaitingForConnections); assert!( manager.pipeline.is_initialized(), @@ -488,11 +503,13 @@ mod tests { // Reconnect: start_sync must skip pipeline.init and resume by sending // GetHeaders from each segment's preserved current_tip_hash. - manager.handle_network_event(&initial_event, &requests).await.unwrap(); + manager.handle_network_event(&initial_event, &network).await.unwrap(); assert_eq!(manager.state(), SyncState::Syncing); - let resumed_locator = match rx.try_recv().expect("resumed GetHeaders not sent") { - NetworkRequest::SendMessage(NetworkMessage::GetHeaders(msg)) => msg.locator_hashes[0], + let sent = mock.sent_messages(); + assert_eq!(sent.len(), 1, "resumed GetHeaders not sent"); + let resumed_locator = match &sent[0] { + NetworkMessage::GetHeaders(msg) => msg.locator_hashes[0], other => panic!("Expected GetHeaders, got {:?}", other), }; assert_eq!( @@ -500,7 +517,6 @@ mod tests { "GetHeaders on reconnect must use the preserved current_tip_hash" ); assert_ne!(resumed_locator, initial_locator); - assert!(rx.try_recv().is_err()); } #[tokio::test] @@ -511,54 +527,55 @@ mod tests { manager.pipeline.mark_tip_complete(); assert!(manager.pipeline.is_tip_complete()); - let (requests, mut rx) = create_test_request_sender(); + let mock = Arc::new(MockNetworkManager::new()); + let network: Arc = mock.clone(); let disconnect_event = NetworkEvent::PeersUpdated { connected_count: 0, - best_height: Some(tip.height()), - addresses: vec![], + best_height: tip.height(), }; - manager.handle_network_event(&disconnect_event, &requests).await.unwrap(); + manager.handle_network_event(&disconnect_event, &network).await.unwrap(); assert_eq!(manager.state(), SyncState::WaitingForConnections); assert!(manager.pipeline.is_initialized()); // Reconnect with a higher peer best_height (a new block was mined). let reconnect_event = NetworkEvent::PeersUpdated { connected_count: 1, - best_height: Some(tip.height() + 1), - addresses: vec![], + best_height: tip.height() + 1, }; - manager.handle_network_event(&reconnect_event, &requests).await.unwrap(); + manager.handle_network_event(&reconnect_event, &network).await.unwrap(); assert_eq!(manager.state(), SyncState::Syncing); - let resumed_locator = match rx.try_recv().expect("resumed GetHeaders not sent") { - NetworkRequest::SendMessage(NetworkMessage::GetHeaders(msg)) => msg.locator_hashes[0], + let sent = mock.sent_messages(); + assert_eq!(sent.len(), 1, "resumed GetHeaders not sent"); + let resumed_locator = match &sent[0] { + NetworkMessage::GetHeaders(msg) => msg.locator_hashes[0], other => panic!("Expected GetHeaders, got {:?}", other), }; assert_eq!(resumed_locator, synced_hash); - assert!(rx.try_recv().is_err()); } #[tokio::test] async fn test_empty_headers_after_tip_announcement_is_harmless() { let mut manager = create_synced_manager().await; manager.pipeline.mark_tip_complete(); - let (requests, mut rx) = create_test_request_sender(); + let mock = Arc::new(MockNetworkManager::new()); + let network: Arc = mock.clone(); // Announce tip to a new peer - let addr: SocketAddr = "1.2.3.4:9999".parse().unwrap(); - let connect = NetworkEvent::PeerConnected { - address: addr, - }; - manager.handle_network_event(&connect, &requests).await.unwrap(); - rx.try_recv().unwrap(); // drain the GetHeaders request + let addr = test_socket_address(1); + let connect = NetworkEvent::PeerConnected(addr); + manager.handle_network_event(&connect, &network).await.unwrap(); + assert_eq!(mock.sent_to_messages().len(), 1); // the GetHeaders announcement + mock.clear_sent(); // Peer responds with empty headers (same height as us) - let events = manager.handle_headers_pipeline(&[], &requests).await.unwrap(); + let events = manager.handle_headers_pipeline(&[], &network).await.unwrap(); // No events emitted, no requests sent, tip segment stays complete assert!(events.is_empty()); - assert!(rx.try_recv().is_err()); + assert!(mock.sent_messages().is_empty()); + assert!(mock.sent_to_messages().is_empty()); assert!(manager.pipeline.is_tip_complete()); } } diff --git a/dash-spv/src/sync/block_headers/pipeline.rs b/dash-spv/src/sync/block_headers/pipeline.rs index cce5baca9..b6a3a453b 100644 --- a/dash-spv/src/sync/block_headers/pipeline.rs +++ b/dash-spv/src/sync/block_headers/pipeline.rs @@ -11,7 +11,7 @@ use dashcore::BlockHash; use crate::chain::CheckpointManager; use crate::error::SyncResult; -use crate::network::RequestSender; +use crate::network::NetworkManager; use crate::sync::block_headers::segment_state::SegmentState; use crate::types::HashedBlockHeader; @@ -117,19 +117,21 @@ impl HeadersPipeline { self.segments.len() } - /// Send pending requests for active segments. - /// Returns the number of requests sent. - pub fn send_pending(&mut self, requests: &RequestSender) -> SyncResult { + /// Declare each active segment's wanted `getheaders` to the network manager. + /// + /// Each non-complete segment wants exactly one locator (`current_tip_hash`). + /// The network manager de-duplicates re-declarations, so this is fired freely + /// (on start, on arrival, on tick) — the broker paces and retries. Returns the + /// number of segments whose want was declared. + pub async fn send_pending(&mut self, network: &Arc) -> SyncResult { let mut sent = 0; for segment in &mut self.segments { // Skip completed segments - if segment.complete { + if !segment.can_send() { continue; } - while segment.can_send() { - segment.send_request(requests)?; - sent += 1; - } + segment.send_request(network).await; + sent += 1; } Ok(sent) } @@ -143,10 +145,7 @@ impl HeadersPipeline { // Route to the tip segment (target_height is None) if it has in-flight requests. // Middle segments complete via checkpoint validation, not empty responses. for segment in &mut self.segments { - if !segment.complete - && segment.target_height.is_none() - && segment.coordinator.active_count() > 0 - { + if !segment.complete && segment.target_height.is_none() { tracing::debug!( "Routing empty response to tip segment {} at height {}", segment.segment_id, @@ -175,8 +174,6 @@ impl HeadersPipeline { if segment.complete && segment.target_height.is_none() { segment.complete = false; self.next_to_store = idx; - // Mark as in-flight so the coordinator accepts these unsolicited headers - segment.coordinator.mark_sent(&[prev_hash]); tracing::debug!( "Tip segment {} receiving post-sync headers, reset for continued processing", segment.segment_id @@ -261,23 +258,17 @@ impl HeadersPipeline { self.segments.iter().map(|s| s.buffered_headers.len() as u32).sum() } - /// Check for timeouts in all segments. - pub fn handle_timeouts(&mut self) { - for segment in &mut self.segments { - segment.handle_timeouts(); - } - } - - /// Drop only per-peer in-flight bookkeeping across every segment. + /// Locator of the active tip segment, if any. /// - /// Buffered headers, segment topology, and per-segment validated tip state - /// are preserved. `next_to_store` and `initialized` stay put so a reconnect - /// can resume sending `GetHeaders` from each segment's preserved - /// `current_tip_hash` without re-fetching what we already have. - pub fn clear_in_flight(&mut self) { - for segment in &mut self.segments { - segment.clear_in_flight(); - } + /// Used to correlate an empty `headers` response (which carries nothing to + /// key on) back to the `RequestKey::Headers(current_tip_hash)` it answers, + /// so the manager can clear it from the network manager. Returns the + /// `current_tip_hash` of the non-complete open-ended (tip) segment. + pub(super) fn active_tip_locator(&self) -> Option { + self.segments + .iter() + .find(|s| !s.complete && s.target_height.is_none()) + .map(|s| s.current_tip_hash) } /// Check if pipeline is initialized. @@ -326,12 +317,14 @@ impl HeadersPipeline { false } - /// Check if the tip segment has active requests in flight. + /// Check if the tip segment is actively catching up. + /// + /// A non-complete open-ended (tip) segment is declaring its `getheaders` each + /// tick, so it has a request in flight from the broker's point of view. Used + /// to avoid firing a redundant catch-up `getheaders` (or an empty-response + /// that would prematurely complete the tip segment). pub fn tip_segment_has_pending_request(&self) -> bool { - self.segments - .iter() - .find(|s| s.target_height.is_none()) - .is_some_and(|s| !s.complete && s.coordinator.active_count() > 0) + self.segments.iter().find(|s| s.target_height.is_none()).is_some_and(|s| !s.complete) } } @@ -339,9 +332,7 @@ impl HeadersPipeline { mod tests { use super::*; use crate::chain::checkpoints::{mainnet_checkpoints, testnet_checkpoints}; - use tokio::sync::mpsc::unbounded_channel; - use crate::network::{NetworkRequest, RequestSender}; use crate::sync::block_headers::segment_state::SegmentState; fn create_test_checkpoint_manager(is_testnet: bool) -> Arc { @@ -353,12 +344,6 @@ mod tests { Arc::new(CheckpointManager::new(checkpoints)) } - fn create_test_request_sender( - ) -> (RequestSender, tokio::sync::mpsc::UnboundedReceiver) { - let (tx, rx) = unbounded_channel(); - (RequestSender::new(tx), rx) - } - #[test] fn test_pipeline_new() { let cm = create_test_checkpoint_manager(true); @@ -397,29 +382,6 @@ mod tests { assert!(pipeline.segment_count() >= 2); } - #[test] - fn test_pipeline_send_pending() { - let cm = create_test_checkpoint_manager(true); - let mut pipeline = HeadersPipeline::new(cm.clone()); - - let genesis = cm.get_checkpoint(0).unwrap(); - pipeline.init(0, genesis.block_hash, 1_200_000); - - let (sender, mut rx) = create_test_request_sender(); - - let sent = pipeline.send_pending(&sender).unwrap(); - - // Should send at least one request per segment - assert!(sent >= pipeline.segment_count()); - - // Verify messages were queued - let mut count = 0; - while rx.try_recv().is_ok() { - count += 1; - } - assert_eq!(count, sent); - } - #[test] fn test_pipeline_is_complete_initially() { let cm = create_test_checkpoint_manager(true); @@ -498,9 +460,6 @@ mod tests { let mut header = Header::dummy(1); header.prev_blockhash = shared_hash; - // Mark segment 1 request as in-flight so receive works - pipeline.segments[1].coordinator.mark_sent(&[shared_hash]); - // Route headers should go to segment 1, not the completed segment 0 let matched = pipeline.receive_headers(&[header]).unwrap(); assert_eq!(matched, Some(1), "Headers should route to segment 1, not completed segment 0"); @@ -536,7 +495,10 @@ mod tests { } #[test] - fn test_clear_in_flight_preserves_buffers_across_segments() { + fn test_disconnect_preserves_segment_chain_state() { + // On disconnect the network manager re-queues in-flight requests; the + // pipeline keeps every segment's validated chain state so a reconnect + // resumes from each `current_tip_hash` without re-fetching what we have. let shared_hash = BlockHash::dummy(42); let mut completed = @@ -544,25 +506,20 @@ mod tests { completed.complete = true; completed.current_height = 100; completed.current_tip_hash = shared_hash; - // Buffered headers on a complete-but-not-yet-drained segment must survive. let mut completed_header = Header::dummy(1); completed_header.prev_blockhash = BlockHash::dummy(0); completed.buffered_headers.push(HashedBlockHeader::from(completed_header)); let mut mid = SegmentState::new(1, 100, shared_hash, Some(200), None); - mid.coordinator.mark_sent(&[shared_hash]); let mut mid_header = Header::dummy(2); mid_header.prev_blockhash = shared_hash; mid.receive_headers(&[mid_header]).unwrap(); let mid_preserved_tip = mid.current_tip_hash; let mid_preserved_height = mid.current_height; let mid_preserved_buffered = mid.buffered_headers.len(); - // Simulate a fresh in-flight follow-up request for this segment. - mid.coordinator.mark_sent(&[mid_preserved_tip]); let tip_hash = BlockHash::dummy(99); - let mut tip = SegmentState::new(2, 500, tip_hash, None, None); - tip.coordinator.mark_sent(&[tip_hash]); + let tip = SegmentState::new(2, 500, tip_hash, None, None); let cm = create_test_checkpoint_manager(true); let mut pipeline = HeadersPipeline::new(cm); @@ -570,9 +527,7 @@ mod tests { pipeline.next_to_store = 0; pipeline.segments = vec![completed, mid, tip]; - pipeline.clear_in_flight(); - - // initialized and next_to_store stay put. + // initialized and next_to_store stay put across a disconnect. assert!(pipeline.is_initialized()); assert_eq!(pipeline.next_to_store, 0); @@ -581,19 +536,16 @@ mod tests { assert_eq!(pipeline.segments[0].buffered_headers.len(), 1); assert_eq!(pipeline.segments[0].current_tip_hash, shared_hash); - // Mid-download segment: validated chain state preserved; coordinator wiped. + // Mid-download segment: validated chain state preserved. assert_eq!(pipeline.segments[1].current_tip_hash, mid_preserved_tip); assert_eq!(pipeline.segments[1].current_height, mid_preserved_height); assert_eq!(pipeline.segments[1].buffered_headers.len(), mid_preserved_buffered); assert!(!pipeline.segments[1].complete); - assert_eq!(pipeline.segments[1].coordinator.active_count(), 0); - assert_eq!(pipeline.segments[1].coordinator.pending_count(), 0); // can_send returns true so a fresh GetHeaders can resume from preserved tip. assert!(pipeline.segments[1].can_send()); - // Tip segment: in-flight cleared, preserved hash/height intact. + // Tip segment: preserved hash intact, still wants its locator. assert_eq!(pipeline.segments[2].current_tip_hash, tip_hash); - assert_eq!(pipeline.segments[2].coordinator.active_count(), 0); assert!(pipeline.segments[2].can_send()); } diff --git a/dash-spv/src/sync/block_headers/segment_state.rs b/dash-spv/src/sync/block_headers/segment_state.rs index 9829d1250..e870b1011 100644 --- a/dash-spv/src/sync/block_headers/segment_state.rs +++ b/dash-spv/src/sync/block_headers/segment_state.rs @@ -1,14 +1,18 @@ use crate::error::{SyncError, SyncResult}; -use crate::network::RequestSender; -use crate::sync::download_coordinator::{DownloadConfig, DownloadCoordinator}; +use crate::network::NetworkManager; use crate::types::HashedBlockHeader; +use dashcore::network::message::NetworkMessage; +use dashcore::network::message_blockdata::GetHeadersMessage; use dashcore::{BlockHash, Header}; -use std::time::Duration; - -/// Timeout for header requests. -const HEADERS_TIMEOUT: Duration = Duration::from_secs(30); +use dashcore_hashes::Hash; +use std::sync::Arc; /// State for a single download segment between two checkpoints. +/// +/// The segment declares the single `getheaders` it wants (a locator from its +/// `current_tip_hash`) to the network manager, which de-duplicates, paces, times +/// out and retries it. The segment keeps no in-flight bookkeeping of its own: it +/// simply wants `current_tip_hash` for as long as it is not `complete`. #[derive(Debug)] pub(super) struct SegmentState { /// Unique segment identifier (index in segments array). @@ -23,8 +27,6 @@ pub(super) struct SegmentState { pub(super) current_tip_hash: BlockHash, /// Current height reached in this segment. pub(super) current_height: u32, - /// Download coordinator for tracking in-flight requests. - pub(super) coordinator: DownloadCoordinator, /// Buffered headers waiting to be stored. pub(super) buffered_headers: Vec, /// Whether this segment has completed downloading. @@ -47,33 +49,38 @@ impl SegmentState { target_hash, current_tip_hash: start_hash, current_height: start_height, - coordinator: DownloadCoordinator::new( - DownloadConfig::default() - .with_max_concurrent(1) // Only 1 request at a time (sequential getheaders) - .with_timeout(HEADERS_TIMEOUT), - ), buffered_headers: Vec::new(), complete: false, } } - /// Check if the segment can send more requests. - /// Only one getheaders request can be in-flight at a time (sequential protocol). + /// Check if the segment still wants a `getheaders`. + /// + /// A segment wants its `current_tip_hash` locator declared for as long as it + /// is not complete. The network manager de-duplicates re-declarations, so the + /// pipeline can (re-)declare each tick without tracking what is on the wire. pub(super) fn can_send(&self) -> bool { - !self.complete && !self.coordinator.is_in_flight(&self.current_tip_hash) + !self.complete } - /// Send a GetHeaders request for this segment. - pub(super) fn send_request(&mut self, requests: &RequestSender) -> SyncResult<()> { - requests.request_block_headers(self.current_tip_hash)?; - self.coordinator.mark_sent(&[self.current_tip_hash]); + /// Declare this segment's `getheaders` to the network manager. + /// + /// The broker de-duplicates by `RequestKey::Headers(current_tip_hash)`, paces + /// the request across peers and retries it on timeout, so this may be a no-op + /// if the request is already in play. + pub(super) async fn send_request(&mut self, network: &Arc) { + network + .send(NetworkMessage::GetHeaders(GetHeadersMessage::new( + vec![self.current_tip_hash], + BlockHash::all_zeros(), + ))) + .await; tracing::debug!( - "Segment {}: sent GetHeaders from height {} hash {}", + "Segment {}: declared GetHeaders from height {} hash {}", self.segment_id, self.current_height, self.current_tip_hash ); - Ok(()) } /// Try to match incoming headers to this segment. @@ -89,8 +96,6 @@ impl SegmentState { if headers.is_empty() { // Empty response means we've reached the peer's tip for this segment self.complete = true; - // Clear in-flight tracking for the current tip hash - self.coordinator.receive(&self.current_tip_hash); tracing::info!( "Segment {}: complete (empty response at height {})", self.segment_id, @@ -109,15 +114,6 @@ impl SegmentState { ))); } - // Mark the request as received, reject if we never requested this hash - let prev_hash = headers[0].prev_blockhash; - if !self.coordinator.receive(&prev_hash) { - return Err(SyncError::InvalidState(format!( - "Segment {}: received unrequested headers (prev_hash {})", - self.segment_id, prev_hash - ))); - } - // Process headers let mut processed = 0; for header in headers { @@ -180,29 +176,6 @@ impl SegmentState { pub(super) fn take_buffered(&mut self) -> Vec { std::mem::take(&mut self.buffered_headers) } - - /// Check for timed out requests and handle retries. - pub(super) fn handle_timeouts(&mut self) { - let timed_out = self.coordinator.check_timeouts(); - for hash in timed_out { - tracing::warn!( - "Segment {}: request timed out for hash {}, will retry", - self.segment_id, - hash - ); - // Re-enqueue for retry - self.coordinator.enqueue_retry(hash); - } - } - - /// Drop only per-peer in-flight bookkeeping. - /// - /// Buffered headers and the validated `current_tip_hash` / `current_height` - /// are preserved so a reconnect can resume from where the last peer left off - /// without re-fetching headers we already have. - pub(super) fn clear_in_flight(&mut self) { - self.coordinator.clear(); - } } #[cfg(test)] @@ -229,6 +202,7 @@ mod tests { let hash = BlockHash::dummy(0); let segment = SegmentState::new(0, 0, hash, Some(1000), None); + // A fresh, incomplete segment wants its locator declared. assert!(segment.can_send()); } @@ -250,13 +224,14 @@ mod tests { assert_eq!(processed, 0); assert!(segment.complete); + // An empty response no longer allows sending — the segment is done. + assert!(!segment.can_send()); } #[test] fn test_segment_receive_headers() { let hash = BlockHash::dummy(1); let mut segment = SegmentState::new(0, 0, hash, None, None); - segment.coordinator.mark_sent(&[hash]); // Create dummy headers that chain from all-zeros let headers: Vec
= (1..=10).map(Header::dummy).collect(); @@ -271,6 +246,8 @@ mod tests { assert_eq!(segment.buffered_headers.len(), 1); assert_eq!(segment.current_height, 1); assert!(!segment.complete); + // The tip advanced to the received header's hash for the next locator. + assert_eq!(segment.current_tip_hash, first.block_hash()); } #[test] @@ -280,7 +257,6 @@ mod tests { let expected_checkpoint_hash = BlockHash::dummy(99); let mut segment = SegmentState::new(0, 0, start_hash, Some(1), Some(expected_checkpoint_hash)); - segment.coordinator.mark_sent(&[start_hash]); // Create a header that will be at height 1 but with a different hash let mut header = Header::dummy(1); @@ -320,7 +296,6 @@ mod tests { // Create segment with checkpoint matching the header's hash let mut segment = SegmentState::new(0, 0, start_hash, Some(1), Some(header_hash)); - segment.coordinator.mark_sent(&[start_hash]); // Receiving this header should succeed and complete the segment let result = segment.receive_headers(&[header]); @@ -332,25 +307,6 @@ mod tests { assert_eq!(segment.buffered_headers.len(), 1); } - #[test] - fn test_unrequested_headers_returns_error() { - let start_hash = BlockHash::dummy(0); - let mut segment = SegmentState::new(0, 0, start_hash, None, None); - - let mut header = Header::dummy(1); - header.prev_blockhash = start_hash; - - let result = segment.receive_headers(&[header]); - assert!(result.is_err()); - match result.unwrap_err() { - SyncError::InvalidState(msg) => { - assert!(msg.contains("unrequested headers")); - } - other => panic!("Expected SyncError::InvalidState, got {:?}", other), - } - assert!(segment.buffered_headers.is_empty()); - } - #[test] fn test_completed_segment_rejects_new_headers() { let start_hash = BlockHash::dummy(0); @@ -375,41 +331,4 @@ mod tests { } assert!(segment.buffered_headers.is_empty()); } - - #[test] - fn test_clear_in_flight_preserves_chain_state() { - let start_hash = BlockHash::dummy(0); - let mut segment = SegmentState::new(0, 0, start_hash, None, None); - segment.coordinator.mark_sent(&[start_hash]); - - let mut header = Header::dummy(1); - header.prev_blockhash = start_hash; - segment.receive_headers(&[header]).unwrap(); - - let preserved_tip_hash = segment.current_tip_hash; - let preserved_height = segment.current_height; - let preserved_buffered = segment.buffered_headers.len(); - assert_ne!(preserved_tip_hash, start_hash); - assert_eq!(preserved_height, 1); - assert_eq!(preserved_buffered, 1); - - // Simulate a fresh in-flight request, then clear it. - segment.coordinator.mark_sent(&[preserved_tip_hash]); - assert!(segment.coordinator.is_in_flight(&preserved_tip_hash)); - - segment.clear_in_flight(); - - assert!(!segment.coordinator.is_in_flight(&preserved_tip_hash)); - assert_eq!(segment.coordinator.active_count(), 0); - assert_eq!(segment.coordinator.pending_count(), 0); - - assert_eq!(segment.current_tip_hash, preserved_tip_hash); - assert_eq!(segment.current_height, preserved_height); - assert_eq!(segment.buffered_headers.len(), preserved_buffered); - assert!(!segment.complete); - - // After clearing, can_send should be true again so a fresh GetHeaders - // can resume from the preserved tip hash without re-fetching what we have. - assert!(segment.can_send()); - } } diff --git a/dash-spv/src/sync/block_headers/sync_manager.rs b/dash-spv/src/sync/block_headers/sync_manager.rs index d47e2cfe8..b1dceb85b 100644 --- a/dash-spv/src/sync/block_headers/sync_manager.rs +++ b/dash-spv/src/sync/block_headers/sync_manager.rs @@ -1,5 +1,5 @@ use crate::error::SyncResult; -use crate::network::{Message, MessageType, NetworkEvent, RequestSender}; +use crate::network::{MessageType, NetworkEvent, NetworkManager}; use crate::storage::{BlockHeaderStorage, MetadataStorage}; use crate::sync::sync_manager::ensure_not_started; use crate::sync::{ @@ -8,7 +8,11 @@ use crate::sync::{ }; use async_trait::async_trait; use dashcore::network::message::NetworkMessage; +use dashcore::network::message_blockdata::GetHeadersMessage; use dashcore::BlockHash; +use dashcore_hashes::Hash; +use std::net::SocketAddr; +use std::sync::Arc; use std::time::{Duration, Instant}; /// Timeout waiting for unsolicited header messages after a block announcement. @@ -37,17 +41,20 @@ impl SyncManager for BlockHeadersMana } fn on_disconnect(&mut self) { - // Drop only per-peer in-flight bookkeeping. Segment topology and - // validated chain state per segment (current_tip_hash, current_height, - // buffered_headers, complete) are preserved so a reconnect can resume - // from where the disconnected peer left off without re-fetching headers - // we already have. - self.pipeline.clear_in_flight(); + // The network manager re-queues in-flight requests itself and paces the + // re-declared `getheaders` to the new peer. Segment topology and validated + // chain state per segment (current_tip_hash, current_height, + // buffered_headers, complete) are preserved so a reconnect can resume from + // where the disconnected peer left off without re-fetching headers we + // already have. Only the peer-scoped announcement bookkeeping is dropped. self.pending_announcements.clear(); self.announced_peers.clear(); } - async fn start_sync(&mut self, requests: &RequestSender) -> SyncResult> { + async fn start_sync( + &mut self, + network: &Arc, + ) -> SyncResult> { ensure_not_started(self.state(), self.identifier())?; self.progress.set_state(SyncState::Syncing); @@ -78,7 +85,7 @@ impl SyncManager for BlockHeadersMana } // Send initial batch of requests - let sent = self.pipeline.send_pending(requests)?; + let sent = self.pipeline.send_pending(network).await?; tracing::info!("Pipeline: sent {} initial requests", sent); Ok(vec![SyncEvent::SyncStart { @@ -88,17 +95,18 @@ impl SyncManager for BlockHeadersMana async fn handle_message( &mut self, - msg: Message, - requests: &RequestSender, + _peer: SocketAddr, + msg: NetworkMessage, + network: &Arc, ) -> SyncResult> { - match msg.inner() { + match &msg { NetworkMessage::Headers(headers) => { // Always route through pipeline when initialized - self.handle_headers_pipeline(headers, requests).await + self.handle_headers_pipeline(headers, network).await } NetworkMessage::Inv(inv) => { - self.handle_inventory(inv, requests).await?; + self.handle_inventory(inv, network).await?; Ok(vec![]) } @@ -109,26 +117,18 @@ impl SyncManager for BlockHeadersMana async fn handle_sync_event( &mut self, _event: &SyncEvent, - _requests: &RequestSender, + _network: &Arc, ) -> SyncResult> { // BlockHeadersManager doesn't react to events from other managers Ok(vec![]) } - async fn tick(&mut self, requests: &RequestSender) -> SyncResult> { + async fn tick(&mut self, network: &Arc) -> SyncResult> { if !self.pipeline.is_initialized() { return Ok(vec![]); } - self.pipeline.handle_timeouts(); - - // During initial sync, send more requests and log progress if self.state() == SyncState::Syncing { - let sent = self.pipeline.send_pending(requests)?; - if sent > 0 { - tracing::debug!("Tick: pipeline sent {} more requests", sent); - } - return Ok(vec![]); } @@ -152,7 +152,7 @@ impl SyncManager for BlockHeadersMana // Reset tip segment and send requests via pipeline self.pipeline.reset_tip_segment(); - self.pipeline.send_pending(requests)?; + self.pipeline.send_pending(network).await?; for hash in stale { self.pending_announcements.remove(&hash); @@ -166,12 +166,10 @@ impl SyncManager for BlockHeadersMana async fn handle_network_event( &mut self, event: &NetworkEvent, - requests: &RequestSender, + network: &Arc, ) -> SyncResult> { match event { - NetworkEvent::PeerConnected { - address, - } => { + NetworkEvent::PeerConnected(address) => { // When synced, send GetHeaders to new peers so Dash Core learns our tip // and sends header announcements instead of inv. Skip when the // pipeline has an active catch-up request to avoid the empty @@ -183,47 +181,51 @@ impl SyncManager for BlockHeadersMana { let tip = self.tip().await?; tracing::info!("Announcing tip {} to new peer {}", tip.height(), address); - requests.request_block_headers_from_peer(*tip.hash(), *address)?; + // Peer-pinned send: this must reach THIS peer so it learns our + // tip, not whichever peer the router would otherwise pick. + network + .send_to( + *address, + NetworkMessage::GetHeaders(GetHeadersMessage::new( + vec![*tip.hash()], + BlockHash::all_zeros(), + )), + ) + .await; self.announced_peers.insert(*address); } } - NetworkEvent::PeerDisconnected { - address, - } => { + NetworkEvent::PeerDisconnected(address) => { self.announced_peers.remove(address); } NetworkEvent::PeersUpdated { connected_count, best_height, - .. } => { - if let Some(best_height) = best_height { - self.progress.update_target_height(*best_height); + self.progress.update_target_height(*best_height); + { let mut metadata_storage = self.metadata_storage.write().await; metadata_storage.store_last_target_height(*best_height).await?; } if *connected_count == 0 { self.stop_sync(); - } else if *connected_count > 0 { + } else { if self.state() == SyncState::WaitingForConnections { - return self.start_sync(requests).await; + return self.start_sync(network).await; } // When already synced but behind peer height, request missing headers - if self.state() == SyncState::Synced { - if let Some(best_height) = best_height { - if *best_height > self.progress.tip_height() - && !self.pipeline.tip_segment_has_pending_request() - { - tracing::info!( - "Peer height {} > our height {}, requesting headers to catch up", - best_height, - self.progress.tip_height() - ); - // Reset tip segment and send requests via pipeline - self.pipeline.reset_tip_segment(); - self.pipeline.send_pending(requests)?; - } - } + if self.state() == SyncState::Synced + && *best_height > self.progress.tip_height() + && !self.pipeline.tip_segment_has_pending_request() + { + tracing::info!( + "Peer height {} > our height {}, requesting headers to catch up", + best_height, + self.progress.tip_height() + ); + // Reset tip segment and send requests via pipeline + self.pipeline.reset_tip_segment(); + self.pipeline.send_pending(network).await?; } } } diff --git a/dash-spv/src/sync/blocks/manager.rs b/dash-spv/src/sync/blocks/manager.rs index 70df8624b..8b5a8a0b8 100644 --- a/dash-spv/src/sync/blocks/manager.rs +++ b/dash-spv/src/sync/blocks/manager.rs @@ -9,7 +9,7 @@ use tokio::sync::RwLock; use super::pipeline::BlocksPipeline; use crate::error::SyncResult; -use crate::network::RequestSender; +use crate::network::NetworkManager; use crate::storage::{BlockHeaderStorage, BlockStorage}; use crate::sync::{BlocksProgress, SyncEvent, SyncManager, SyncState}; use key_wallet_manager::WalletInterface; @@ -63,11 +63,11 @@ impl BlocksManager SyncResult<()> { - let sent = self.pipeline.send_pending(requests).await?; - if sent > 0 { - self.progress.add_requested(sent as u32); - } + pub(super) async fn send_pending( + &mut self, + network: &Arc, + ) -> SyncResult<()> { + self.pipeline.send_pending(network).await?; Ok(()) } @@ -167,16 +167,14 @@ impl std::fmt::Debug #[cfg(test)] mod tests { use super::*; - use crate::network::{MessageType, NetworkManager}; + use crate::network::MessageType; use crate::storage::{ DiskStorageManager, PersistentBlockHeaderStorage, PersistentBlockStorage, StorageManager, }; use crate::sync::{ManagerIdentifier, SyncEvent, SyncManagerProgress}; - use crate::test_utils::MockNetworkManager; use crate::types::HashedBlock; use key_wallet_manager::test_utils::{MockWallet, MOCK_WALLET_ID}; - use key_wallet_manager::FilterMatchKey; - use std::collections::{BTreeMap, BTreeSet}; + use std::collections::BTreeSet; type TestBlocksManager = BlocksManager; @@ -193,7 +191,7 @@ mod tests { let manager = create_test_manager().await; assert_eq!(manager.identifier(), ManagerIdentifier::Block); assert_eq!(manager.state(), SyncState::WaitForEvents); - assert_eq!(manager.wanted_message_types(), vec![MessageType::Block]); + assert_eq!(manager.wanted_message_types(), [MessageType::Block]); } #[tokio::test] @@ -214,24 +212,37 @@ mod tests { #[tokio::test] async fn test_blocks_manager_handle_blocks_needed_event() { + use crate::network::NetworkManager; + use crate::test_utils::MockNetworkManager; + use key_wallet_manager::FilterMatchKey; + use std::collections::BTreeMap; + let mut manager = create_test_manager().await; manager.progress.set_state(SyncState::Synced); - let network = MockNetworkManager::new(); - let requests = network.request_sender(); + let mock = Arc::new(MockNetworkManager::new()); + let network: Arc = mock.clone(); - let block_hash = dashcore::BlockHash::dummy(0); + // Block is not in storage, so the manager queues it for download and + // declares a getdata to the broker. + let block_hash = dashcore::block::Header::dummy(0).block_hash(); let mut blocks = BTreeMap::new(); blocks.insert(FilterMatchKey::new(100, block_hash), BTreeSet::from([MOCK_WALLET_ID])); let event = SyncEvent::BlocksNeeded { blocks, }; - let events = manager.handle_sync_event(&event, &requests).await.unwrap(); + let events = manager.handle_sync_event(&event, &network).await.unwrap(); - // Should queue the block + // Should queue the block and transition to Syncing. assert_eq!(manager.state(), SyncState::Syncing); assert!(events.is_empty()); + + // The queued block was declared to the network as a getdata request. + assert!( + !mock.sent_messages().is_empty(), + "expected a getdata to be declared for the needed block" + ); } /// `process_buffered_blocks` must call `process_block_for_wallets` with diff --git a/dash-spv/src/sync/blocks/pipeline.rs b/dash-spv/src/sync/blocks/pipeline.rs index 02b29aac2..97a484c3c 100644 --- a/dash-spv/src/sync/blocks/pipeline.rs +++ b/dash-spv/src/sync/blocks/pipeline.rs @@ -1,40 +1,33 @@ //! Blocks pipeline implementation. //! -//! Handles concurrent block downloads with timeout and retry logic. -//! Uses the generic DownloadCoordinator for core mechanics. +//! Declares wanted blocks to the network manager (the broker) and buffers the +//! arrivals for height-ordered processing. The broker owns pacing, timeouts and +//! retries — this pipeline keeps no in-flight queue of its own. use std::collections::{BTreeMap, BTreeSet, HashMap}; -use std::time::Duration; +use std::sync::Arc; use crate::error::SyncResult; -use crate::network::RequestSender; -use crate::sync::download_coordinator::{DownloadConfig, DownloadCoordinator}; +use crate::network::NetworkManager; use crate::types::HashedBlock; +use dashcore::network::message::NetworkMessage; +use dashcore::network::message_blockdata::Inventory; use dashcore::BlockHash; use key_wallet_manager::{FilterMatchKey, WalletId}; -/// Maximum number of concurrent block downloads. -const MAX_CONCURRENT_BLOCK_DOWNLOADS: usize = 20; - -/// Timeout for block downloads before retry. -const BLOCK_TIMEOUT: Duration = Duration::from_secs(30); - -/// Maximum blocks per GetData request, kept a bit lower for better download distribution to multiple peers -const BLOCKS_PER_REQUEST: usize = 8; - /// Pipeline for downloading blocks with height-ordered processing. /// -/// Uses DownloadCoordinator for core download mechanics. -/// This is a thin wrapper that handles building GetData inventory messages. -/// Tracks block heights to enable ordered processing and buffers downloaded blocks. +/// Holds no request queue of its own: it declares the blocks it wants to the +/// network manager (the broker de-duplicates, paces, times out and retries), and +/// buffers the arrivals for height-ordered processing. A block is "wanted" for +/// exactly as long as it sits in `hash_to_height`. pub(super) struct BlocksPipeline { - /// Core download coordinator (handles pending, in-flight, timeouts). - coordinator: DownloadCoordinator, - /// Heights queued or in-flight (waiting for download). + /// Heights still wanted (block requested, not yet downloaded). pending_heights: BTreeSet, /// Downloaded blocks ready to process (height -> block, with its cached hash). downloaded: BTreeMap, - /// Map hash -> height for looking up height when block arrives. + /// Wanted blocks: hash -> height. A block leaves this map once downloaded. + /// Doubles as the "is this block wanted?" set for validating arrivals. hash_to_height: HashMap, /// Per-block interested wallets, populated when the block is queued. /// Only those wallets get the block processed. @@ -44,9 +37,9 @@ pub(super) struct BlocksPipeline { impl std::fmt::Debug for BlocksPipeline { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("BlocksPipeline") - .field("coordinator", &self.coordinator) .field("pending_heights", &self.pending_heights.len()) .field("downloaded", &self.downloaded.len()) + .field("wanted", &self.hash_to_height.len()) .finish() } } @@ -61,11 +54,6 @@ impl BlocksPipeline { /// Create a new blocks pipeline. pub(super) fn new() -> Self { Self { - coordinator: DownloadCoordinator::new( - DownloadConfig::default() - .with_max_concurrent(MAX_CONCURRENT_BLOCK_DOWNLOADS) - .with_timeout(BLOCK_TIMEOUT), - ), pending_heights: BTreeSet::new(), downloaded: BTreeMap::new(), hash_to_height: HashMap::new(), @@ -74,81 +62,81 @@ impl BlocksPipeline { } /// Queue blocks with their heights and per-block interested wallet sets. + /// + /// Returns how many were not already tracked, which is what "blocks requested" + /// means for progress: a block enters the wanted set once, however many times it + /// is later re-declared to the broker. pub(super) fn queue( &mut self, blocks: impl IntoIterator)>, - ) { + ) -> usize { + let mut newly_tracked = 0; for (key, wallets) in blocks { let hash = *key.hash(); let already_tracked = self.hash_to_height.contains_key(&hash) || self.hash_to_wallets.contains_key(&hash); if !already_tracked { - self.coordinator.enqueue([hash]); self.pending_heights.insert(key.height()); self.hash_to_height.insert(hash, key.height()); + newly_tracked += 1; } self.hash_to_wallets.entry(hash).or_default().extend(wallets); } + newly_tracked } /// Check if the pipeline has completed all work. /// - /// Returns true when no blocks are pending, downloading, or waiting to be processed. + /// Returns true when no blocks are wanted, downloading, or waiting to be processed. pub(super) fn is_complete(&self) -> bool { - self.coordinator.is_empty() && self.downloaded.is_empty() && self.pending_heights.is_empty() + self.hash_to_height.is_empty() + && self.downloaded.is_empty() + && self.pending_heights.is_empty() } - /// Check if there are pending requests to make. + /// Check if there are blocks still to download. pub(super) fn has_pending_requests(&self) -> bool { - self.coordinator.available_to_send() > 0 + !self.hash_to_height.is_empty() } - /// Send pending block requests up to the concurrency limit. + /// Declare every wanted block to the network manager. /// - /// Sends multiple smaller GetData messages to distribute requests across peers. - /// Returns the number of blocks requested. - pub(super) async fn send_pending(&mut self, requests: &RequestSender) -> SyncResult { - let mut total_sent = 0; - - while self.coordinator.available_to_send() > 0 { - // Take a batch of up to BLOCKS_PER_REQUEST - let count = self.coordinator.available_to_send().min(BLOCKS_PER_REQUEST); - let hashes = self.coordinator.take_pending(count); - if hashes.is_empty() { - break; - } - - requests.request_blocks(hashes.clone())?; - self.coordinator.mark_sent(&hashes); - total_sent += hashes.len(); - - tracing::debug!( - "Requested {} blocks ({} downloading, {} pending)", - hashes.len(), - self.coordinator.active_count(), - self.coordinator.pending_count() - ); + /// Fired freely (on queue, on arrival, on tick): the broker de-duplicates, so + /// re-declaring a block already queued or on the wire is a no-op, and it owns + /// pacing (one `getdata` per block, throttled by each peer's measured capacity) + /// and retry (re-inject on timeout after dropping the dead peer). Re-declaring + /// each tick is the safety net if a peer drops before the broker retries. + /// + /// Returns the number of blocks declared (offered, not necessarily newly sent). + pub(super) async fn send_pending( + &mut self, + network: &Arc, + ) -> SyncResult { + if self.hash_to_height.is_empty() { + return Ok(0); } - - Ok(total_sent) + let hashes: Vec = self.hash_to_height.keys().copied().collect(); + for hash in &hashes { + network.send(NetworkMessage::GetData(vec![Inventory::Block(*hash)])).await; + } + tracing::debug!("Declared {} wanted block(s) to the broker", hashes.len()); + Ok(hashes.len()) } /// Handle a received block using internal height mapping. /// - /// Looks up the height from the internal hash_to_height map and stores - /// the block in the downloaded buffer for height-ordered processing. - /// Returns `true` if this was a tracked block, `false` if unrequested. + /// Looks up the height from the internal `hash_to_height` map and stores the + /// block in the downloaded buffer for height-ordered processing. + /// Returns `true` if this was a wanted block, `false` if unrequested. pub(super) fn receive_block(&mut self, block: &HashedBlock) -> bool { let hash = *block.hash(); - if !self.coordinator.receive(&hash) { + // Not in the wanted set => unrequested or already downloaded; ignore. + let Some(height) = self.hash_to_height.remove(&hash) else { tracing::debug!("Ignoring unrequested block: {}", hash); return false; - } - - if let Some(height) = self.hash_to_height.remove(&hash) { - self.pending_heights.remove(&height); - self.downloaded.insert(height, block.clone()); - } + }; + self.pending_heights.remove(&height); + self.downloaded.insert(height, block.clone()); true } @@ -188,20 +176,6 @@ impl BlocksPipeline { self.hash_to_wallets.entry(hash).or_default().extend(wallets); self.downloaded.insert(height, block); } - - /// Check for timed out downloads and re-queue them. - pub(super) fn handle_timeouts(&mut self) { - self.coordinator.check_and_retry_timeouts(); - } - - /// Move in-flight `getdata` requests back to pending after a peer - /// disconnect so the next `send_pending` reissues them to the new peer. - /// `pending_heights`, `downloaded`, `hash_to_height`, and `hash_to_wallets` - /// are preserved so already-received blocks are not re-fetched and the - /// per-block wallet routing stays intact. - pub(super) fn requeue_in_flight(&mut self) { - self.coordinator.requeue_in_flight(); - } } #[cfg(test)] @@ -211,10 +185,6 @@ mod tests { use super::*; - fn test_hash(n: u8) -> BlockHash { - BlockHash::from_byte_array([n; 32]) - } - fn make_test_block(n: u8) -> Block { use dashcore::blockdata::block::Header; let header = Header { @@ -234,8 +204,7 @@ mod tests { #[test] fn test_blocks_pipeline_new() { let pipeline = BlocksPipeline::new(); - assert_eq!(pipeline.coordinator.pending_count(), 0); - assert_eq!(pipeline.coordinator.active_count(), 0); + assert!(pipeline.hash_to_height.is_empty()); assert!(pipeline.is_complete()); } @@ -245,7 +214,7 @@ mod tests { let block = make_test_block(1); pipeline.queue([(FilterMatchKey::new(100, block.block_hash()), BTreeSet::new())]); - assert_eq!(pipeline.coordinator.pending_count(), 1); + assert_eq!(pipeline.hash_to_height.len(), 1); assert!(!pipeline.is_complete()); assert!(pipeline.has_pending_requests()); } @@ -262,7 +231,7 @@ mod tests { (FilterMatchKey::new(102, block3.block_hash()), BTreeSet::new()), ]); - assert_eq!(pipeline.coordinator.pending_count(), 3); + assert_eq!(pipeline.hash_to_height.len(), 3); assert_eq!(pipeline.pending_heights.len(), 3); assert!(pipeline.pending_heights.contains(&100)); assert!(pipeline.pending_heights.contains(&101)); @@ -275,17 +244,10 @@ mod tests { let block = make_test_block(1); let hash = block.block_hash(); - // Queue with height tracking pipeline.queue([(FilterMatchKey::new(100, block.block_hash()), BTreeSet::new())]); - // Simulate sending via coordinator - let hashes = pipeline.coordinator.take_pending(1); - pipeline.coordinator.mark_sent(&hashes); - assert_eq!(pipeline.coordinator.active_count(), 1); - - // Receive block assert!(pipeline.receive_block(&HashedBlock::from(&block))); - assert_eq!(pipeline.coordinator.active_count(), 0); + assert!(pipeline.hash_to_height.is_empty()); assert_eq!(pipeline.downloaded.len(), 1); assert!(pipeline.pending_heights.is_empty()); assert_eq!(*pipeline.downloaded.get(&100).unwrap().hash(), hash); @@ -300,84 +262,6 @@ mod tests { assert!(pipeline.downloaded.is_empty()); } - #[test] - fn test_max_concurrent() { - let mut pipeline = BlocksPipeline::new(); - - // Queue more blocks than max concurrent - for i in 0..=MAX_CONCURRENT_BLOCK_DOWNLOADS { - let block = make_test_block(i as u8); - pipeline.queue([(FilterMatchKey::new(i as u32, block.block_hash()), BTreeSet::new())]); - } - - // Take and mark as downloading up to limit - let to_send = pipeline.coordinator.available_to_send(); - let hashes = pipeline.coordinator.take_pending(to_send); - pipeline.coordinator.mark_sent(&hashes); - - assert_eq!(pipeline.coordinator.active_count(), MAX_CONCURRENT_BLOCK_DOWNLOADS); - assert_eq!(pipeline.coordinator.pending_count(), 1); - assert!(!pipeline.has_pending_requests()); - } - - #[test] - fn test_requeue_in_flight_preserves_downloaded_and_pending_heights() { - let mut pipeline = BlocksPipeline::new(); - let block_a = make_test_block(1); - let block_b = make_test_block(2); - let hash_a = block_a.block_hash(); - let hash_b = block_b.block_hash(); - - // A: queued and sent — will be requeued. - pipeline.queue([(FilterMatchKey::new(100, hash_a), BTreeSet::from([[1u8; 32]]))]); - let sent = pipeline.coordinator.take_pending(1); - pipeline.coordinator.mark_sent(&sent); - assert_eq!(pipeline.coordinator.active_count(), 1); - - // B: already received, sitting in `downloaded` — must survive requeue. - pipeline.add_from_storage(HashedBlock::from(&block_b), 200, BTreeSet::from([[2u8; 32]])); - - pipeline.requeue_in_flight(); - - assert_eq!(pipeline.coordinator.active_count(), 0); - assert_eq!(pipeline.coordinator.pending_count(), 1); - assert!(pipeline.pending_heights.contains(&100)); - assert_eq!(pipeline.hash_to_height.get(&hash_a), Some(&100)); - assert!(pipeline.hash_to_wallets.contains_key(&hash_a)); - assert!(pipeline.downloaded.contains_key(&200)); - assert!(pipeline.hash_to_wallets.contains_key(&hash_b)); - } - - #[test] - fn test_timeout_requeues() { - // Create pipeline with very short timeout for testing - let mut pipeline = BlocksPipeline { - coordinator: DownloadCoordinator::new( - DownloadConfig::default() - .with_max_concurrent(MAX_CONCURRENT_BLOCK_DOWNLOADS) - .with_timeout(Duration::from_millis(10)), - ), - pending_heights: BTreeSet::new(), - downloaded: BTreeMap::new(), - hash_to_height: HashMap::new(), - hash_to_wallets: HashMap::new(), - }; - - // Use coordinator directly to set up in-flight state - let hash = test_hash(1); - pipeline.coordinator.enqueue([hash]); - let hashes = pipeline.coordinator.take_pending(1); - pipeline.coordinator.mark_sent(&hashes); - - // Wait for timeout - std::thread::sleep(Duration::from_millis(20)); - - pipeline.handle_timeouts(); - - assert_eq!(pipeline.coordinator.active_count(), 0); - assert_eq!(pipeline.coordinator.pending_count(), 1); - } - #[test] fn test_take_next_ordered_block_in_order() { let mut pipeline = BlocksPipeline::new(); @@ -386,30 +270,23 @@ mod tests { let hash1 = block1.block_hash(); let hash2 = block2.block_hash(); - // Use add_from_storage to test ordering logic without network - // Add block 2 first (out of order) pipeline.add_from_storage(HashedBlock::from(&block2), 101, BTreeSet::new()); - // Also track height 100 as pending to simulate waiting pipeline.pending_heights.insert(100); // Cannot take block 2 yet - waiting for block at height 100 assert!(pipeline.take_next_ordered_block().is_none()); - // Add block 1 pipeline.pending_heights.remove(&100); pipeline.add_from_storage(HashedBlock::from(&block1), 100, BTreeSet::new()); - // Now block 1 is ready (lowest height) let (block, height, _) = pipeline.take_next_ordered_block().unwrap(); assert_eq!(height, 100); assert_eq!(*block.hash(), hash1); - // Block 2 is now ready let (block, height, _) = pipeline.take_next_ordered_block().unwrap(); assert_eq!(height, 101); assert_eq!(*block.hash(), hash2); - // No more blocks assert!(pipeline.take_next_ordered_block().is_none()); } @@ -418,17 +295,13 @@ mod tests { let mut pipeline = BlocksPipeline::new(); let block2 = make_test_block(2); - // Add block at height 101, but height 100 is still pending pipeline.pending_heights.insert(100); pipeline.add_from_storage(HashedBlock::from(&block2), 101, BTreeSet::new()); - // Cannot take block 2 - block at height 100 is still pending assert!(pipeline.take_next_ordered_block().is_none()); - // Clear the pending height pipeline.pending_heights.remove(&100); - // Now block 2 is ready let (_, height, _) = pipeline.take_next_ordered_block().unwrap(); assert_eq!(height, 101); } @@ -440,7 +313,6 @@ mod tests { let hash = block.block_hash(); pipeline.add_from_storage(HashedBlock::from(&block), 100, BTreeSet::new()); - assert_eq!(pipeline.downloaded.len(), 1); let (taken_block, height, _) = pipeline.take_next_ordered_block().unwrap(); @@ -453,12 +325,10 @@ mod tests { let mut pipeline = BlocksPipeline::new(); assert!(pipeline.is_complete()); - // Adding to downloaded makes it incomplete let block = make_test_block(1); pipeline.add_from_storage(HashedBlock::from(&block), 100, BTreeSet::new()); assert!(!pipeline.is_complete()); - // Take the block pipeline.take_next_ordered_block(); assert!(pipeline.is_complete()); } @@ -468,7 +338,6 @@ mod tests { let mut pipeline = BlocksPipeline::new(); assert!(pipeline.is_complete()); - // Pending heights make it incomplete pipeline.pending_heights.insert(100); assert!(!pipeline.is_complete()); @@ -478,18 +347,12 @@ mod tests { #[test] fn test_queue_propagates_wallet_set_through_take_next() { - // A block queued with a non-empty wallet set must yield that exact - // wallet set when taken in height order via `take_next_ordered_block`. let mut pipeline = BlocksPipeline::new(); let block = make_test_block(1); let hash = block.block_hash(); let wallets: BTreeSet = BTreeSet::from([[1u8; 32], [2u8; 32]]); pipeline.queue([(FilterMatchKey::new(100, hash), wallets.clone())]); - - // Drive the block through receive_block to land it in `downloaded`. - let hashes = pipeline.coordinator.take_pending(1); - pipeline.coordinator.mark_sent(&hashes); assert!(pipeline.receive_block(&HashedBlock::from(&block))); let (taken_block, height, taken_wallets) = pipeline.take_next_ordered_block().unwrap(); @@ -500,9 +363,6 @@ mod tests { #[test] fn test_queue_merges_wallet_sets_for_repeat_hashes() { - // Queueing the same block hash twice with different wallet sets must - // produce the union when the block is later taken from the pipeline, - // and must not double-count it in the coordinator's pending state. let mut pipeline = BlocksPipeline::new(); let block = make_test_block(1); let hash = block.block_hash(); @@ -510,15 +370,11 @@ mod tests { let wallets_b: BTreeSet = BTreeSet::from([[2u8; 32], [3u8; 32]]); pipeline.queue([(FilterMatchKey::new(100, hash), wallets_a.clone())]); - assert_eq!(pipeline.coordinator.pending_count(), 1); + assert_eq!(pipeline.hash_to_height.len(), 1); pipeline.queue([(FilterMatchKey::new(100, hash), wallets_b.clone())]); - // Re-queueing must not double the coordinator's pending count. - assert_eq!(pipeline.coordinator.pending_count(), 1); + // Re-queueing must not double the wanted count. + assert_eq!(pipeline.hash_to_height.len(), 1); - // Land the block in `downloaded` to retrieve it. - let hashes = pipeline.coordinator.take_pending(1); - assert_eq!(hashes.len(), 1); - pipeline.coordinator.mark_sent(&hashes); assert!(pipeline.receive_block(&HashedBlock::from(&block))); let (_, _, taken_wallets) = pipeline.take_next_ordered_block().unwrap(); @@ -527,44 +383,8 @@ mod tests { assert_eq!(taken_wallets, expected); } - #[test] - fn test_queue_does_not_re_enqueue_in_flight_hash() { - // A late-arriving wallet match for a block already in flight must - // merge the wallet id without re-enqueueing the hash. Re-enqueueing - // would cause a duplicate request and corrupt the coordinator's - // pending/in-flight state. - let mut pipeline = BlocksPipeline::new(); - let block = make_test_block(1); - let hash = block.block_hash(); - let wallets_a: BTreeSet = BTreeSet::from([[1u8; 32]]); - let wallets_b: BTreeSet = BTreeSet::from([[2u8; 32]]); - - pipeline.queue([(FilterMatchKey::new(100, hash), wallets_a.clone())]); - // Move the hash to in-flight. - let hashes = pipeline.coordinator.take_pending(1); - pipeline.coordinator.mark_sent(&hashes); - assert_eq!(pipeline.coordinator.pending_count(), 0); - assert_eq!(pipeline.coordinator.active_count(), 1); - - // A second queue call for the same hash must not push it back to - // pending while it is in flight. - pipeline.queue([(FilterMatchKey::new(100, hash), wallets_b.clone())]); - assert_eq!(pipeline.coordinator.pending_count(), 0); - assert_eq!(pipeline.coordinator.active_count(), 1); - - // Late wallet ids are still merged for when the block arrives. - assert!(pipeline.receive_block(&HashedBlock::from(&block))); - let (_, _, taken_wallets) = pipeline.take_next_ordered_block().unwrap(); - let mut expected = wallets_a; - expected.extend(wallets_b); - assert_eq!(taken_wallets, expected); - } - #[test] fn test_queue_does_not_re_enqueue_downloaded_hash() { - // A late-arriving wallet match for a block already received and sitting - // in `downloaded` (but not yet consumed by `take_next_ordered_block`) - // must merge the wallet id without re-enqueueing the hash. let mut pipeline = BlocksPipeline::new(); let block = make_test_block(1); let hash = block.block_hash(); @@ -572,20 +392,15 @@ mod tests { let wallets_b: BTreeSet = BTreeSet::from([[2u8; 32]]); pipeline.queue([(FilterMatchKey::new(100, hash), wallets_a.clone())]); - let hashes = pipeline.coordinator.take_pending(1); - pipeline.coordinator.mark_sent(&hashes); assert!(pipeline.receive_block(&HashedBlock::from(&block))); assert_eq!(pipeline.downloaded.len(), 1); - assert_eq!(pipeline.coordinator.pending_count(), 0); - assert_eq!(pipeline.coordinator.active_count(), 0); + assert!(pipeline.hash_to_height.is_empty()); // Late-arriving match for the same hash must not re-enqueue. pipeline.queue([(FilterMatchKey::new(100, hash), wallets_b.clone())]); - assert_eq!(pipeline.coordinator.pending_count(), 0); - assert_eq!(pipeline.coordinator.active_count(), 0); + assert!(pipeline.hash_to_height.is_empty()); assert_eq!(pipeline.downloaded.len(), 1); - // Late wallet ids are still merged for when the block is taken. let (_, _, taken_wallets) = pipeline.take_next_ordered_block().unwrap(); let mut expected = wallets_a; expected.extend(wallets_b); @@ -594,8 +409,6 @@ mod tests { #[test] fn test_add_from_storage_merges_wallet_sets() { - // The `add_from_storage` path must merge wallet sets for repeat - // additions of the same block hash, matching `queue`'s semantics. let mut pipeline = BlocksPipeline::new(); let block = make_test_block(1); let wallets_a: BTreeSet = BTreeSet::from([[1u8; 32]]); @@ -615,19 +428,14 @@ mod tests { let mut pipeline = BlocksPipeline::new(); let block = make_test_block(1); - // Queue and mark as sent via coordinator pipeline.queue([(FilterMatchKey::new(100, block.block_hash()), BTreeSet::new())]); - let hashes = pipeline.coordinator.take_pending(1); - pipeline.coordinator.mark_sent(&hashes); - // First receive - let result = pipeline.receive_block(&HashedBlock::from(&block)); - assert!(result); + // First receive returns the height. + assert!(pipeline.receive_block(&HashedBlock::from(&block))); assert_eq!(pipeline.downloaded.len(), 1); - // Duplicate receive (not tracked anymore since already completed) - let result = pipeline.receive_block(&HashedBlock::from(&block)); - assert!(!result); + // Duplicate receive: no longer wanted. + assert!(!pipeline.receive_block(&HashedBlock::from(&block))); assert_eq!(pipeline.downloaded.len(), 1); } } diff --git a/dash-spv/src/sync/blocks/sync_manager.rs b/dash-spv/src/sync/blocks/sync_manager.rs index e7ecbc68a..d6f992e3a 100644 --- a/dash-spv/src/sync/blocks/sync_manager.rs +++ b/dash-spv/src/sync/blocks/sync_manager.rs @@ -1,5 +1,5 @@ use crate::error::SyncResult; -use crate::network::{Message, MessageType, RequestSender}; +use crate::network::{MessageType, NetworkManager, RequestKey}; use crate::storage::{BlockHeaderStorage, BlockStorage}; use crate::sync::sync_manager::ensure_not_started; use crate::sync::{ @@ -11,6 +11,8 @@ use async_trait::async_trait; use dashcore::network::message::NetworkMessage; use key_wallet_manager::{FilterMatchKey, WalletId, WalletInterface}; use std::collections::BTreeSet; +use std::net::SocketAddr; +use std::sync::Arc; #[async_trait] impl SyncManager @@ -32,7 +34,10 @@ impl SyncM &[MessageType::Block] } - async fn start_sync(&mut self, _requests: &RequestSender) -> SyncResult> { + async fn start_sync( + &mut self, + _network: &Arc, + ) -> SyncResult> { ensure_not_started(self.state(), self.identifier())?; // Check if filters already completed (event received before start_sync) if self.filters_sync_complete && self.pipeline.is_complete() { @@ -52,22 +57,22 @@ impl SyncM Ok(vec![]) } - /// Keep the entire pipeline (downloaded blocks, pending queue, per-block - /// wallet routing) and the `filters_sync_complete` flag, and move in-flight - /// `getdata`s back to the front of `pending` so the next `send_pending` - /// reissues them to the new peer immediately. Without this preservation, - /// `FiltersManager`'s tracker would re-track the same block hashes after a - /// re-scan and leak `pending_blocks` counters that never reach zero. - fn on_disconnect(&mut self) { - self.pipeline.requeue_in_flight(); - } + /// Keep the entire pipeline (downloaded blocks, wanted set, per-block wallet + /// routing) and the `filters_sync_complete` flag across a peer disconnect. + /// In-flight `getdata`s are re-queued by the network manager itself, and the + /// pipeline's wanted set is preserved so `send_pending` reissues them to the + /// new peer. Without this preservation, `FiltersManager`'s tracker would + /// re-track the same block hashes after a re-scan and leak `pending_blocks` + /// counters that never reach zero. + fn on_disconnect(&mut self) {} async fn handle_message( &mut self, - msg: Message, - requests: &RequestSender, + _peer: SocketAddr, + msg: NetworkMessage, + network: &Arc, ) -> SyncResult> { - let NetworkMessage::Block(block) = msg.inner() else { + let NetworkMessage::Block(block) = &msg else { return Ok(vec![]); }; @@ -79,6 +84,10 @@ impl SyncM return Ok(vec![]); } + // Response correlated: tell the network manager to stop tracking this + // request for timeout/retry. + network.request_answered(RequestKey::Block(*hashed_block.hash())).await; + // Look up height for storage let height = self .header_storage @@ -100,20 +109,16 @@ impl SyncM self.progress.add_downloaded(1); - // Process buffered blocks - let events = self.process_buffered_blocks().await?; - - if self.pipeline.has_pending_requests() { - self.send_pending(requests).await?; - } - - Ok(events) + // Process buffered blocks. No `send_pending` here: the wanted blocks are + // already declared to the broker, which paces them out as capacity frees. + // New work is declared on `BlocksNeeded` and topped up on tick. + self.process_buffered_blocks().await } async fn handle_sync_event( &mut self, event: &SyncEvent, - requests: &RequestSender, + network: &Arc, ) -> SyncResult> { // React to BlocksNeeded events if let SyncEvent::BlocksNeeded { @@ -157,13 +162,14 @@ impl SyncM drop(block_storage); // Queue all blocks that need downloading - self.pipeline.queue(to_download); + let newly_wanted = self.pipeline.queue(to_download); + self.progress.add_requested(newly_wanted as u32); self.progress.set_state(SyncState::Syncing); - // Send batched request for blocks not in storage + // Declare blocks not in storage to the broker. if self.pipeline.has_pending_requests() { - self.send_pending(requests).await?; + self.send_pending(network).await?; } // Process any blocks we loaded from storage @@ -192,13 +198,7 @@ impl SyncM Ok(vec![]) } - async fn tick(&mut self, requests: &RequestSender) -> SyncResult> { - // Handle timeouts - self.pipeline.handle_timeouts(); - - self.send_pending(requests).await?; - - // Try to process any buffered blocks + async fn tick(&mut self, _network: &Arc) -> SyncResult> { self.process_buffered_blocks().await } diff --git a/dash-spv/src/sync/chainlock/manager.rs b/dash-spv/src/sync/chainlock/manager.rs index c211919df..0e75e292b 100644 --- a/dash-spv/src/sync/chainlock/manager.rs +++ b/dash-spv/src/sync/chainlock/manager.rs @@ -333,7 +333,7 @@ mod tests { let manager = create_test_manager().await; assert_eq!(manager.identifier(), ManagerIdentifier::ChainLock); assert_eq!(manager.state(), SyncState::WaitForEvents); - assert_eq!(manager.wanted_message_types(), vec![MessageType::CLSig, MessageType::Inv]); + assert_eq!(manager.wanted_message_types(), [MessageType::ChainLock, MessageType::Inv]); } /// Buffered `MasternodeStateUpdated` events delivered during @@ -342,9 +342,9 @@ mod tests { /// sync cycle after reconnect, so dropping it here is safe. #[tokio::test] async fn test_handle_sync_event_drops_masternode_state_updated_in_waiting_for_connections() { - use crate::network::RequestSender; + use crate::network::NetworkManager; use crate::sync::SyncEvent; - use tokio::sync::mpsc::unbounded_channel; + use crate::test_utils::MockNetworkManager; let mut manager = create_test_manager().await; manager.set_state(SyncState::WaitingForConnections); @@ -353,8 +353,8 @@ mod tests { height: 100, qr_info_result: None, }; - let (tx, _rx) = unbounded_channel(); - let events = manager.handle_sync_event(&event, &RequestSender::new(tx)).await.unwrap(); + let network: Arc = Arc::new(MockNetworkManager::new()); + let events = manager.handle_sync_event(&event, &network).await.unwrap(); assert!(events.is_empty()); assert_eq!(manager.state(), SyncState::WaitingForConnections); diff --git a/dash-spv/src/sync/chainlock/sync_manager.rs b/dash-spv/src/sync/chainlock/sync_manager.rs index b750410a5..b0f81900e 100644 --- a/dash-spv/src/sync/chainlock/sync_manager.rs +++ b/dash-spv/src/sync/chainlock/sync_manager.rs @@ -1,5 +1,5 @@ use crate::error::SyncResult; -use crate::network::{Message, MessageType, RequestSender}; +use crate::network::{MessageType, NetworkManager}; use crate::storage::{BlockHeaderStorage, MetadataStorage}; use crate::sync::{ ChainLockManager, ManagerIdentifier, SyncEvent, SyncManager, SyncManagerProgress, SyncState, @@ -7,6 +7,8 @@ use crate::sync::{ use async_trait::async_trait; use dashcore::network::message::NetworkMessage; use dashcore::network::message_blockdata::Inventory; +use std::net::SocketAddr; +use std::sync::Arc; #[async_trait] impl SyncManager for ChainLockManager { @@ -23,7 +25,7 @@ impl SyncManager for ChainLockManager } fn wanted_message_types(&self) -> &'static [MessageType] { - &[MessageType::CLSig, MessageType::Inv] + &[MessageType::ChainLock, MessageType::Inv] } fn on_disconnect(&mut self) { @@ -33,10 +35,11 @@ impl SyncManager for ChainLockManager async fn handle_message( &mut self, - msg: Message, - requests: &RequestSender, + peer: SocketAddr, + msg: NetworkMessage, + network: &Arc, ) -> SyncResult> { - match msg.inner() { + match &msg { NetworkMessage::CLSig(chainlock) => self.process_chainlock(chainlock).await, NetworkMessage::Inv(inv) => { // Check for ChainLock inventory items, filtering out already-requested ones @@ -58,12 +61,14 @@ impl SyncManager for ChainLockManager "Received {} ChainLock announcements, requesting via getdata", chainlocks_to_request.len() ); - requests - .request_inventory(chainlocks_to_request.clone(), msg.peer_address())?; - - for item in &chainlocks_to_request { - if let Inventory::ChainLock(hash) = item { - self.requested_chainlocks.insert(*hash); + if network + .send_to(peer, NetworkMessage::GetData(chainlocks_to_request.clone())) + .await + { + for item in &chainlocks_to_request { + if let Inventory::ChainLock(hash) = item { + self.requested_chainlocks.insert(*hash); + } } } } @@ -76,7 +81,7 @@ impl SyncManager for ChainLockManager async fn handle_sync_event( &mut self, event: &SyncEvent, - _requests: &RequestSender, + _network: &Arc, ) -> SyncResult> { // `MasternodeStateUpdated` fires on every MnListDiff / QRInfo // update; the work below is strictly one-shot startup work, so @@ -112,11 +117,6 @@ impl SyncManager for ChainLockManager Ok(vec![]) } - async fn tick(&mut self, _requests: &RequestSender) -> SyncResult> { - // No periodic work needed - Ok(vec![]) - } - fn progress(&self) -> SyncManagerProgress { SyncManagerProgress::ChainLock(self.progress.clone()) } diff --git a/dash-spv/src/sync/download_coordinator.rs b/dash-spv/src/sync/download_coordinator.rs deleted file mode 100644 index e36753b6d..000000000 --- a/dash-spv/src/sync/download_coordinator.rs +++ /dev/null @@ -1,465 +0,0 @@ -//! Generic download coordinator for pipelined downloads. -//! -//! Provides a single abstraction for managing concurrent downloads with: -//! - Pending queue management -//! - In-flight tracking with timestamps -//! - Timeout detection and retry logic -//! - Configurable concurrency limits - -use std::collections::{HashMap, VecDeque}; -use std::hash::Hash; -use std::time::{Duration, Instant}; - -/// Configuration for download coordination. -#[derive(Debug, Clone)] -pub struct DownloadConfig { - /// Maximum concurrent in-flight requests. - max_concurrent: usize, - /// Timeout duration for requests. - timeout: Duration, -} - -impl Default for DownloadConfig { - fn default() -> Self { - Self { - max_concurrent: 10, - timeout: Duration::from_secs(30), - } - } -} - -impl DownloadConfig { - /// Create config with custom max concurrent. - pub(crate) fn with_max_concurrent(mut self, max: usize) -> Self { - self.max_concurrent = max; - self - } - - /// Create config with custom timeout. - pub(crate) fn with_timeout(mut self, timeout: Duration) -> Self { - self.timeout = timeout; - self - } -} - -/// Generic download coordinator. -/// -/// Handles the common mechanics of pipelined downloads: -/// - Queue management (pending items) -/// - In-flight tracking with timestamps -/// - Timeout detection and retry -/// - Concurrency limits -/// -/// Generic over the key type `K` which identifies download items. -/// Use `u32` for height-based downloads, `BlockHash` for hash-based. -#[derive(Debug)] -pub(crate) struct DownloadCoordinator { - /// Items waiting to be requested. - pending: VecDeque, - /// Items currently in-flight (key -> sent time). - in_flight: HashMap, - /// Retry counts per key. - retry_counts: HashMap, - /// Configuration. - config: DownloadConfig, - /// Last time progress was made. - last_progress: Instant, -} - -impl Default for DownloadCoordinator { - fn default() -> Self { - Self::new(DownloadConfig::default()) - } -} - -impl DownloadCoordinator { - /// Create a new coordinator with the given configuration. - pub(crate) fn new(config: DownloadConfig) -> Self { - Self { - pending: VecDeque::new(), - in_flight: HashMap::new(), - retry_counts: HashMap::new(), - config, - last_progress: Instant::now(), - } - } - - /// Clear all state. - pub(crate) fn clear(&mut self) { - self.pending.clear(); - self.in_flight.clear(); - self.retry_counts.clear(); - self.last_progress = Instant::now(); - } - - /// Move all in-flight items back to the front of the pending queue. - /// - /// Used on peer disconnect: the requests went to a now-dead peer, but the - /// items themselves are still wanted. Retry counts are preserved so a peer - /// that consistently fails to deliver an item still trips the normal retry - /// budget. Without this hook, items would only be retried once their - /// timeout elapsed. - pub(crate) fn requeue_in_flight(&mut self) { - let items: Vec = self.in_flight.drain().map(|(k, _)| k).collect(); - if items.is_empty() { - return; - } - for item in items.into_iter().rev() { - self.pending.push_front(item); - } - } - - /// Queue items for download. - pub(crate) fn enqueue(&mut self, items: impl IntoIterator) { - for item in items { - self.pending.push_back(item); - } - } - - /// Queue an item for retry (goes to front of queue). - pub(crate) fn enqueue_retry(&mut self, item: K) { - let count = self.retry_counts.entry(item.clone()).or_insert(0); - *count += 1; - tracing::warn!("Retrying item (attempt {})", count); - self.pending.push_front(item); - } - - /// Get the number of items available to send (respecting concurrency limit). - pub(crate) fn available_to_send(&self) -> usize { - self.config.max_concurrent.saturating_sub(self.in_flight.len()).min(self.pending.len()) - } - - /// Take items from the pending queue (up to count). - /// - /// Items are removed from pending but NOT yet marked as in-flight. - /// Call `mark_sent` after successfully sending the request. - pub(crate) fn take_pending(&mut self, count: usize) -> Vec { - let actual = count.min(self.pending.len()); - let mut items = Vec::with_capacity(actual); - for _ in 0..actual { - if let Some(item) = self.pending.pop_front() { - items.push(item); - } - } - items - } - - /// Mark items as sent (now in-flight). - pub(crate) fn mark_sent(&mut self, items: &[K]) { - let now = Instant::now(); - for item in items { - self.in_flight.insert(item.clone(), now); - } - } - - /// Handle a received item. - /// - /// Returns true if the item was being tracked, false if unexpected. - pub(crate) fn receive(&mut self, key: &K) -> bool { - if self.in_flight.remove(key).is_some() { - self.retry_counts.remove(key); - self.last_progress = Instant::now(); - true - } else { - false - } - } - - /// Drop a key from the pending queue without touching in-flight state. - /// - /// Used when a pending item is satisfied through a side channel: a late - /// response from a disconnected peer can complete a batch that - /// `requeue_in_flight` just moved from in-flight back to pending. Without - /// this hook, the key would stay in `pending` with no tracker, and the - /// next `take_pending` would resurrect a finished batch. - pub(crate) fn cancel_pending(&mut self, key: &K) { - self.pending.retain(|k| k != key); - self.retry_counts.remove(key); - } - - /// Check if an item is currently in-flight. - pub(crate) fn is_in_flight(&self, key: &K) -> bool { - self.in_flight.contains_key(key) - } - - /// Check for timed-out items. - /// - /// Returns items that have timed out. They are removed from in-flight tracking. - /// Caller should call `enqueue_retry` for items that should be retried. - pub(crate) fn check_timeouts(&mut self) -> Vec { - let now = Instant::now(); - let timed_out: Vec = self - .in_flight - .iter() - .filter(|(_, sent_time)| now.duration_since(**sent_time) > self.config.timeout) - .map(|(key, _)| key.clone()) - .collect(); - - for key in &timed_out { - self.in_flight.remove(key); - } - - if !timed_out.is_empty() { - tracing::debug!("{} items timed out after {:?}", timed_out.len(), self.config.timeout); - } - - timed_out - } - - /// Check for timed-out items and re-enqueue them for retry. - /// - /// Combines `check_timeouts()` and `enqueue_retry()` in one call. - /// Returns all timed-out items that were re-queued. - pub(crate) fn check_and_retry_timeouts(&mut self) -> Vec { - let timed_out = self.check_timeouts(); - for item in &timed_out { - self.enqueue_retry(item.clone()); - } - timed_out - } - - /// Check if the coordinator has no work (empty pending and in-flight). - pub(crate) fn is_empty(&self) -> bool { - self.pending.is_empty() && self.in_flight.is_empty() - } - - /// Get the number of pending items. - pub(crate) fn pending_count(&self) -> usize { - self.pending.len() - } - - /// Get the number of in-flight items. - pub(crate) fn active_count(&self) -> usize { - self.in_flight.len() - } - - /// Get the total remaining items (pending + in-flight). - pub(crate) fn remaining(&self) -> usize { - self.pending.len() + self.in_flight.len() - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_new_coordinator() { - let coord: DownloadCoordinator = DownloadCoordinator::default(); - assert!(coord.is_empty()); - assert_eq!(coord.pending_count(), 0); - assert_eq!(coord.active_count(), 0); - } - - #[test] - fn test_enqueue() { - let mut coord: DownloadCoordinator = DownloadCoordinator::default(); - coord.enqueue([1, 2, 3, 4, 5]); - - assert_eq!(coord.pending_count(), 5); - } - - #[test] - fn test_enqueue_retry_goes_to_front() { - let mut coord: DownloadCoordinator = DownloadCoordinator::default(); - coord.enqueue([1, 2]); - coord.enqueue_retry(99); - - let items = coord.take_pending(3); - assert_eq!(items, vec![99, 1, 2]); - } - - #[test] - fn test_take_pending() { - let mut coord: DownloadCoordinator = DownloadCoordinator::default(); - coord.enqueue([1, 2, 3, 4, 5]); - - let items = coord.take_pending(3); - assert_eq!(items, vec![1, 2, 3]); - assert_eq!(coord.pending_count(), 2); - } - - #[test] - fn test_mark_sent() { - let mut coord: DownloadCoordinator = DownloadCoordinator::default(); - coord.enqueue([1, 2, 3]); - - let items = coord.take_pending(2); - coord.mark_sent(&items); - - assert_eq!(coord.pending_count(), 1); - assert_eq!(coord.active_count(), 2); - assert!(coord.is_in_flight(&1)); - assert!(coord.is_in_flight(&2)); - assert!(!coord.is_in_flight(&3)); - } - - #[test] - fn test_receive() { - let mut coord: DownloadCoordinator = DownloadCoordinator::default(); - coord.mark_sent(&[1]); - coord.mark_sent(&[2]); - - assert!(coord.receive(&1)); - assert_eq!(coord.active_count(), 1); - - assert!(!coord.receive(&99)); // Not tracked - assert_eq!(coord.active_count(), 1); - } - - #[test] - fn test_available_to_send() { - let mut coord: DownloadCoordinator = - DownloadCoordinator::new(DownloadConfig::default().with_max_concurrent(3)); - - coord.enqueue([1, 2, 3, 4, 5]); - assert_eq!(coord.available_to_send(), 3); - - coord.mark_sent(&[1]); - coord.mark_sent(&[2]); - assert_eq!(coord.available_to_send(), 1); - - coord.mark_sent(&[3]); - assert_eq!(coord.available_to_send(), 0); - } - - #[test] - fn test_check_timeouts() { - let mut coord: DownloadCoordinator = DownloadCoordinator::new( - DownloadConfig::default().with_timeout(Duration::from_millis(10)), - ); - - coord.mark_sent(&[1]); - coord.mark_sent(&[2]); - - // Immediately, nothing timed out - let timed_out = coord.check_timeouts(); - assert!(timed_out.is_empty()); - - // Wait for timeout - std::thread::sleep(Duration::from_millis(20)); - - let timed_out = coord.check_timeouts(); - assert_eq!(timed_out.len(), 2); - assert!(coord.in_flight.is_empty()); - } - - #[test] - fn test_requeue_in_flight_moves_items_to_pending_front() { - let mut coord: DownloadCoordinator = DownloadCoordinator::default(); - coord.enqueue([10, 11]); - coord.mark_sent(&[1, 2, 3]); - - coord.requeue_in_flight(); - - assert_eq!(coord.active_count(), 0); - // Requeued items go to the front, original pending follows. - let items = coord.take_pending(5); - assert_eq!(items.len(), 5); - assert_eq!(&items[3..], &[10, 11]); - let mut requeued = items[..3].to_vec(); - requeued.sort(); - assert_eq!(requeued, vec![1, 2, 3]); - } - - #[test] - fn test_requeue_in_flight_preserves_retry_counts() { - let mut coord: DownloadCoordinator = DownloadCoordinator::default(); - coord.enqueue_retry(7); - let items = coord.take_pending(1); - coord.mark_sent(&items); - assert_eq!(coord.retry_counts.get(&7), Some(&1)); - - coord.requeue_in_flight(); - - assert_eq!(coord.retry_counts.get(&7), Some(&1)); - assert!(!coord.is_in_flight(&7)); - assert_eq!(coord.pending_count(), 1); - } - - #[test] - fn test_requeue_in_flight_no_op_when_empty() { - let mut coord: DownloadCoordinator = DownloadCoordinator::default(); - coord.enqueue([1, 2]); - - coord.requeue_in_flight(); - - assert_eq!(coord.pending_count(), 2); - assert_eq!(coord.active_count(), 0); - } - - #[test] - fn test_cancel_pending_removes_from_pending_only() { - let mut coord: DownloadCoordinator = DownloadCoordinator::default(); - coord.enqueue([1, 2, 3]); - coord.mark_sent(&[10]); - coord.enqueue_retry(2); - assert_eq!(coord.retry_counts.get(&2), Some(&1)); - - coord.cancel_pending(&2); - - assert_eq!(coord.pending_count(), 2); - assert!(coord.is_in_flight(&10)); - assert_eq!(coord.retry_counts.get(&2), None); - - let items = coord.take_pending(2); - assert_eq!(items, vec![1, 3]); - } - - #[test] - fn test_cancel_pending_unknown_key_is_noop() { - let mut coord: DownloadCoordinator = DownloadCoordinator::default(); - coord.enqueue([1, 2]); - coord.mark_sent(&[5]); - - coord.cancel_pending(&99); - - assert_eq!(coord.pending_count(), 2); - assert_eq!(coord.active_count(), 1); - } - - #[test] - fn test_clear() { - let mut coord: DownloadCoordinator = DownloadCoordinator::default(); - coord.enqueue([1, 2, 3]); - coord.mark_sent(&[4]); - coord.enqueue_retry(5); - - coord.clear(); - - assert!(coord.is_empty()); - assert_eq!(coord.pending_count(), 0); - assert_eq!(coord.active_count(), 0); - } - - #[test] - fn test_remaining() { - let mut coord: DownloadCoordinator = DownloadCoordinator::default(); - coord.enqueue([1, 2, 3]); - coord.mark_sent(&[4]); - coord.mark_sent(&[5]); - - assert_eq!(coord.remaining(), 5); - } - - #[test] - fn test_config_builders() { - let config = - DownloadConfig::default().with_max_concurrent(20).with_timeout(Duration::from_secs(60)); - - assert_eq!(config.max_concurrent, 20); - assert_eq!(config.timeout, Duration::from_secs(60)); - } - - #[test] - fn test_with_string_keys() { - let mut coord: DownloadCoordinator = DownloadCoordinator::default(); - coord.enqueue(["block_a".to_string(), "block_b".to_string()]); - - let items = coord.take_pending(1); - coord.mark_sent(&items); - - assert!(coord.receive(&"block_a".to_string())); - assert!(!coord.receive(&"block_c".to_string())); - } -} diff --git a/dash-spv/src/sync/filter_headers/manager.rs b/dash-spv/src/sync/filter_headers/manager.rs index eb4f3600f..7c1185594 100644 --- a/dash-spv/src/sync/filter_headers/manager.rs +++ b/dash-spv/src/sync/filter_headers/manager.rs @@ -10,7 +10,7 @@ use tokio::sync::RwLock; use super::pipeline::FilterHeadersPipeline; use crate::error::SyncResult; -use crate::network::RequestSender; +use crate::network::NetworkManager; use crate::storage::{BlockHeaderStorage, FilterHeaderStorage}; use crate::sync::filter_headers::util::compute_filter_headers; use crate::sync::progress::ProgressPercentage; @@ -133,7 +133,10 @@ impl FilterHeadersManager } /// Start or resume filter header download. - async fn start_download(&mut self, requests: &RequestSender) -> SyncResult> { + async fn start_download( + &mut self, + network: &Arc, + ) -> SyncResult> { // Get current filter tip let filter_headers_tip = self.filter_header_storage.read().await.get_filter_tip_height().await?.unwrap_or(0); @@ -178,8 +181,8 @@ impl FilterHeadersManager .await?; drop(header_storage); - // Send initial requests - self.pipeline.send_pending(requests)?; + // Declare initial batches to the broker + self.pipeline.send_pending(network).await?; self.set_state(SyncState::Syncing); @@ -193,7 +196,7 @@ impl FilterHeadersManager pub(super) async fn handle_new_headers( &mut self, tip_height: u32, - requests: &RequestSender, + network: &Arc, ) -> SyncResult> { self.progress.update_block_header_tip_height(tip_height); self.update_target_height(tip_height); @@ -227,12 +230,12 @@ impl FilterHeadersManager .await?; } drop(header_storage); - self.pipeline.send_pending(requests)?; + self.pipeline.send_pending(network).await?; Ok(vec![]) } SyncState::WaitingForConnections | SyncState::WaitForEvents => { // Need full startup (calculates start from storage, handles checkpoints) - self.start_download(requests).await + self.start_download(network).await } _ => Ok(vec![]), } @@ -267,18 +270,12 @@ mod tests { .expect("Failed to create FilterHeadersManager") } - fn create_test_request_sender( - ) -> (RequestSender, tokio::sync::mpsc::UnboundedReceiver) { - let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); - (RequestSender::new(tx), rx) - } - #[tokio::test] async fn test_filter_headers_manager_new() { let manager = create_test_manager().await; assert_eq!(manager.identifier(), ManagerIdentifier::FilterHeader); assert_eq!(manager.state(), SyncState::WaitForEvents); - assert_eq!(manager.wanted_message_types(), vec![MessageType::CFHeaders]); + assert_eq!(manager.wanted_message_types(), [MessageType::CfHeaders]); assert!(!manager.block_headers_synced); } @@ -331,8 +328,11 @@ mod tests { #[tokio::test] async fn test_block_headers_synced_event_gating() { + use crate::network::NetworkManager; + use crate::test_utils::MockNetworkManager; + let mut manager = create_test_manager().await; - let (sender, _rx) = create_test_request_sender(); + let network: Arc = Arc::new(MockNetworkManager::new()); // Filter headers caught up to block header tip and target manager.progress.update_current_height(1000); @@ -344,7 +344,7 @@ mod tests { let event = SyncEvent::BlockHeadersStored { tip_height: 1000, }; - let events = manager.handle_sync_event(&event, &sender).await.unwrap(); + let events = manager.handle_sync_event(&event, &network).await.unwrap(); assert!(!manager.block_headers_synced); assert!(!events.iter().any(|e| matches!(e, SyncEvent::FilterHeadersSyncComplete { .. }))); @@ -352,7 +352,7 @@ mod tests { let event = SyncEvent::BlockHeaderSyncComplete { tip_height: 1000, }; - let events = manager.handle_sync_event(&event, &sender).await.unwrap(); + let events = manager.handle_sync_event(&event, &network).await.unwrap(); assert!(manager.block_headers_synced); assert!(events.iter().any(|e| matches!(e, SyncEvent::FilterHeadersSyncComplete { .. }))); assert_eq!(manager.state(), SyncState::Synced); @@ -360,8 +360,11 @@ mod tests { #[tokio::test] async fn test_block_header_sync_complete_during_active_download() { + use crate::network::NetworkManager; + use crate::test_utils::MockNetworkManager; + let mut manager = create_test_manager().await; - let (sender, _rx) = create_test_request_sender(); + let network: Arc = Arc::new(MockNetworkManager::new()); // Filter headers caught up to block tip, but target is higher (more headers coming) manager.progress.update_current_height(1000); @@ -373,7 +376,7 @@ mod tests { let event = SyncEvent::BlockHeaderSyncComplete { tip_height: 1000, }; - let events = manager.handle_sync_event(&event, &sender).await.unwrap(); + let events = manager.handle_sync_event(&event, &network).await.unwrap(); assert!(manager.block_headers_synced); assert!(!events.iter().any(|e| matches!(e, SyncEvent::FilterHeadersSyncComplete { .. }))); diff --git a/dash-spv/src/sync/filter_headers/pipeline.rs b/dash-spv/src/sync/filter_headers/pipeline.rs index 309b28ca0..c7d3ac50b 100644 --- a/dash-spv/src/sync/filter_headers/pipeline.rs +++ b/dash-spv/src/sync/filter_headers/pipeline.rs @@ -1,38 +1,34 @@ //! CFHeaders pipeline implementation. //! -//! Handles pipelined download of compact block filter headers (BIP 157/158). -//! Uses DownloadCoordinator for batch tracking with out-of-order buffering. +//! Declares wanted compact block filter header batches (BIP 157/158) to the +//! network manager (the broker) and buffers out-of-order responses for +//! sequential processing. The broker owns pacing, timeouts and retries — this +//! pipeline keeps no in-flight queue of its own. + +use std::collections::HashMap; +use std::sync::Arc; use dashcore::network::message::NetworkMessage; -use dashcore::network::message_filter::CFHeaders; +use dashcore::network::message_filter::{CFHeaders, GetCFHeaders}; use dashcore::BlockHash; -use std::collections::HashMap; -use std::time::Duration; use crate::error::{SyncError, SyncResult}; -use crate::network::RequestSender; +use crate::network::NetworkManager; use crate::storage::BlockHeaderStorage; -use crate::sync::download_coordinator::{DownloadConfig, DownloadCoordinator}; /// Batch size for filter header requests. const FILTER_HEADERS_BATCH_SIZE: u32 = 2000; -/// Maximum concurrent CFHeaders requests. -const MAX_CONCURRENT_CFHEADERS_REQUESTS: usize = 10; - -/// Timeout for CFHeaders requests (shorter for faster retry on multi-peer). -/// Timeout for CFHeaders requests. Single response but allow time for network latency. -const FILTER_HEADERS_TIMEOUT: Duration = Duration::from_secs(20); - /// Pipeline for downloading compact block filter headers. /// -/// Uses DownloadCoordinator for batch-level tracking (keyed by stop_hash), -/// with a HashMap buffer for out-of-order responses that need sequential processing. +/// Holds no request queue of its own: the batches it wants are exactly the +/// entries of `batch_starts` (keyed by stop_hash). It declares those to the +/// network manager (which de-duplicates, paces, times out and retries) and +/// buffers out-of-order responses for sequential processing. #[derive(Debug)] pub(super) struct FilterHeadersPipeline { - /// Core coordinator tracks batches by stop_hash. - coordinator: DownloadCoordinator, - /// Maps stop_hash -> start_height for each batch. + /// Wanted batches: stop_hash -> start_height. A batch leaves this map once + /// received. Doubles as the "is this batch wanted?" set for arrivals. batch_starts: HashMap, /// Out-of-order response buffer (start_height -> data). buffered: HashMap, @@ -52,11 +48,6 @@ impl FilterHeadersPipeline { /// Create a new CFHeaders pipeline. pub(super) fn new() -> Self { Self { - coordinator: DownloadCoordinator::new( - DownloadConfig::default() - .with_max_concurrent(MAX_CONCURRENT_CFHEADERS_REQUESTS) - .with_timeout(FILTER_HEADERS_TIMEOUT), - ), batch_starts: HashMap::new(), buffered: HashMap::new(), next_expected: 0, @@ -92,7 +83,6 @@ impl FilterHeadersPipeline { SyncError::Storage(format!("Missing header at height {}", batch_end)) })?; - self.coordinator.enqueue([stop_hash]); self.batch_starts.insert(stop_hash, current); added += 1; @@ -118,7 +108,7 @@ impl FilterHeadersPipeline { /// Check if the pipeline is complete. pub(super) fn is_complete(&self) -> bool { - self.coordinator.is_empty() + self.batch_starts.is_empty() && self.buffered.is_empty() && (self.target_height == 0 || self.next_expected > self.target_height) } @@ -130,7 +120,6 @@ impl FilterHeadersPipeline { start_height: u32, target_height: u32, ) -> SyncResult<()> { - self.coordinator.clear(); self.batch_starts.clear(); self.buffered.clear(); self.next_expected = start_height; @@ -147,7 +136,6 @@ impl FilterHeadersPipeline { SyncError::Storage(format!("Missing header at height {}", batch_end)) })?; - self.coordinator.enqueue([stop_hash]); self.batch_starts.insert(stop_hash, current); current = batch_end + 1; @@ -155,7 +143,7 @@ impl FilterHeadersPipeline { tracing::info!( "Built CFHeaders request queue: {} batches for heights {} to {}", - self.coordinator.pending_count(), + self.batch_starts.len(), start_height, target_height ); @@ -163,40 +151,38 @@ impl FilterHeadersPipeline { Ok(()) } - /// Send pending requests using a RequestSender (synchronous). - pub(super) fn send_pending(&mut self, requests: &RequestSender) -> SyncResult { - let count = self.coordinator.available_to_send(); - if count == 0 { + /// Declare every wanted CFHeaders batch to the network manager. + /// + /// Fired freely (on init, on extend, on arrival, on tick): the broker + /// de-duplicates, so re-declaring a batch already queued or on the wire is a + /// no-op, and it owns pacing and retry. Re-declaring each tick is the safety + /// net if a peer drops before the broker retries. + /// + /// Returns the number of batches declared (offered, not necessarily newly sent). + pub(super) async fn send_pending( + &mut self, + network: &Arc, + ) -> SyncResult { + if self.batch_starts.is_empty() { return Ok(0); } - let stop_hashes = self.coordinator.take_pending(count); - let mut sent = 0; - - for stop_hash in stop_hashes { - let Some(&start_height) = self.batch_starts.get(&stop_hash) else { - return Err(SyncError::InvalidState(format!( - "No batch_starts entry for pending stop_hash {}", - stop_hash - ))); - }; - - requests.request_filter_headers(start_height, stop_hash)?; - - self.coordinator.mark_sent(&[stop_hash]); - - tracing::debug!( - "Sent GetCFHeaders: start={}, stop={} ({} active, {} pending)", - start_height, - stop_hash, - self.coordinator.active_count(), - self.coordinator.pending_count() - ); - - sent += 1; + let batches: Vec<(BlockHash, u32)> = + self.batch_starts.iter().map(|(stop_hash, start)| (*stop_hash, *start)).collect(); + + for (stop_hash, start_height) in &batches { + network + .send(NetworkMessage::GetCFHeaders(GetCFHeaders { + filter_type: 0u8, + start_height: *start_height, + stop_hash: *stop_hash, + })) + .await; } - Ok(sent) + tracing::debug!("Declared {} wanted CFHeaders batch(es) to the broker", batches.len()); + + Ok(batches.len()) } /// Try to match an incoming message to a pipeline response. @@ -211,11 +197,8 @@ impl FilterHeadersPipeline { return None; } - // Match by stop_hash - the response includes it - if !self.coordinator.is_in_flight(&cfheaders.stop_hash) { - return None; - } - + // Match by stop_hash - the response includes it. A batch is "wanted" + // exactly while it sits in `batch_starts`. let start_height = *self.batch_starts.get(&cfheaders.stop_hash)?; Some((start_height, cfheaders.clone())) } @@ -225,7 +208,8 @@ impl FilterHeadersPipeline { /// Returns `Some(data)` if this response is the next expected and should /// be processed immediately. Returns `None` if buffered for later. pub(super) fn receive(&mut self, start_height: u32, data: CFHeaders) -> Option { - self.coordinator.receive(&data.stop_hash); + // Drop the batch from the wanted set; the broker is told the request was + // answered by the manager via `request_answered(RequestKey::CfHeaders)`. self.batch_starts.remove(&data.stop_hash); if start_height == self.next_expected { @@ -253,13 +237,6 @@ impl FilterHeadersPipeline { } ready } - - /// Re-enqueue timed out requests for retry. - pub(super) fn handle_timeouts(&mut self) { - for stop_hash in self.coordinator.check_timeouts() { - self.coordinator.enqueue_retry(stop_hash); - } - } } #[cfg(test)] @@ -307,8 +284,7 @@ mod tests { let stop_hash = BlockHash::all_zeros(); - // Mark batch as in-flight (by stop_hash) - pipeline.coordinator.mark_sent(&[stop_hash]); + // Mark batch as wanted (by stop_hash) pipeline.batch_starts.insert(stop_hash, 1); let cfheaders = CFHeaders { @@ -333,8 +309,7 @@ mod tests { let stop_hash = BlockHash::all_zeros(); - // Mark batch as in-flight (by stop_hash) - pipeline.coordinator.mark_sent(&[stop_hash]); + // Mark batch as wanted (by stop_hash) pipeline.batch_starts.insert(stop_hash, 2000); let cfheaders = CFHeaders { @@ -373,76 +348,4 @@ mod tests { assert_eq!(ready[0].0, 2000); assert_eq!(pipeline.buffered.len(), 0); } - - #[test] - fn test_handle_timeouts_basic_retry() { - use std::time::Duration; - - let mut pipeline = FilterHeadersPipeline { - coordinator: DownloadCoordinator::new( - DownloadConfig::default().with_timeout(Duration::from_millis(1)), - ), - batch_starts: HashMap::new(), - buffered: HashMap::new(), - next_expected: 1, - target_height: 2000, - }; - - let stop_hash = BlockHash::all_zeros(); - pipeline.coordinator.mark_sent(&[stop_hash]); - pipeline.batch_starts.insert(stop_hash, 1); - - std::thread::sleep(Duration::from_millis(5)); - - pipeline.handle_timeouts(); - assert_eq!(pipeline.coordinator.pending_count(), 1); - } - - #[test] - fn test_send_pending_errors_on_missing_batch_starts() { - let mut pipeline = FilterHeadersPipeline::new(); - pipeline.next_expected = 1; - pipeline.target_height = 2000; - - let hash_without_entry = BlockHash::from_byte_array([0x02; 32]); - - // Enqueue a stop_hash without a corresponding batch_starts entry - pipeline.coordinator.enqueue([hash_without_entry]); - - let (tx, _rx) = tokio::sync::mpsc::unbounded_channel(); - let requests = RequestSender::new(tx); - - let err = pipeline.send_pending(&requests).unwrap_err(); - assert!(matches!(err, SyncError::InvalidState(_))); - } - - #[test] - fn test_handle_timeouts_multiple_batches() { - use std::time::Duration; - - let mut pipeline = FilterHeadersPipeline { - coordinator: DownloadCoordinator::new( - DownloadConfig::default().with_timeout(Duration::from_millis(1)), - ), - batch_starts: HashMap::new(), - buffered: HashMap::new(), - next_expected: 1, - target_height: 4000, - }; - - let hash1 = BlockHash::from_byte_array([0x01; 32]); - let hash2 = BlockHash::from_byte_array([0x02; 32]); - - pipeline.coordinator.mark_sent(&[hash1, hash2]); - pipeline.batch_starts.insert(hash1, 1); - pipeline.batch_starts.insert(hash2, 2001); - - std::thread::sleep(Duration::from_millis(5)); - - pipeline.handle_timeouts(); - // Both batches re-queued - assert_eq!(pipeline.coordinator.pending_count(), 2); - assert!(pipeline.batch_starts.contains_key(&hash1)); - assert!(pipeline.batch_starts.contains_key(&hash2)); - } } diff --git a/dash-spv/src/sync/filter_headers/sync_manager.rs b/dash-spv/src/sync/filter_headers/sync_manager.rs index eae554d57..e5e423d59 100644 --- a/dash-spv/src/sync/filter_headers/sync_manager.rs +++ b/dash-spv/src/sync/filter_headers/sync_manager.rs @@ -1,5 +1,5 @@ use crate::error::SyncResult; -use crate::network::{Message, MessageType, RequestSender}; +use crate::network::{MessageType, NetworkManager, RequestKey}; use crate::storage::{BlockHeaderStorage, FilterHeaderStorage}; use crate::sync::filter_headers::pipeline::FilterHeadersPipeline; use crate::sync::progress::ProgressPercentage; @@ -8,6 +8,9 @@ use crate::sync::{ }; use crate::SyncError; use async_trait::async_trait; +use dashcore::network::message::NetworkMessage; +use std::net::SocketAddr; +use std::sync::Arc; #[async_trait] impl SyncManager for FilterHeadersManager { @@ -28,7 +31,7 @@ impl SyncManager for FilterHeade } fn wanted_message_types(&self) -> &'static [MessageType] { - &[MessageType::CFHeaders] + &[MessageType::CfHeaders] } fn on_disconnect(&mut self) { @@ -39,11 +42,12 @@ impl SyncManager for FilterHeade async fn handle_message( &mut self, - msg: Message, - requests: &RequestSender, + _peer: SocketAddr, + msg: NetworkMessage, + network: &Arc, ) -> SyncResult> { // Match response to get start height - let Some((start_height, cfheaders)) = self.pipeline.match_response(msg.inner()) else { + let Some((start_height, cfheaders)) = self.pipeline.match_response(&msg) else { if self.pipeline.is_complete() { if let Some(event) = self.try_complete_sync() { return Ok(vec![event]); @@ -52,6 +56,10 @@ impl SyncManager for FilterHeade return Ok(vec![]); }; + // Response correlated: tell the network manager to stop tracking this + // batch for timeout/retry. + network.request_answered(RequestKey::CfHeaders(cfheaders.stop_hash)).await; + let mut events = Vec::new(); // Try to receive (may buffer if out of order) @@ -114,8 +122,8 @@ impl SyncManager for FilterHeade ); } - // Send more requests - self.pipeline.send_pending(requests)?; + // Declare any remaining wanted batches to the broker + self.pipeline.send_pending(network).await?; if self.pipeline.is_complete() { if let Some(event) = self.try_complete_sync() { @@ -129,32 +137,22 @@ impl SyncManager for FilterHeade async fn handle_sync_event( &mut self, event: &SyncEvent, - requests: &RequestSender, + network: &Arc, ) -> SyncResult> { match event { SyncEvent::BlockHeaderSyncComplete { tip_height, } => { self.block_headers_synced = true; - self.handle_new_headers(*tip_height, requests).await + self.handle_new_headers(*tip_height, network).await } SyncEvent::BlockHeadersStored { tip_height, - } => self.handle_new_headers(*tip_height, requests).await, + } => self.handle_new_headers(*tip_height, network).await, _ => Ok(vec![]), } } - async fn tick(&mut self, requests: &RequestSender) -> SyncResult> { - // Handle timed out requests (re-queues them for retry) - self.pipeline.handle_timeouts(); - - // Send pending requests (including retries) - self.pipeline.send_pending(requests)?; - - Ok(vec![]) - } - fn progress(&self) -> SyncManagerProgress { SyncManagerProgress::FilterHeaders(self.progress.clone()) } diff --git a/dash-spv/src/sync/filters/manager.rs b/dash-spv/src/sync/filters/manager.rs index 376a80c8c..dec958872 100644 --- a/dash-spv/src/sync/filters/manager.rs +++ b/dash-spv/src/sync/filters/manager.rs @@ -14,7 +14,7 @@ use super::batch::FiltersBatch; use super::block_match_tracker::{BlockMatchTracker, BlockTrackResult}; use super::pipeline::FiltersPipeline; use crate::error::SyncResult; -use crate::network::RequestSender; +use crate::network::NetworkManager; use crate::storage::{BlockHeaderStorage, FilterHeaderStorage, FilterStorage}; use crate::sync::filters::util::get_prev_filter_header; use crate::sync::{FiltersProgress, SyncEvent, SyncManager, SyncState}; @@ -29,19 +29,6 @@ use tokio::sync::RwLock; /// Batch size for processing filters. const BATCH_PROCESSING_SIZE: u32 = 5000; -/// Snapshot of a behind wallet's compact-filter query inputs for a batch scan. -struct WalletScanState { - /// The wallet these inputs belong to. - id: WalletId, - /// The wallet's committed sync checkpoint; heights at or below it are skipped. - synced: u32, - /// Monitored scriptPubKeys. - scripts: Vec, - /// Bare `hash160` filter elements (owner/voting key hashes) a compact - /// filter carries beyond the scriptPubKeys. - elements: Vec>, -} - /// Maximum number of batches to scan ahead while waiting for blocks. const MAX_LOOKAHEAD_BATCHES: usize = 3; @@ -177,7 +164,7 @@ impl, ) -> SyncResult> { debug_assert!(self.is_idle(), "manager should have no in-flight state on start"); @@ -294,7 +281,7 @@ impl = HashMap::new(); - let mut filter_elements: HashMap>> = HashMap::new(); - { + // own progress are skipped during the rescan. + let synced_heights: HashMap = { let wallet = self.wallet.read().await; - for id in new_scripts.keys() { - synced_heights.insert(*id, wallet.wallet_synced_height(id)); - filter_elements.insert(*id, wallet.monitored_filter_elements_for(id)); - } - } + new_scripts.keys().map(|id| (*id, wallet.wallet_synced_height(id))).collect() + }; let mut block_to_wallets: BTreeMap> = BTreeMap::new(); for (wallet_id, scripts) in new_scripts { - let elements = filter_elements.get(wallet_id).map(Vec::as_slice).unwrap_or(&[]); - if scripts.is_empty() && elements.is_empty() { + if scripts.is_empty() { continue; } let scripts_vec: Vec = scripts.iter().cloned().collect(); let min_synced = synced_heights.get(wallet_id).copied().unwrap_or(0); - let matches = check_compact_filters_for_elements( - batch_filters, - &scripts_vec, - elements, - min_synced, - ); + let matches = + check_compact_filters_for_elements(batch_filters, &scripts_vec, &[], min_synced); for key in matches { block_to_wallets.entry(key).or_default().insert(*wallet_id); } @@ -805,20 +780,12 @@ impl = Vec::new(); + let mut wallet_states: Vec<(WalletId, u32, Vec)> = Vec::new(); for wallet_id in &behind { let synced = wallet.wallet_synced_height(wallet_id); let scripts = wallet.monitored_script_pubkeys_for(wallet_id); - // Bare owner/voting key hashes a compact filter carries beyond the - // wallet's scriptPubKeys. - let elements = wallet.monitored_filter_elements_for(wallet_id); - if !scripts.is_empty() || !elements.is_empty() { - wallet_states.push(WalletScanState { - id: *wallet_id, - synced, - scripts, - elements, - }); + if !scripts.is_empty() { + wallet_states.push((*wallet_id, synced, scripts)); } } drop(wallet); @@ -848,25 +815,17 @@ impl = - wallet_states.iter().flat_map(|s| s.scripts.iter().cloned()).collect(); - let union_elements: Vec> = - wallet_states.iter().flat_map(|s| s.elements.iter().cloned()).collect(); - let min_synced = wallet_states.iter().map(|s| s.synced).min().unwrap_or(0); + wallet_states.iter().flat_map(|(_, _, scripts)| scripts.iter().cloned()).collect(); + let min_synced = wallet_states.iter().map(|(_, synced, _)| *synced).min().unwrap_or(0); - // Pre-group each wallet's scripts and bare elements by length once; - // reused across every matched filter. + // Pre-group each wallet's scripts by length once; reused across every matched filter. let wallet_queries: Vec<(WalletId, u32, FilterQuery)> = wallet_states .iter() - .map(|s| { - let mut query: FilterQuery = s.scripts.iter().map(|sp| sp.as_bytes()).collect(); - for element in &s.elements { - query.push(element); - } - (s.id, s.synced, query) + .map(|(id, synced, scripts)| { + (*id, *synced, scripts.iter().map(|s| s.as_bytes()).collect()) }) .collect(); @@ -876,12 +835,8 @@ impl> = BTreeMap::new(); for key in matches { @@ -973,7 +928,7 @@ impl, ) -> SyncResult> { self.progress.update_filter_header_tip_height(tip_height); self.update_target_height(tip_height); @@ -991,7 +946,7 @@ impl {} } @@ -1029,7 +984,7 @@ impl Arc { + Arc::new(crate::test_utils::MockNetworkManager::new()) + } + + /// Like [`test_network`] but keeps the concrete mock so a test can assert on + /// what the manager declared to the broker. + async fn test_network_with_mock( + ) -> (Arc, Arc) { + let mock = Arc::new(crate::test_utils::MockNetworkManager::new()); + let network: Arc = mock.clone(); + (mock, network) + } + + /// Start heights of every `GetCFilters` the manager declared, in order. + fn declared_filter_starts(mock: &crate::test_utils::MockNetworkManager) -> Vec { + mock.sent_messages() + .iter() + .filter_map(|m| match m { + NetworkMessage::GetCFilters(gcf) => Some(gcf.start_height), + _ => None, + }) + .collect() + } type TestFiltersManager = FiltersManager< PersistentBlockHeaderStorage, @@ -1151,7 +1131,7 @@ mod tests { let manager = create_test_manager().await; assert_eq!(manager.identifier(), ManagerIdentifier::Filter); assert_eq!(manager.state(), SyncState::WaitForEvents); - assert_eq!(manager.wanted_message_types(), vec![MessageType::CFilter]); + assert_eq!(manager.wanted_message_types(), [MessageType::CFilter]); assert_eq!(manager.progress.committed_height(), 0); assert_eq!(manager.progress.stored_height(), 0); assert_eq!(manager.progress.target_height(), 0); @@ -1905,8 +1885,8 @@ mod tests { .await .unwrap(); - let (tx, _rx) = unbounded_channel(); - let _ = manager.tick(&RequestSender::new(tx)).await.unwrap(); + let network = test_network().await; + let _ = manager.tick(&network).await.unwrap(); // Batch must start at 151, not at 0. assert!(manager.active_batches.contains_key(&151)); @@ -2153,8 +2133,8 @@ mod tests { // Chain tip higher so the Synced early-return is not taken manager.progress.update_target_height(1000); - let (tx, _rx) = unbounded_channel(); - let events = manager.start_download(&RequestSender::new(tx)).await.unwrap(); + let network = test_network().await; + let events = manager.start_download(&network).await.unwrap(); assert!(events.is_empty()); assert_eq!(manager.state(), SyncState::WaitForEvents); @@ -2182,8 +2162,8 @@ mod tests { manager.progress.update_filter_header_tip_height(100); manager.progress.update_target_height(1000); - let (tx, _rx) = unbounded_channel(); - let events = manager.start_download(&RequestSender::new(tx)).await.unwrap(); + let network = test_network().await; + let events = manager.start_download(&network).await.unwrap(); assert_eq!(manager.state(), SyncState::Syncing); assert!(!manager.is_idle()); @@ -2194,184 +2174,6 @@ mod tests { assert_eq!(batch.end_height(), 100); } - /// Reproduces #892: filters were only ever stored near a high tip (e.g. - /// headers previously synced from a checkpoint), and the wallet's scan - /// start is far below that region. `start_download` must not treat the - /// stored tip watermark as contiguous coverage from `scan_start`: - /// preloading `load_filters(scan_start, ..)` would read never-populated - /// segments (debug abort in `SegmentCache::get_items`, sentinel filter - /// data in release builds). The unreachable tip-region filters are - /// discarded and the download restarts from `scan_start`. - #[tokio::test] - async fn test_start_download_discards_stored_filters_above_scan_start() { - let mut manager = create_test_manager().await; - - // Block headers cover the whole range so send_pending can resolve stop hashes. - let headers = dashcore::block::Header::dummy_batch(0..1001); - manager - .header_storage - .write() - .await - .store_headers( - &headers.iter().map(crate::types::HashedBlockHeader::from).collect::>(), - ) - .await - .unwrap(); - - // Filters were only ever stored near the tip: 900..=1000. The heights - // below 900 were never populated. - { - let mut filter_storage = manager.filter_storage.write().await; - for height in 900..=1000u32 { - filter_storage.store_filter(height, &[height as u8; 8]).await.unwrap(); - } - } - // Restart shape: `new()` seeds stored_height from the tip watermark. - manager.progress.update_stored_height(1000); - manager.progress.update_filter_header_tip_height(1000); - manager.progress.update_target_height(1000); - - // Wallet committed far below the stored region, so scan_start = 100. - manager.wallet.write().await.update_wallet_synced_height(&MOCK_WALLET_ID, 99); - - let (tx, mut rx) = unbounded_channel(); - let events = manager.start_download(&RequestSender::new(tx)).await.unwrap(); - - assert!(events.is_empty()); - assert_eq!(manager.state(), SyncState::Syncing); - - // The unreachable tip-region filters were discarded entirely... - assert_eq!(manager.filter_storage.read().await.filter_tip_height().await.unwrap(), 0); - assert_eq!(manager.filter_storage.read().await.filter_start_height().await, None); - - // ...nothing was preloaded into the initial batch... - let batch = manager.active_batches.get(&100).expect("initial batch at scan_start"); - assert!(batch.filters().is_empty()); - assert!(!batch.verified()); - assert!(!batch.scanned()); - assert_eq!(batch.end_height(), 1000); - - // ...scan gating no longer sees the stale tip watermark... - assert_eq!(manager.progress.stored_height(), 0); - - // ...and both the store cursor and the download restart from - // scan_start rather than stored_filters_tip + 1. - assert_eq!(manager.next_batch_to_store, 100); - match rx.try_recv().expect("a filter request must have been sent") { - NetworkRequest::SendMessage(NetworkMessage::GetCFilters(gcf)) => { - assert_eq!(gcf.start_height, 100); - } - other => panic!("Expected GetCFilters, got {:?}", other), - } - } - - /// Counterpart to the sparse-storage case: when the stored filter range - /// actually reaches down to `scan_start`, the preload happens exactly as - /// before and nothing is discarded. - #[tokio::test] - async fn test_start_download_preloads_when_stored_filters_cover_scan_start() { - let mut manager = create_test_manager().await; - - let headers = dashcore::block::Header::dummy_batch(0..1001); - manager - .header_storage - .write() - .await - .store_headers( - &headers.iter().map(crate::types::HashedBlockHeader::from).collect::>(), - ) - .await - .unwrap(); - - // Filters stored densely from genesis through the tip. - { - let mut filter_storage = manager.filter_storage.write().await; - for height in 0..=1000u32 { - filter_storage.store_filter(height, &[height as u8; 8]).await.unwrap(); - } - } - manager.progress.update_stored_height(1000); - manager.progress.update_filter_header_tip_height(1000); - manager.progress.update_target_height(1000); - - manager.wallet.write().await.update_wallet_synced_height(&MOCK_WALLET_ID, 99); - - let (tx, mut rx) = unbounded_channel(); - manager.start_download(&RequestSender::new(tx)).await.unwrap(); - - assert_eq!(manager.state(), SyncState::Syncing); - - // Storage is untouched. - assert_eq!(manager.filter_storage.read().await.filter_tip_height().await.unwrap(), 1000); - assert_eq!(manager.filter_storage.read().await.filter_start_height().await, Some(0)); - - // The stored range 100..=1000 was preloaded and the batch is verified - // and scanned immediately. - let batch = manager.active_batches.get(&100).expect("initial batch at scan_start"); - assert_eq!(batch.filters().len(), 901); - assert!(batch.verified()); - assert!(batch.scanned()); - assert_eq!(manager.progress.stored_height(), 1000); - - // Nothing left to download: everything through the filter header tip - // is already stored. - assert_eq!(manager.next_batch_to_store, 1001); - assert!(rx.try_recv().is_err(), "no filter request expected"); - } - - /// Genesis-only storage: a single filter at height 0, scanning from 0. - /// `filter_tip_height` collapses to 0 for both an empty store and this - /// one, so gating the preload on `tip > 0` would misread it as empty and - /// needlessly re-download height 0. The stored start (Some(0)) must drive - /// the decision: the filter is preloaded and nothing is discarded or - /// re-requested. Regtest can produce exactly this shape. - #[tokio::test] - async fn test_start_download_preloads_genesis_only_stored_filter() { - let mut manager = create_test_manager().await; - - let headers = dashcore::block::Header::dummy_batch(0..1); - manager - .header_storage - .write() - .await - .store_headers( - &headers.iter().map(crate::types::HashedBlockHeader::from).collect::>(), - ) - .await - .unwrap(); - - // Exactly one filter stored, at height 0. - manager.filter_storage.write().await.store_filter(0, &[0u8; 8]).await.unwrap(); - assert_eq!(manager.filter_storage.read().await.filter_tip_height().await.unwrap(), 0); - assert_eq!(manager.filter_storage.read().await.filter_start_height().await, Some(0)); - - // Restart shape at the genesis tip. - manager.progress.update_stored_height(0); - manager.progress.update_filter_header_tip_height(0); - manager.progress.update_target_height(0); - // Wallet at genesis: scan_start = 0. - - let (tx, mut rx) = unbounded_channel(); - manager.start_download(&RequestSender::new(tx)).await.unwrap(); - - // The lone filter was NOT discarded... - assert_eq!(manager.filter_storage.read().await.filter_tip_height().await.unwrap(), 0); - assert_eq!(manager.filter_storage.read().await.filter_start_height().await, Some(0)); - - // ...it was preloaded into the initial batch, which is verified and - // scanned since the whole (single-height) range is covered... - let batch = manager.active_batches.get(&0).expect("initial batch at scan_start"); - assert_eq!(batch.filters().len(), 1); - assert!(batch.verified()); - assert!(batch.scanned()); - assert_eq!(manager.progress.stored_height(), 0); - - // ...and the download frontier sits above the stored tip, so no - // filter request goes out for the already-stored genesis height. - assert_eq!(manager.next_batch_to_store, 1); - assert!(rx.try_recv().is_err(), "no filter request expected for the genesis-only store"); - } - #[tokio::test] async fn test_handle_new_filter_headers_transitions_synced_to_syncing() { let mut manager = create_test_manager().await; @@ -2390,11 +2192,10 @@ mod tests { // stops it before creating any batch or emitting an event. manager.active_batches.insert(101, FiltersBatch::new(101, 200, HashMap::new())); - let (tx, _rx) = unbounded_channel(); - let requests = RequestSender::new(tx); + let network = test_network().await; // New filter headers arrive at 150: committed(100) < tip(150) - let events = manager.handle_new_filter_headers(150, &requests).await.unwrap(); + let events = manager.handle_new_filter_headers(150, &network).await.unwrap(); assert!(events.is_empty()); assert_eq!(manager.state(), SyncState::Syncing); @@ -2426,10 +2227,9 @@ mod tests { manager.progress.update_filter_header_tip_height(100); manager.progress.update_target_height(100); - let (tx, _rx) = unbounded_channel(); - let requests = RequestSender::new(tx); + let network = test_network().await; - let events = manager.handle_new_filter_headers(100, &requests).await.unwrap(); + let events = manager.handle_new_filter_headers(100, &network).await.unwrap(); assert_eq!(manager.state(), SyncState::Synced); assert!( @@ -2475,12 +2275,11 @@ mod tests { manager.progress.update_filter_header_tip_height(100); manager.progress.update_target_height(100); - let (tx, _rx) = unbounded_channel(); - let requests = RequestSender::new(tx); + let network = test_network().await; // Boot already synced: start_download detects the synced state and // returns early, but must still advance the scan frontier. - let events = manager.start_download(&requests).await.unwrap(); + let events = manager.start_download(&network).await.unwrap(); assert_eq!(manager.state(), SyncState::Synced); assert!(events.iter().any(|e| matches!( e, @@ -2492,7 +2291,7 @@ mod tests { // A new block extends the chain; the lookahead batch must start at the // frontier, not at height 0. - manager.handle_new_filter_headers(101, &requests).await.unwrap(); + manager.handle_new_filter_headers(101, &network).await.unwrap(); assert!(!manager.active_batches.contains_key(&0)); assert_eq!(manager.active_batches.keys().next(), Some(&101)); } @@ -2511,12 +2310,11 @@ mod tests { // Fully-synced restart: `start_sync` requires `WaitingForConnections`. manager.set_state(SyncState::WaitingForConnections); - let (tx, _rx) = unbounded_channel(); - let requests = RequestSender::new(tx); + let network = test_network().await; // Reconnect: the already-synced branch reports completion and anchors the // store/processing cursors at the frontier. - let events = manager.start_sync(&requests).await.unwrap(); + let events = manager.start_sync(&network).await.unwrap(); assert_eq!(manager.state(), SyncState::Synced); assert!(events.iter().any(|e| matches!( e, @@ -2534,7 +2332,7 @@ mod tests { end_height: 101, tip_height: 101, }, - &requests, + &network, ) .await .unwrap(); @@ -2546,13 +2344,10 @@ mod tests { block_hash: headers[101].block_hash(), filter: boundary_filter.content.clone(), }; - manager - .handle_message(Message::new(peer, NetworkMessage::CFilter(cfilter)), &requests) - .await - .unwrap(); + manager.handle_message(peer, NetworkMessage::CFilter(cfilter), &network).await.unwrap(); // A trailing tick drives any residual processing to completion. - manager.tick(&requests).await.unwrap(); + manager.tick(&network).await.unwrap(); assert_eq!(manager.progress.committed_height(), 101); assert_eq!(manager.state(), SyncState::Synced); @@ -2575,8 +2370,7 @@ mod tests { // Fully-synced boot leaves the manager in its default state. assert_eq!(manager.state(), SyncState::WaitForEvents); - let (tx, mut rx) = unbounded_channel(); - let requests = RequestSender::new(tx); + let network = test_network().await; // Filter header sync completing at the stored tip reports Synced via // `start_download`'s early return, anchoring the frontier cursors. @@ -2585,7 +2379,7 @@ mod tests { &SyncEvent::FilterHeadersSyncComplete { tip_height: 100, }, - &requests, + &network, ) .await .unwrap(); @@ -2606,22 +2400,11 @@ mod tests { end_height: 101, tip_height: 101, }, - &requests, + &network, ) .await .unwrap(); - // Only the boundary body may be requested: a pipeline left unparked - // would re-request every filter from height 1 here. - while let Ok(request) = rx.try_recv() { - let (NetworkRequest::SendMessage(msg) - | NetworkRequest::SendMessageToPeer(msg, _) - | NetworkRequest::BroadcastMessage(msg)) = request; - if let NetworkMessage::GetCFilters(get) = msg { - assert_eq!(get.start_height, 101, "unexpected filter re-download"); - } - } - // The peer answers with the boundary filter body over the real path. let peer: SocketAddr = "127.0.0.1:19999".parse().unwrap(); let cfilter = CFilter { @@ -2629,13 +2412,10 @@ mod tests { block_hash: headers[101].block_hash(), filter: boundary_filter.content.clone(), }; - manager - .handle_message(Message::new(peer, NetworkMessage::CFilter(cfilter)), &requests) - .await - .unwrap(); + manager.handle_message(peer, NetworkMessage::CFilter(cfilter), &network).await.unwrap(); // A trailing tick drives any residual processing to completion. - manager.tick(&requests).await.unwrap(); + manager.tick(&network).await.unwrap(); assert_eq!(manager.progress.committed_height(), 101); assert_eq!(manager.state(), SyncState::Synced); @@ -2656,28 +2436,15 @@ mod tests { manager.progress.update_filter_header_tip_height(101); manager.set_state(SyncState::WaitingForConnections); - let (tx, mut rx) = unbounded_channel(); - let requests = RequestSender::new(tx); + let network = test_network().await; // Reconnect delegates to the download path, which must request the // boundary body at 101 rather than parking without asking for it. - let events = manager.start_sync(&requests).await.unwrap(); + let events = manager.start_sync(&network).await.unwrap(); assert_eq!(manager.state(), SyncState::Syncing); assert!(!events.iter().any(|e| matches!(e, SyncEvent::SyncStart { .. }))); assert!(manager.active_batches.contains_key(&101)); - let mut requested_boundary = false; - while let Ok(request) = rx.try_recv() { - let (NetworkRequest::SendMessage(msg) - | NetworkRequest::SendMessageToPeer(msg, _) - | NetworkRequest::BroadcastMessage(msg)) = request; - if let NetworkMessage::GetCFilters(get) = msg { - assert_eq!(get.start_height, 101, "unexpected filter re-download"); - requested_boundary = true; - } - } - assert!(requested_boundary, "boundary filter body 101 must be requested"); - // The peer answers with the boundary filter body over the real path. let peer: SocketAddr = "127.0.0.1:19999".parse().unwrap(); let cfilter = CFilter { @@ -2685,12 +2452,9 @@ mod tests { block_hash: headers[101].block_hash(), filter: boundary_filter.content.clone(), }; - manager - .handle_message(Message::new(peer, NetworkMessage::CFilter(cfilter)), &requests) - .await - .unwrap(); + manager.handle_message(peer, NetworkMessage::CFilter(cfilter), &network).await.unwrap(); - manager.tick(&requests).await.unwrap(); + manager.tick(&network).await.unwrap(); assert_eq!(manager.progress.committed_height(), 101); assert_eq!(manager.state(), SyncState::Synced); @@ -2710,10 +2474,9 @@ mod tests { manager.progress.update_target_height(100); manager.filter_pipeline.init(101, 100); - let (tx, _rx) = unbounded_channel(); - let requests = RequestSender::new(tx); + let network = test_network().await; - let events = manager.handle_new_filter_headers(100, &requests).await.unwrap(); + let events = manager.handle_new_filter_headers(100, &network).await.unwrap(); assert_eq!(manager.state(), SyncState::Synced); assert!(events.is_empty()); @@ -2784,8 +2547,7 @@ mod tests { // MockWallet defaults to synced_height=0, so wallets_behind(100) = {MOCK_WALLET_ID}. assert_eq!(manager.wallet.read().await.synced_height(), 0); - let (tx, _rx) = unbounded_channel(); - let requests = RequestSender::new(tx); + let network = test_network().await; // Sanity: the pre-populated stale processed record is present, so // `track` for the same wallet would short-circuit to AlreadyProcessed. @@ -2798,7 +2560,7 @@ mod tests { manager.tracker.clear(); manager.tracker.record_processed(150, stale_hash, &BTreeSet::from([MOCK_WALLET_ID])); - let events = manager.tick(&requests).await.unwrap(); + let events = manager.tick(&network).await.unwrap(); // Old in-flight state was cleared and a fresh batch was created at scan_start=0. assert!(!manager.active_batches.contains_key(&101)); @@ -2846,10 +2608,9 @@ mod tests { manager.progress.update_filter_header_tip_height(200); manager.progress.update_target_height(200); - let (tx, _rx) = unbounded_channel(); - let requests = RequestSender::new(tx); + let network = test_network().await; - let events = manager.tick(&requests).await.unwrap(); + let events = manager.tick(&network).await.unwrap(); assert!(events.is_empty()); assert_eq!(manager.progress.committed_height(), 100); @@ -2857,6 +2618,57 @@ mod tests { assert!(manager.active_batches.is_empty()); } + /// A wallet reporting `synced_height = 0` is "behind" any positive + /// `committed_height`, but on a checkpoint sync the scan cannot reach below + /// the stored headers' start anyway, so a restart would resume above the + /// frontier and change nothing. The trigger must compare against that floor + /// rather than the raw `synced_height`: comparing the raw value is + /// level-sensitive, so it would fire on every tick and wipe the in-flight + /// filter batches before any could complete. + #[tokio::test] + async fn test_tick_does_not_rescan_when_restart_would_land_above_committed() { + let mut manager = create_test_manager().await; + + // Checkpoint sync: headers start at 500, so no scan can reach below it. + let headers = dashcore::block::Header::dummy_batch(500..600); + manager + .header_storage + .write() + .await + .store_headers_at_height( + &headers.iter().map(crate::types::HashedBlockHeader::from).collect::>(), + 500, + ) + .await + .unwrap(); + assert_eq!(manager.header_storage.read().await.get_start_height().await, Some(500)); + + // MockWallet defaults to synced_height=0, so wallets_behind(400) lists it + // even though it needs no coverage below the 500 floor. + assert_eq!(manager.wallet.read().await.synced_height(), 0); + assert!(!manager.wallet.read().await.wallets_behind(400).is_empty()); + + manager.set_state(SyncState::Syncing); + manager.progress.update_committed_height(400); + manager.progress.update_stored_height(400); + manager.progress.update_filter_header_tip_height(600); + manager.progress.update_target_height(600); + + // In-flight work that a spurious `reset_for_rescan` would wipe. + manager.active_batches.insert(401, FiltersBatch::new(401, 500, HashMap::new())); + manager.filter_pipeline.init(401, 500); + + let network = test_network().await; + manager.tick(&network).await.unwrap(); + + // restart_at = max(0 + 1, 500) = 500 > 400, so the scan was left alone. + assert_eq!(manager.progress.committed_height(), 400); + assert!( + manager.active_batches.contains_key(&401), + "in-flight batch must survive a restart that could not reach below the floor" + ); + } + /// `committed_height = 0` on a fresh manager must not falsely trip the /// rescan trigger. `wallets_behind(0)` returns an empty set since heights /// are unsigned, so no wallet can be strictly less than 0. @@ -2867,10 +2679,9 @@ mod tests { assert_eq!(manager.progress.committed_height(), 0); assert_eq!(manager.state(), SyncState::WaitForEvents); - let (tx, _rx) = unbounded_channel(); - let requests = RequestSender::new(tx); + let network = test_network().await; - let events = manager.tick(&requests).await.unwrap(); + let events = manager.tick(&network).await.unwrap(); assert!(events.is_empty()); assert!(manager.is_idle()); @@ -2888,10 +2699,9 @@ mod tests { // Wallet behind committed — would normally trip the trigger. assert!(!manager.wallet.read().await.wallets_behind(100).is_empty()); - let (tx, _rx) = unbounded_channel(); - let requests = RequestSender::new(tx); + let network = test_network().await; - let events = manager.tick(&requests).await.unwrap(); + let events = manager.tick(&network).await.unwrap(); assert!(events.is_empty()); // committed_height not lowered, no batches created. @@ -2978,12 +2788,11 @@ mod tests { manager.progress.update_target_height(1000); manager.filter_pipeline.init(101, 100); - let (tx, _rx) = unbounded_channel(); - let requests = RequestSender::new(tx); + let network = test_network().await; // committed(0) < tip(100) fires the guard even though stored == tip. // Because stored >= tip, send_pending is skipped (no downloads needed). - let _events = manager.handle_new_filter_headers(100, &requests).await.unwrap(); + let _events = manager.handle_new_filter_headers(100, &network).await.unwrap(); assert_eq!(manager.state(), SyncState::Syncing); assert!( @@ -3030,11 +2839,10 @@ mod tests { manager.progress.update_target_height(101); manager.active_batches.insert(0, FiltersBatch::new(0, 100, HashMap::new())); - let (tx, _rx) = unbounded_channel(); - let requests = RequestSender::new(tx); + let network = test_network().await; // New block 101 arrives: tip grows past the in-flight rescan boundary. - manager.handle_new_filter_headers(101, &requests).await.unwrap(); + manager.handle_new_filter_headers(101, &network).await.unwrap(); assert!( manager.active_batches.contains_key(&101), @@ -3056,10 +2864,9 @@ mod tests { manager.progress.update_target_height(1000); manager.filter_pipeline.init(101, 100); - let (tx, _rx) = unbounded_channel(); - let requests = RequestSender::new(tx); + let network = test_network().await; - let events = manager.handle_new_filter_headers(100, &requests).await.unwrap(); + let events = manager.handle_new_filter_headers(100, &network).await.unwrap(); assert_eq!(manager.state(), SyncState::Synced); assert!(events.is_empty()); @@ -3119,4 +2926,186 @@ mod tests { events ); } + + /// A restart can leave filters stored only near the tip while the wallet + /// still needs a range far below them. Those tip-region filters are + /// unreachable for the scan, so `start_download` must discard them and + /// restart the download from `scan_start` rather than from + /// `stored_filters_tip + 1`. + #[tokio::test] + async fn test_start_download_discards_stored_filters_above_scan_start() { + let mut manager = create_test_manager().await; + + // Block headers cover the whole range so send_pending can resolve stop hashes. + let headers = dashcore::block::Header::dummy_batch(0..1001); + manager + .header_storage + .write() + .await + .store_headers( + &headers.iter().map(crate::types::HashedBlockHeader::from).collect::>(), + ) + .await + .unwrap(); + + // Filters were only ever stored near the tip: 900..=1000. The heights + // below 900 were never populated. + { + let mut filter_storage = manager.filter_storage.write().await; + for height in 900..=1000u32 { + filter_storage.store_filter(height, &[height as u8; 8]).await.unwrap(); + } + } + // Restart shape: `new()` seeds stored_height from the tip watermark. + manager.progress.update_stored_height(1000); + manager.progress.update_filter_header_tip_height(1000); + manager.progress.update_target_height(1000); + + // Wallet committed far below the stored region, so scan_start = 100. + manager.wallet.write().await.update_wallet_synced_height(&MOCK_WALLET_ID, 99); + + let (mock, network) = test_network_with_mock().await; + let events = manager.start_download(&network).await.unwrap(); + + assert!(events.is_empty()); + assert_eq!(manager.state(), SyncState::Syncing); + + // The unreachable tip-region filters were discarded entirely... + assert_eq!(manager.filter_storage.read().await.filter_tip_height().await.unwrap(), 0); + assert_eq!(manager.filter_storage.read().await.filter_start_height().await, None); + + // ...nothing was preloaded into the initial batch... + let batch = manager.active_batches.get(&100).expect("initial batch at scan_start"); + assert!(batch.filters().is_empty()); + assert!(!batch.verified()); + assert!(!batch.scanned()); + assert_eq!(batch.end_height(), 1000); + + // ...scan gating no longer sees the stale tip watermark... + assert_eq!(manager.progress.stored_height(), 0); + + // ...and both the store cursor and the download restart from + // scan_start rather than stored_filters_tip + 1. + assert_eq!(manager.next_batch_to_store, 100); + assert!( + declared_filter_starts(&mock).contains(&100), + "expected a GetCFilters declared from scan_start, got {:?}", + declared_filter_starts(&mock) + ); + } + + /// Counterpart to the sparse-storage case: when the stored filter range + /// actually reaches down to `scan_start`, the preload happens exactly as + /// before and nothing is discarded. + #[tokio::test] + async fn test_start_download_preloads_when_stored_filters_cover_scan_start() { + let mut manager = create_test_manager().await; + + let headers = dashcore::block::Header::dummy_batch(0..1001); + manager + .header_storage + .write() + .await + .store_headers( + &headers.iter().map(crate::types::HashedBlockHeader::from).collect::>(), + ) + .await + .unwrap(); + + // Filters stored densely from genesis through the tip. + { + let mut filter_storage = manager.filter_storage.write().await; + for height in 0..=1000u32 { + filter_storage.store_filter(height, &[height as u8; 8]).await.unwrap(); + } + } + manager.progress.update_stored_height(1000); + manager.progress.update_filter_header_tip_height(1000); + manager.progress.update_target_height(1000); + + manager.wallet.write().await.update_wallet_synced_height(&MOCK_WALLET_ID, 99); + + let (mock, network) = test_network_with_mock().await; + manager.start_download(&network).await.unwrap(); + + assert_eq!(manager.state(), SyncState::Syncing); + + // Storage is untouched. + assert_eq!(manager.filter_storage.read().await.filter_tip_height().await.unwrap(), 1000); + assert_eq!(manager.filter_storage.read().await.filter_start_height().await, Some(0)); + + // The stored range 100..=1000 was preloaded and the batch is verified + // and scanned immediately. + let batch = manager.active_batches.get(&100).expect("initial batch at scan_start"); + assert_eq!(batch.filters().len(), 901); + assert!(batch.verified()); + assert!(batch.scanned()); + assert_eq!(manager.progress.stored_height(), 1000); + + // Nothing left to download: everything through the filter header tip + // is already stored. + assert_eq!(manager.next_batch_to_store, 1001); + assert!( + declared_filter_starts(&mock).is_empty(), + "no filter request expected, got {:?}", + declared_filter_starts(&mock) + ); + } + + /// Genesis-only storage: a single filter at height 0, scanning from 0. + /// `filter_tip_height` collapses to 0 for both an empty store and this + /// one, so gating the preload on `tip > 0` would misread it as empty and + /// needlessly re-download height 0. The stored start (Some(0)) must drive + /// the decision: the filter is preloaded and nothing is discarded or + /// re-requested. Regtest can produce exactly this shape. + #[tokio::test] + async fn test_start_download_preloads_genesis_only_stored_filter() { + let mut manager = create_test_manager().await; + + let headers = dashcore::block::Header::dummy_batch(0..1); + manager + .header_storage + .write() + .await + .store_headers( + &headers.iter().map(crate::types::HashedBlockHeader::from).collect::>(), + ) + .await + .unwrap(); + + // Exactly one filter stored, at height 0. + manager.filter_storage.write().await.store_filter(0, &[0u8; 8]).await.unwrap(); + assert_eq!(manager.filter_storage.read().await.filter_tip_height().await.unwrap(), 0); + assert_eq!(manager.filter_storage.read().await.filter_start_height().await, Some(0)); + + // Restart shape at the genesis tip. + manager.progress.update_stored_height(0); + manager.progress.update_filter_header_tip_height(0); + manager.progress.update_target_height(0); + // Wallet at genesis: scan_start = 0. + + let (mock, network) = test_network_with_mock().await; + manager.start_download(&network).await.unwrap(); + + // The lone filter was NOT discarded... + assert_eq!(manager.filter_storage.read().await.filter_tip_height().await.unwrap(), 0); + assert_eq!(manager.filter_storage.read().await.filter_start_height().await, Some(0)); + + // ...it was preloaded into the initial batch, which is verified and + // scanned since the whole (single-height) range is covered... + let batch = manager.active_batches.get(&0).expect("initial batch at scan_start"); + assert_eq!(batch.filters().len(), 1); + assert!(batch.verified()); + assert!(batch.scanned()); + assert_eq!(manager.progress.stored_height(), 0); + + // ...and the download frontier sits above the stored tip, so no + // filter request goes out for the already-stored genesis height. + assert_eq!(manager.next_batch_to_store, 1); + assert!( + declared_filter_starts(&mock).is_empty(), + "no filter request expected for the genesis-only store, got {:?}", + declared_filter_starts(&mock) + ); + } } diff --git a/dash-spv/src/sync/filters/pipeline.rs b/dash-spv/src/sync/filters/pipeline.rs index 03c0bba9f..1a85c8872 100644 --- a/dash-spv/src/sync/filters/pipeline.rs +++ b/dash-spv/src/sync/filters/pipeline.rs @@ -1,48 +1,38 @@ //! CFilters pipeline implementation. //! //! Handles pipelined download of compact block filters (BIP 157/158). -//! Uses DownloadCoordinator for batch-level tracking, with additional +//! Declares wanted batches to the network manager's request broker, plus //! per-batch tracking for individual filter responses. //! //! Filters are buffered in a HashMap until the entire batch //! is complete, enabling batch verification and direct wallet matching. use std::collections::{BTreeSet, HashMap}; -use std::time::Duration; +use std::sync::Arc; +use dashcore::network::message::NetworkMessage; +use dashcore::network::message_filter::GetCFilters; use dashcore::BlockHash; -use crate::error::{SyncError, SyncResult}; -use crate::network::RequestSender; +use crate::error::SyncResult; +use crate::network::NetworkManager; use crate::storage::BlockHeaderStorage; -use crate::sync::download_coordinator::{DownloadConfig, DownloadCoordinator}; use crate::sync::filters::batch::FiltersBatch; use crate::sync::filters::batch_tracker::BatchTracker; /// Batch size for filter requests. const FILTER_BATCH_SIZE: u32 = 1000; -/// Maximum concurrent filter batch requests. -const MAX_CONCURRENT_FILTER_BATCHES: usize = 20; - -/// Timeout for filter batch requests. -/// Each batch requires 1000 individual filter messages, so allow plenty of time. -const FILTER_TIMEOUT: Duration = Duration::from_secs(30); - /// Pipeline for downloading compact block filters. -/// -/// Uses DownloadCoordinator for batch-level download mechanics, -/// with BatchTracker for tracking individual filters within -/// each batch. -/// -/// Filters are buffered until the entire batch is complete, then returned -/// via `take_completed_batches()` for verification and matching. #[derive(Debug)] pub(super) struct FiltersPipeline { - /// Core coordinator tracks batch start heights. - coordinator: DownloadCoordinator, /// Tracks individual filter receipts per batch (start_height -> tracker). + /// Doubles as the "which batches are still wanted?" set. batch_trackers: HashMap, + /// Cached stop hash per wanted batch (start_height -> stop_hash), so + /// re-declaring the wanted set each tick doesn't re-read storage per batch. + /// Filled lazily in `send_pending` as headers become available. + batch_stops: HashMap, /// Completed filter batches. completed_batches: BTreeSet, /// Target height for sync. @@ -63,12 +53,8 @@ impl FiltersPipeline { /// Create a new CFilters pipeline. pub(super) fn new() -> Self { Self { - coordinator: DownloadCoordinator::new( - DownloadConfig::default() - .with_max_concurrent(MAX_CONCURRENT_FILTER_BATCHES) - .with_timeout(FILTER_TIMEOUT), - ), batch_trackers: HashMap::new(), + batch_stops: HashMap::new(), completed_batches: BTreeSet::new(), target_height: 0, filters_received: 0, @@ -76,9 +62,9 @@ impl FiltersPipeline { } } - /// Returns true if the pipeline has no in-flight or pending work. + /// Returns true if the pipeline has no wanted batches left. pub(super) fn is_idle(&self) -> bool { - self.coordinator.active_count() == 0 && self.coordinator.pending_count() == 0 + self.batch_trackers.is_empty() } /// Take completed batches with their buffered filter data for processing. @@ -88,19 +74,18 @@ impl FiltersPipeline { /// Initialize the pipeline for a sync range. /// - /// Pre-queues all batches for the range using the coordinator's pending queue. + /// Creates a tracker for every batch in the range; `send_pending` then + /// declares them to the broker. pub(super) fn init(&mut self, start_height: u32, target_height: u32) { - self.coordinator.clear(); self.batch_trackers.clear(); + self.batch_stops.clear(); self.completed_batches.clear(); self.target_height = target_height; self.highest_received = start_height.saturating_sub(1); self.filters_received = 0; - // Pre-queue all batches let mut current = start_height; while current <= target_height { - self.coordinator.enqueue([current]); let batch_end = (current + FILTER_BATCH_SIZE - 1).min(target_height); self.batch_trackers.insert(current, BatchTracker::new(batch_end)); current = batch_end + 1; @@ -109,7 +94,7 @@ impl FiltersPipeline { /// Extend the target height without resetting pipeline state. /// - /// Queues additional batches from the old target boundary to the new target. + /// Adds trackers for the batches from the old target boundary to the new target. pub(super) fn extend_target(&mut self, new_target: u32) { if new_target <= self.target_height { return; @@ -118,87 +103,85 @@ impl FiltersPipeline { let old_target = self.target_height; self.target_height = new_target; - // Queue new batches from (old_target + 1) to new_target let mut current = old_target + 1; while current <= new_target { - self.coordinator.enqueue([current]); let batch_end = (current + FILTER_BATCH_SIZE - 1).min(new_target); self.batch_trackers.insert(current, BatchTracker::new(batch_end)); current = batch_end + 1; } } - /// Send pending filter requests up to the concurrency limit. + /// Declare every wanted batch to the network manager. + /// + /// Resolves (and caches) each batch's stop hash from storage the first time + /// it becomes available, then declares all resolved batches to the broker, + /// which de-duplicates in-flight ones. Batches whose stop header isn't stored + /// yet are simply skipped this round and retried next tick. pub(super) async fn send_pending( &mut self, - requests: &RequestSender, + network: &Arc, storage: &impl BlockHeaderStorage, ) -> SyncResult { - let count = self.coordinator.available_to_send(); - if count == 0 { + if self.batch_trackers.is_empty() { return Ok(0); } - let start_heights = self.coordinator.take_pending(count); - let mut sent = 0; - - for start_height in start_heights { - let batch_end = match self.batch_trackers.get(&start_height) { - Some(tracker) => tracker.end_height(), - None => { - return Err(SyncError::InvalidState(format!( - "missing batch tracker for start_height {}", - start_height - ))); + // Resolve stop hashes for any wanted batch we haven't cached yet. + let uncached: Vec<(u32, u32)> = self + .batch_trackers + .iter() + .filter(|(start, _)| !self.batch_stops.contains_key(start)) + .map(|(&start, tracker)| (start, tracker.end_height())) + .collect(); + for (start, batch_end) in uncached { + match storage.get_header(batch_end).await { + Ok(Some(h)) => { + self.batch_stops.insert(start, *h.hash()); } - }; - - // Get stop hash for this batch. If the header isn't available yet, - // re-queue for the next tick instead of losing the batch permanently. - let stop_hash = match storage.get_header(batch_end).await { - Ok(Some(h)) => *h.hash(), Ok(None) => { tracing::debug!( - "Header at height {} not yet available, re-queuing filter batch {}", + "Header at height {} not yet available, deferring filter batch {}", batch_end, - start_height + start ); - self.coordinator.enqueue([start_height]); - continue; } Err(e) => { - tracing::warn!( - "Error reading header at height {}, re-queuing filter batch {}: {}", - batch_end, - start_height, - e - ); - self.coordinator.enqueue([start_height]); - continue; + tracing::warn!("Error reading header at height {}: {}", batch_end, e); } - }; - - requests.request_filters(start_height, stop_hash)?; - - self.coordinator.mark_sent(&[start_height]); - - tracing::trace!( - "Sent GetCFilters: {} to {} ({} active batches)", - start_height, - batch_end, - self.coordinator.active_count() - ); + } + } - sent += 1; + // Declare every wanted batch whose stop hash is known, lowest-first for a + // deterministic fan-out. The broker de-duplicates, so re-declaring one + // already in flight is a no-op. + let mut ready: Vec<(u32, BlockHash)> = self + .batch_stops + .iter() + .filter(|(start, _)| self.batch_trackers.contains_key(start)) + .map(|(&start, &stop)| (start, stop)) + .collect(); + ready.sort_unstable_by_key(|(start, _)| *start); + + let n = ready.len(); + for (start_height, stop_hash) in ready { + network + .send(NetworkMessage::GetCFilters(GetCFilters { + filter_type: 0u8, + start_height, + stop_hash, + })) + .await; } - Ok(sent) + Ok(n) } /// Handle a received CFilter message with filter data. /// /// Buffers the filter data for batch verification and wallet matching. - /// Returns `Some(height)` when a batch completes, `None` otherwise. + /// Returns `Some(batch_start)` when a batch completes (the `GetCFilters` + /// request key, so the caller can tell the network manager it was answered), + /// `None` otherwise. pub(super) fn receive_with_data( &mut self, height: u32, @@ -234,10 +217,11 @@ impl FiltersPipeline { let filters = self.batch_trackers.get_mut(&batch_start).map(|t| t.take_filters()).unwrap_or_default(); + // Out of the wanted set: drop its tracker and cached stop hash so the next + // `send_pending` won't re-declare it (the manager also tells the broker it + // was answered, stopping any in-flight retry). self.batch_trackers.remove(&batch_start); - if !self.coordinator.receive(&batch_start) { - self.coordinator.cancel_pending(&batch_start); - } + self.batch_stops.remove(&batch_start); tracing::info!( "Filter batch {}-{} complete ({} filters)", @@ -248,7 +232,7 @@ impl FiltersPipeline { let batch = FiltersBatch::new(batch_start, end_height, filters); self.completed_batches.insert(batch); - Some(height) + Some(batch_start) } /// Find which batch a filter height belongs to. @@ -260,92 +244,23 @@ impl FiltersPipeline { } None } - - /// Check for timed out batches and handle retries. - /// - /// Does not remove batch trackers — keeps them to receive any late-arriving filters. - pub(super) fn handle_timeouts(&mut self) { - for start in self.coordinator.check_timeouts() { - self.coordinator.enqueue_retry(start); - } - } - - /// Move in-flight `getcfilters` requests back to pending after a peer - /// disconnect so the next `send_pending` reissues them to the new peer. - /// Per-batch trackers and any partially-received filters within them are - /// preserved — `BatchTracker::insert_filter` is idempotent, so duplicates - /// from the new peer are harmless. - pub(super) fn requeue_in_flight(&mut self) { - self.coordinator.requeue_in_flight(); - } } #[cfg(test)] mod tests { use super::*; - use crate::network::{NetworkRequest, RequestSender}; - use crate::storage::{PersistentBlockHeaderStorage, PersistentStorage}; - use dashcore::bip158::BlockFilter; use dashcore::block::Header; - use dashcore::network::message::NetworkMessage; use dashcore_hashes::Hash; use key_wallet_manager::FilterMatchKey; - use std::time::Duration; - use tempfile::TempDir; - use tokio::sync::mpsc::unbounded_channel; - // ========================================================================= - // Helper functions - // ========================================================================= - - /// Create a pipeline with short timeout for testing timeouts. - fn create_pipeline_with_short_timeout() -> FiltersPipeline { - FiltersPipeline { - coordinator: DownloadCoordinator::new( - DownloadConfig::default().with_timeout(Duration::from_millis(1)), - ), - batch_trackers: HashMap::new(), - completed_batches: BTreeSet::new(), - target_height: 0, - filters_received: 0, - highest_received: 0, - } - } - - /// Create a pipeline with max_concurrent=2 for testing deferred sends. - fn create_pipeline_with_low_concurrency() -> FiltersPipeline { - FiltersPipeline { - coordinator: DownloadCoordinator::new( - DownloadConfig::default().with_max_concurrent(2).with_timeout(FILTER_TIMEOUT), - ), - batch_trackers: HashMap::new(), - completed_batches: BTreeSet::new(), - target_height: 0, - filters_received: 0, - highest_received: 0, - } - } - - /// Create a test request sender with its receiver. - fn create_test_request_sender( - ) -> (RequestSender, tokio::sync::mpsc::UnboundedReceiver) { - let (tx, rx) = unbounded_channel(); - (RequestSender::new(tx), rx) - } /// Generate dummy filter data for testing. fn dummy_filter_data(height: u32) -> Vec { vec![height as u8, (height >> 8) as u8, 0x01, 0x02] } - // ========================================================================= - // FiltersPipeline Construction Tests - // ========================================================================= - #[test] fn test_pipeline_new() { let pipeline = FiltersPipeline::new(); - - assert_eq!(pipeline.coordinator.active_count(), 0); assert!(pipeline.batch_trackers.is_empty()); assert!(pipeline.completed_batches.is_empty()); assert_eq!(pipeline.target_height, 0); @@ -362,26 +277,13 @@ mod tests { assert!(!pipeline.is_idle()); } - #[test] - fn test_pipeline_default_trait() { - let default_pipeline = FiltersPipeline::default(); - let new_pipeline = FiltersPipeline::new(); - - assert_eq!( - default_pipeline.coordinator.active_count(), - new_pipeline.coordinator.active_count() - ); - assert_eq!(default_pipeline.target_height, new_pipeline.target_height); - } - #[test] fn test_pipeline_init() { let mut pipeline = FiltersPipeline::new(); - pipeline.init(100, 500); - // Should have 1 batch queued (100-500 is 401 filters, fits in 1 batch) - assert_eq!(pipeline.coordinator.pending_count(), 1); + // 100..=500 fits in a single batch. + assert_eq!(pipeline.batch_trackers.len(), 1); assert_eq!(pipeline.target_height, 500); assert_eq!(pipeline.highest_received, 99); assert_eq!(pipeline.filters_received, 0); @@ -391,78 +293,51 @@ mod tests { fn test_pipeline_init_resets_state() { let mut pipeline = FiltersPipeline::new(); - // Add some state pipeline.batch_trackers.insert(0, BatchTracker::new(99)); pipeline.completed_batches.insert(FiltersBatch::new(100, 199, HashMap::new())); - pipeline.coordinator.mark_sent(&[0]); pipeline.filters_received = 50; - // Init should clear old state and set up new batches pipeline.init(200, 300); assert!(pipeline.completed_batches.is_empty()); - assert_eq!(pipeline.coordinator.active_count(), 0); assert_eq!(pipeline.filters_received, 0); - // 1 batch queued for heights 200-300 - assert_eq!(pipeline.coordinator.pending_count(), 1); assert_eq!(pipeline.batch_trackers.len(), 1); assert_eq!(pipeline.batch_trackers.get(&200).unwrap().end_height(), 300); assert_eq!(pipeline.target_height, 300); } - // ========================================================================= - // Target Extension Tests - // ========================================================================= + #[test] + fn test_init_creates_contiguous_batches() { + let mut pipeline = FiltersPipeline::new(); + pipeline.init(0, 2500); + + // 0-999, 1000-1999, 2000-2500 + let mut ranges: Vec<(u32, u32)> = + pipeline.batch_trackers.iter().map(|(&s, t)| (s, t.end_height())).collect(); + ranges.sort_unstable(); + assert_eq!(ranges, vec![(0, 999), (1000, 1999), (2000, 2500)]); + } #[test] fn test_extend_target_increases() { let mut pipeline = FiltersPipeline::new(); pipeline.init(0, 100); - pipeline.extend_target(200); - assert_eq!(pipeline.target_height, 200); } - #[tokio::test] - async fn test_extend_target_contiguous_batches() { - // init's last batch is truncated (3000-3500), extend_target fills from 3501. - // Verify all batches are contiguous after sending. - let headers = Header::dummy_batch(0..6000); - let tmp_dir = TempDir::new().unwrap(); - let mut storage = PersistentBlockHeaderStorage::open(tmp_dir.path()).await.unwrap(); - storage - .store_headers( - &headers.iter().map(crate::types::HashedBlockHeader::from).collect::>(), - ) - .await - .unwrap(); - + #[test] + fn test_extend_target_contiguous_batches() { let mut pipeline = FiltersPipeline::new(); pipeline.init(0, 3500); pipeline.extend_target(5000); - let (sender, _rx) = create_test_request_sender(); - pipeline.send_pending(&sender, &storage).await.unwrap(); - - let mut ranges: Vec<(u32, u32)> = pipeline - .batch_trackers - .iter() - .map(|(&start, tracker)| (start, tracker.end_height())) - .collect(); - ranges.sort_by_key(|&(start, _)| start); + let mut ranges: Vec<(u32, u32)> = + pipeline.batch_trackers.iter().map(|(&s, t)| (s, t.end_height())).collect(); + ranges.sort_unstable_by_key(|&(s, _)| s); - // Verify contiguous: 0-999, 1000-1999, 2000-2999, 3000-3500, 3501-4500, 4501-5000 for window in ranges.windows(2) { - assert_eq!( - window[0].1 + 1, - window[1].0, - "gap or overlap between batches: {}-{} and {}-{}", - window[0].0, - window[0].1, - window[1].0, - window[1].1 - ); + assert_eq!(window[0].1 + 1, window[1].0, "gap or overlap between batches"); } assert_eq!(ranges[3], (3000, 3500)); assert_eq!(ranges[4], (3501, 4500)); @@ -472,97 +347,24 @@ mod tests { fn test_extend_target_ignores_lower() { let mut pipeline = FiltersPipeline::new(); pipeline.init(0, 100); - pipeline.extend_target(50); - assert_eq!(pipeline.target_height, 100); - pipeline.extend_target(100); - assert_eq!(pipeline.target_height, 100); } - // ========================================================================= - // Receive Tests - // ========================================================================= - - #[test] - fn test_requeue_in_flight_preserves_partial_batch_receipts() { - let mut pipeline = FiltersPipeline::new(); - pipeline.target_height = 99; - - // One batch in-flight (start_height 0). Receive a filter so the - // tracker has partial state. - pipeline.batch_trackers.insert(0, BatchTracker::new(99)); - pipeline.coordinator.mark_sent(&[0]); - let hash = Header::dummy(50).block_hash(); - pipeline.receive_with_data(50, hash, &dummy_filter_data(50)); - assert_eq!(pipeline.filters_received, 1); - assert_eq!(pipeline.coordinator.active_count(), 1); - - pipeline.requeue_in_flight(); - - // Batch is back in pending; tracker (and the partial filter inside it) - // is preserved so the new peer's response merges idempotently. - assert_eq!(pipeline.coordinator.active_count(), 0); - assert_eq!(pipeline.coordinator.pending_count(), 1); - let tracker = pipeline.batch_trackers.get(&0).expect("tracker preserved"); - assert_eq!(tracker.received(), 1); - assert_eq!(pipeline.filters_received, 1); - assert_eq!(pipeline.highest_received, 50); - } - - #[test] - fn test_late_filter_after_requeue_completes_batch_without_orphaning_pending() { - // Regression: a late `cfilter` from the disconnected peer can complete - // a batch after `requeue_in_flight` moved it back to pending. Without - // the cancel-pending hook, the key would linger in `pending` while the - // tracker was gone, and the next `send_pending` would error with - // `SyncError::InvalidState`. - let mut pipeline = FiltersPipeline::new(); - pipeline.target_height = 2; - - pipeline.batch_trackers.insert(0, BatchTracker::new(2)); - pipeline.coordinator.mark_sent(&[0]); - - // Two filters arrive before disconnect. - for h in 0..=1 { - let hash = Header::dummy(h).block_hash(); - pipeline.receive_with_data(h, hash, &dummy_filter_data(h)); - } - - pipeline.requeue_in_flight(); - assert_eq!(pipeline.coordinator.pending_count(), 1); - assert_eq!(pipeline.coordinator.active_count(), 0); - - // Late buffered filter from old peer completes the batch. - let hash = Header::dummy(2).block_hash(); - pipeline.receive_with_data(2, hash, &dummy_filter_data(2)); - - assert_eq!(pipeline.completed_batches.len(), 1); - assert!(pipeline.batch_trackers.is_empty()); - // The orphaned pending key must be gone so `send_pending` does not - // resurrect a finished batch. - assert_eq!(pipeline.coordinator.pending_count(), 0); - assert_eq!(pipeline.coordinator.active_count(), 0); - } - #[test] fn test_receive_single_filter() { let mut pipeline = FiltersPipeline::new(); pipeline.target_height = 99; - - // Set up batch tracker manually (simulating an in-flight batch) pipeline.batch_trackers.insert(0, BatchTracker::new(99)); - pipeline.coordinator.mark_sent(&[0]); let height = 50; let hash = Header::dummy(height).block_hash(); let result = pipeline.receive_with_data(height, hash, &dummy_filter_data(height)); - // Returns None since batch is not complete (only 1 of 100 filters received) + // Batch not complete (1 of 100 filters). assert_eq!(result, None); - // But counters are updated assert_eq!(pipeline.filters_received, 1); assert_eq!(pipeline.highest_received, 50); } @@ -572,7 +374,6 @@ mod tests { let mut pipeline = FiltersPipeline::new(); pipeline.target_height = 99; - // No batch tracker set up - filter is unexpected let hash = Header::dummy(50).block_hash(); let result = pipeline.receive_with_data(50, hash, &dummy_filter_data(50)); @@ -581,26 +382,26 @@ mod tests { } #[test] - fn test_receive_batch_completion() { + fn test_receive_batch_completion_returns_batch_start() { let mut pipeline = FiltersPipeline::new(); pipeline.target_height = 2; - - // Set up a small batch (3 filters: 0, 1, 2) pipeline.batch_trackers.insert(0, BatchTracker::new(2)); - pipeline.coordinator.mark_sent(&[0]); - // Receive all filters + let mut completed_at = None; for h in 0..=2 { let hash = Header::dummy(h).block_hash(); - pipeline.receive_with_data(h, hash, &dummy_filter_data(h)); + if let Some(start) = pipeline.receive_with_data(h, hash, &dummy_filter_data(h)) { + completed_at = Some(start); + } } - // Batch should be complete and moved to completed_batches + // Completing filter returns the batch START (the GetCFilters request key). + assert_eq!(completed_at, Some(0)); assert!(pipeline.batch_trackers.is_empty()); + assert!(!pipeline.batch_stops.contains_key(&0)); assert_eq!(pipeline.completed_batches.len(), 1); let completed = pipeline.take_completed_batches(); - assert_eq!(completed.len(), 1); let batch = completed.into_iter().next().unwrap(); assert_eq!(batch.start_height(), 0); assert_eq!(batch.end_height(), 2); @@ -611,111 +412,48 @@ mod tests { fn test_receive_out_of_order() { let mut pipeline = FiltersPipeline::new(); pipeline.target_height = 4; - pipeline.batch_trackers.insert(0, BatchTracker::new(4)); - pipeline.coordinator.mark_sent(&[0]); - // Receive out of order for h in [3, 1, 4, 0, 2] { let hash = Header::dummy(h).block_hash(); pipeline.receive_with_data(h, hash, &dummy_filter_data(h)); } - // Should complete successfully assert!(pipeline.batch_trackers.is_empty()); assert_eq!(pipeline.completed_batches.len(), 1); } - #[test] - fn test_receive_updates_counters() { - let mut pipeline = FiltersPipeline::new(); - pipeline.target_height = 99; - - pipeline.batch_trackers.insert(0, BatchTracker::new(99)); - pipeline.coordinator.mark_sent(&[0]); - - // Receive some filters - for h in [10, 5, 20, 15] { - let hash = Header::dummy(h).block_hash(); - pipeline.receive_with_data(h, hash, &dummy_filter_data(h)); - } - - assert_eq!(pipeline.filters_received, 4); - assert_eq!(pipeline.highest_received, 20); - } - - #[test] - fn test_receive_small_batch_at_target() { - let mut pipeline = FiltersPipeline::new(); - pipeline.target_height = 1005; - - // Small batch of 6 filters (1000-1005) - pipeline.batch_trackers.insert(1000, BatchTracker::new(1005)); - pipeline.coordinator.mark_sent(&[1000]); - - // Receive all 6 filters - for h in 1000..=1005 { - let hash = Header::dummy(h).block_hash(); - pipeline.receive_with_data(h, hash, &dummy_filter_data(h)); - } - - assert_eq!(pipeline.completed_batches.len(), 1); - let batch = pipeline.completed_batches.iter().next().unwrap(); - assert_eq!(batch.filters().len(), 6); - } - #[test] fn test_receive_multiple_batches() { let mut pipeline = FiltersPipeline::new(); pipeline.target_height = 9; - - // Set up two batches manually pipeline.batch_trackers.insert(0, BatchTracker::new(4)); pipeline.batch_trackers.insert(5, BatchTracker::new(9)); - pipeline.coordinator.mark_sent(&[0, 5]); - // Receive first batch for h in 0..=4 { let hash = Header::dummy(h).block_hash(); pipeline.receive_with_data(h, hash, &dummy_filter_data(h)); } - assert_eq!(pipeline.completed_batches.len(), 1); assert_eq!(pipeline.batch_trackers.len(), 1); - // Receive second batch for h in 5..=9 { let hash = Header::dummy(h).block_hash(); pipeline.receive_with_data(h, hash, &dummy_filter_data(h)); } - assert_eq!(pipeline.completed_batches.len(), 2); assert!(pipeline.batch_trackers.is_empty()); } - // ========================================================================= - // find_batch_for_height Tests - // ========================================================================= - #[test] - fn test_find_batch_for_height_found() { + fn test_find_batch_for_height() { let mut pipeline = FiltersPipeline::new(); pipeline.batch_trackers.insert(0, BatchTracker::new(999)); pipeline.batch_trackers.insert(1000, BatchTracker::new(1999)); assert_eq!(pipeline.find_batch_for_height(500), Some(0)); assert_eq!(pipeline.find_batch_for_height(1500), Some(1000)); - } - - #[test] - fn test_find_batch_for_height_none() { - let mut pipeline = FiltersPipeline::new(); - pipeline.batch_trackers.insert(100, BatchTracker::new(199)); - - // Below range - assert_eq!(pipeline.find_batch_for_height(50), None); - // Above range - assert_eq!(pipeline.find_batch_for_height(250), None); + assert_eq!(pipeline.find_batch_for_height(5000), None); } #[test] @@ -723,349 +461,15 @@ mod tests { let mut pipeline = FiltersPipeline::new(); pipeline.batch_trackers.insert(100, BatchTracker::new(199)); - // First height in batch assert_eq!(pipeline.find_batch_for_height(100), Some(100)); - // Last height in batch assert_eq!(pipeline.find_batch_for_height(199), Some(100)); - } - - // ========================================================================= - // Timeout Tests - // ========================================================================= - - #[test] - fn test_handle_timeouts_no_batches() { - let mut pipeline = FiltersPipeline::new(); - pipeline.handle_timeouts(); - } - - #[test] - fn test_handle_timeouts_requeue() { - let mut pipeline = create_pipeline_with_short_timeout(); - pipeline.target_height = 999; - - // Set up batch and mark as in-flight (simulating a sent request) - pipeline.batch_trackers.insert(0, BatchTracker::new(999)); - pipeline.coordinator.mark_sent(&[0]); - - // Wait for timeout - std::thread::sleep(Duration::from_millis(5)); - - pipeline.handle_timeouts(); - - // Batch should be re-queued in coordinator's pending queue - assert_eq!(pipeline.coordinator.pending_count(), 1); - assert_eq!(pipeline.coordinator.active_count(), 0); - } - - #[test] - fn test_handle_timeouts_keeps_tracker() { - let mut pipeline = create_pipeline_with_short_timeout(); - pipeline.target_height = 99; - - pipeline.batch_trackers.insert(0, BatchTracker::new(99)); - pipeline.coordinator.mark_sent(&[0]); - - // Receive some filters before timeout - for h in 0..10 { - let hash = Header::dummy(h).block_hash(); - pipeline.receive_with_data(h, hash, &dummy_filter_data(h)); - } - - std::thread::sleep(Duration::from_millis(5)); - - pipeline.handle_timeouts(); - - // Should timeout but tracker is preserved for late arrivals - assert!(pipeline.batch_trackers.contains_key(&0)); - assert_eq!(pipeline.batch_trackers.get(&0).unwrap().received(), 10); - } - - #[test] - fn test_timeout_does_not_duplicate_inflight_batches() { - // This test verifies the bug fix: when an early batch times out, - // only that batch is re-queued, not later in-flight batches. - let mut pipeline = FiltersPipeline { - coordinator: DownloadCoordinator::new( - DownloadConfig::default() - .with_timeout(Duration::from_millis(1)) - .with_max_concurrent(10), - ), - batch_trackers: HashMap::new(), - completed_batches: BTreeSet::new(), - target_height: 2999, - filters_received: 0, - highest_received: 0, - }; - - // Simulate 3 in-flight batches: 0-999, 1000-1999, 2000-2999 - pipeline.batch_trackers.insert(0, BatchTracker::new(999)); - pipeline.batch_trackers.insert(1000, BatchTracker::new(1999)); - pipeline.batch_trackers.insert(2000, BatchTracker::new(2999)); - pipeline.coordinator.mark_sent(&[0, 1000, 2000]); - - assert_eq!(pipeline.coordinator.active_count(), 3); - assert_eq!(pipeline.coordinator.pending_count(), 0); - - // Wait for timeout - std::thread::sleep(Duration::from_millis(5)); - - // Handle timeouts - all 3 should timeout and be re-queued - pipeline.handle_timeouts(); - - // All 3 batches should be in the pending queue, not duplicated - assert_eq!(pipeline.coordinator.pending_count(), 3); - assert_eq!(pipeline.coordinator.active_count(), 0); - - // Take pending items - should get exactly 3, not more - let pending = pipeline.coordinator.take_pending(10); - assert_eq!(pending.len(), 3); - assert!(pending.contains(&0)); - assert!(pending.contains(&1000)); - assert!(pending.contains(&2000)); - } - - // ========================================================================= - // send_pending Tests - // ========================================================================= - - #[tokio::test] - async fn test_send_pending_single_batch() { - let headers = Header::dummy_batch(0..1000); - let tmp_dir = TempDir::new().unwrap(); - let mut storage = PersistentBlockHeaderStorage::open(tmp_dir.path()).await.unwrap(); - storage - .store_headers( - &headers.iter().map(crate::types::HashedBlockHeader::from).collect::>(), - ) - .await - .unwrap(); - - let mut pipeline = FiltersPipeline::new(); - pipeline.init(0, 999); - - let (sender, mut rx) = create_test_request_sender(); - - let count = pipeline.send_pending(&sender, &storage).await.unwrap(); - - assert_eq!(count, 1); - assert_eq!(pipeline.coordinator.active_count(), 1); - assert!(pipeline.batch_trackers.contains_key(&0)); - // No more pending since the single batch was sent - assert_eq!(pipeline.coordinator.pending_count(), 0); - - // Verify message was sent - let request = rx.try_recv().unwrap(); - let NetworkRequest::SendMessage(msg) = request else { - panic!("Expected SendMessage variant"); - }; - if let NetworkMessage::GetCFilters(gcf) = msg { - assert_eq!(gcf.start_height, 0); - assert_eq!(gcf.filter_type, 0); - } else { - panic!("Expected GetCFilters message"); - } - } - - #[tokio::test] - async fn test_send_pending_respects_limit() { - // Create enough headers for many batches - let headers = Header::dummy_batch(0..25000); - let tmp_dir = TempDir::new().unwrap(); - let mut storage = PersistentBlockHeaderStorage::open(tmp_dir.path()).await.unwrap(); - storage - .store_headers( - &headers.iter().map(crate::types::HashedBlockHeader::from).collect::>(), - ) - .await - .unwrap(); - - let mut pipeline = FiltersPipeline::new(); - pipeline.init(0, 24999); - - let (sender, _rx) = create_test_request_sender(); - - let count = pipeline.send_pending(&sender, &storage).await.unwrap(); - - // 25 batches needed, but only 20 can be in-flight at once - assert_eq!(count, MAX_CONCURRENT_FILTER_BATCHES); - assert_eq!(pipeline.coordinator.active_count(), MAX_CONCURRENT_FILTER_BATCHES); - assert_eq!(pipeline.batch_trackers.len(), 25); - assert_eq!(pipeline.coordinator.pending_count(), 5); - } - - #[tokio::test] - async fn test_send_pending_calculates_end() { - let headers = Header::dummy_batch(0..1500); - let tmp_dir = TempDir::new().unwrap(); - let mut storage = PersistentBlockHeaderStorage::open(tmp_dir.path()).await.unwrap(); - storage - .store_headers( - &headers.iter().map(crate::types::HashedBlockHeader::from).collect::>(), - ) - .await - .unwrap(); - - let mut pipeline = FiltersPipeline::new(); - // Target is 1200, so second batch ends at 1200 not 1999 - pipeline.init(0, 1200); - - let (sender, _rx) = create_test_request_sender(); - - let count = pipeline.send_pending(&sender, &storage).await.unwrap(); - - assert_eq!(count, 2); - - // First batch: 0-999 - assert!(pipeline.batch_trackers.contains_key(&0)); - assert_eq!(pipeline.batch_trackers.get(&0).unwrap().end_height(), 999); - - // Second batch: 1000-1200 (capped by target) - assert!(pipeline.batch_trackers.contains_key(&1000)); - assert_eq!(pipeline.batch_trackers.get(&1000).unwrap().end_height(), 1200); - } - - #[tokio::test] - async fn test_send_pending_sends_all_queued() { - let headers = Header::dummy_batch(0..3000); - let tmp_dir = TempDir::new().unwrap(); - let mut storage = PersistentBlockHeaderStorage::open(tmp_dir.path()).await.unwrap(); - storage - .store_headers( - &headers.iter().map(crate::types::HashedBlockHeader::from).collect::>(), - ) - .await - .unwrap(); - - let mut pipeline = FiltersPipeline::new(); - pipeline.init(0, 2500); - - let (sender, _rx) = create_test_request_sender(); - - let count = pipeline.send_pending(&sender, &storage).await.unwrap(); - - // Should send all 3 batches: 0-999, 1000-1999, 2000-2500 - assert_eq!(count, 3); - assert_eq!(pipeline.coordinator.active_count(), 3); - assert_eq!(pipeline.coordinator.pending_count(), 0); - } - - #[tokio::test] - async fn test_send_pending_no_work_when_queue_empty() { - let headers = Header::dummy_batch(0..100); - let tmp_dir = TempDir::new().unwrap(); - let mut storage = PersistentBlockHeaderStorage::open(tmp_dir.path()).await.unwrap(); - storage - .store_headers( - &headers.iter().map(crate::types::HashedBlockHeader::from).collect::>(), - ) - .await - .unwrap(); - - let mut pipeline = FiltersPipeline::new(); - pipeline.init(0, 50); - - let (sender, _rx) = create_test_request_sender(); - - // First send exhausts the queue - let count = pipeline.send_pending(&sender, &storage).await.unwrap(); - assert_eq!(count, 1); - - // Second send has nothing to do - let count = pipeline.send_pending(&sender, &storage).await.unwrap(); - assert_eq!(count, 0); - } - - // ========================================================================= - // Integration Tests - // ========================================================================= - - #[tokio::test] - async fn test_full_batch_lifecycle() { - let headers = Header::dummy_batch(0..100); - let tmp_dir = TempDir::new().unwrap(); - let mut storage = PersistentBlockHeaderStorage::open(tmp_dir.path()).await.unwrap(); - storage - .store_headers( - &headers.iter().map(crate::types::HashedBlockHeader::from).collect::>(), - ) - .await - .unwrap(); - - let mut pipeline = FiltersPipeline::new(); - pipeline.init(0, 99); - - let (sender, _rx) = create_test_request_sender(); - - // Send request - let sent = pipeline.send_pending(&sender, &storage).await.unwrap(); - assert_eq!(sent, 1); - assert_eq!(pipeline.coordinator.active_count(), 1); - - // Receive all filters - for h in 0..=99 { - let hash = Header::dummy(h).block_hash(); - pipeline.receive_with_data(h, hash, &dummy_filter_data(h)); - } - - // Batch should be complete - assert_eq!(pipeline.coordinator.active_count(), 0); - assert_eq!(pipeline.completed_batches.len(), 1); - assert_eq!(pipeline.filters_received, 100); - assert_eq!(pipeline.highest_received, 99); - - // Take completed - let completed = pipeline.take_completed_batches(); - assert_eq!(completed.len(), 1); - assert!(pipeline.completed_batches.is_empty()); - } - - #[tokio::test] - async fn test_timeout_and_retry_flow() { - let headers = Header::dummy_batch(0..1000); - let tmp_dir = TempDir::new().unwrap(); - let mut storage = PersistentBlockHeaderStorage::open(tmp_dir.path()).await.unwrap(); - storage - .store_headers( - &headers.iter().map(crate::types::HashedBlockHeader::from).collect::>(), - ) - .await - .unwrap(); - - let mut pipeline = create_pipeline_with_short_timeout(); - pipeline.init(0, 999); - - let (sender, _rx) = create_test_request_sender(); - - // Send initial request - pipeline.send_pending(&sender, &storage).await.unwrap(); - assert_eq!(pipeline.coordinator.active_count(), 1); - assert_eq!(pipeline.coordinator.pending_count(), 0); - - // Wait for timeout - std::thread::sleep(Duration::from_millis(5)); - - // Handle timeout - should re-queue the batch via coordinator - pipeline.handle_timeouts(); - assert_eq!(pipeline.coordinator.pending_count(), 1); - assert_eq!(pipeline.coordinator.active_count(), 0); - - // Tracker should still exist for late arrivals - assert!(pipeline.batch_trackers.contains_key(&0)); - - // Can retry by sending again - pipeline.send_pending(&sender, &storage).await.unwrap(); - assert_eq!(pipeline.coordinator.active_count(), 1); - - // Existing tracker is reused (not replaced) - assert!(pipeline.batch_trackers.contains_key(&0)); + assert_eq!(pipeline.find_batch_for_height(50), None); + assert_eq!(pipeline.find_batch_for_height(250), None); } #[test] fn test_take_completed_batches_clears() { let mut pipeline = FiltersPipeline::new(); - - // Add some completed batches pipeline.completed_batches.insert(FiltersBatch::new(0, 99, HashMap::new())); pipeline.completed_batches.insert(FiltersBatch::new(100, 199, HashMap::new())); @@ -1076,62 +480,59 @@ mod tests { #[test] fn test_filters_batch_filters_mut() { + use dashcore::bip158::BlockFilter; let mut batch = FiltersBatch::new(0, 0, HashMap::new()); - batch .filters_mut() .insert(FilterMatchKey::new(0, BlockHash::all_zeros()), BlockFilter::new(&[0x01])); - assert_eq!(batch.filters().len(), 1); } - #[tokio::test] - async fn test_deferred_batch_keeps_end_height_after_extend() { - // init(0, 2500) creates 3 batches but only 2 can be sent (max concurrent=2). - // The boundary batch (2000-2500) stays queued. After extend_target changes - // target_height to 4000, the deferred batch must still use end_height=2500. - let headers = Header::dummy_batch(0..5000); - let tmp_dir = TempDir::new().unwrap(); - let mut storage = PersistentBlockHeaderStorage::open(tmp_dir.path()).await.unwrap(); - storage - .store_headers( - &headers.iter().map(crate::types::HashedBlockHeader::from).collect::>(), - ) - .await - .unwrap(); + /// `Default` must be the same thing as `new`, since the pipeline is created + /// through both paths. + #[test] + fn test_pipeline_default_trait() { + let default_pipeline = FiltersPipeline::default(); + let new_pipeline = FiltersPipeline::new(); - let mut pipeline = create_pipeline_with_low_concurrency(); - pipeline.init(0, 2500); - assert_eq!(pipeline.coordinator.pending_count(), 3); // 0, 1000, 2000 + assert_eq!(default_pipeline.is_idle(), new_pipeline.is_idle()); + assert_eq!(default_pipeline.target_height, new_pipeline.target_height); + assert_eq!(default_pipeline.filters_received, new_pipeline.filters_received); + assert_eq!(default_pipeline.highest_received, new_pipeline.highest_received); + } - let (sender, _rx) = create_test_request_sender(); + /// `extend_target` adds trackers for the new range without disturbing the + /// ones already wanted: an existing batch keeps the `end_height` it was + /// created with, even when the new target moves far beyond it. + #[test] + fn test_extend_target_preserves_existing_batch_end_heights() { + let mut pipeline = FiltersPipeline::new(); + pipeline.init(0, 2500); - // Only 2 batches sent, batch 2000 stays queued - pipeline.send_pending(&sender, &storage).await.unwrap(); - assert_eq!(pipeline.coordinator.active_count(), 2); - assert_eq!(pipeline.coordinator.pending_count(), 1); + // The boundary batch was truncated at the old target. + assert_eq!(pipeline.batch_trackers.get(&2000).unwrap().end_height(), 2500); - // Extend target — batch 2000's tracker must keep end_height=2500 pipeline.extend_target(4000); - // Complete batch 0 to free a slot, then send deferred batch - for h in 0..1000 { - let hash = headers[h as usize].block_hash(); - pipeline.receive_with_data(h, hash, &dummy_filter_data(h)); - } - pipeline.send_pending(&sender, &storage).await.unwrap(); - assert_eq!( pipeline.batch_trackers.get(&2000).unwrap().end_height(), 2500, - "deferred batch should use its original end height" + "an already-wanted batch must keep its original end height" ); + // ...and the extension picks up from the old boundary. + assert_eq!(pipeline.batch_trackers.get(&2501).unwrap().end_height(), 3500); } + /// End-to-end over the whole batch: declare it to the broker, feed every + /// filter in, and take the completed batch out. This is the only coverage + /// of `send_pending`, which resolves stop hashes from header storage. #[tokio::test] - async fn test_send_pending_requeues_on_missing_header() { - // Headers 0..999 exist, but NOT 1999 (stop hash for batch 1000-1999). - let headers = Header::dummy_batch(0..1000); + async fn test_full_batch_lifecycle() { + use crate::storage::{PersistentBlockHeaderStorage, PersistentStorage}; + use crate::test_utils::MockNetworkManager; + use tempfile::TempDir; + + let headers = Header::dummy_batch(0..100); let tmp_dir = TempDir::new().unwrap(); let mut storage = PersistentBlockHeaderStorage::open(tmp_dir.path()).await.unwrap(); storage @@ -1142,38 +543,39 @@ mod tests { .unwrap(); let mut pipeline = FiltersPipeline::new(); - pipeline.init(0, 1999); - assert_eq!(pipeline.coordinator.pending_count(), 2); + pipeline.init(0, 99); + assert!(!pipeline.is_idle()); - let (sender, mut rx) = create_test_request_sender(); + let mock = Arc::new(MockNetworkManager::new()); + let network: Arc = mock.clone(); - // Batch 0 succeeds (header 999 exists), batch 1000 re-queued (header 1999 missing) - let sent = pipeline.send_pending(&sender, &storage).await.unwrap(); - assert_eq!(sent, 1); - assert_eq!(pipeline.coordinator.active_count(), 1); - assert_eq!(pipeline.coordinator.pending_count(), 1); + // The single wanted batch is declared to the broker. + let declared = pipeline.send_pending(&network, &storage).await.unwrap(); + assert_eq!(declared, 1); + let starts: Vec = mock + .sent_messages() + .iter() + .filter_map(|m| match m { + NetworkMessage::GetCFilters(gcf) => Some(gcf.start_height), + _ => None, + }) + .collect(); + assert_eq!(starts, vec![0]); - let request = rx.try_recv().unwrap(); - match request { - NetworkRequest::SendMessage(NetworkMessage::GetCFilters(gcf)) => { - assert_eq!(gcf.start_height, 0); - } - other => panic!("Expected GetCFilters, got {:?}", other), + // Receive all filters + for h in 0..=99 { + let hash = Header::dummy(h).block_hash(); + pipeline.receive_with_data(h, hash, &dummy_filter_data(h)); } - assert!(rx.try_recv().is_err(), "should not have sent second request"); - // Store the missing headers and retry - let more_headers = Header::dummy_batch(1000..2000); - storage - .store_headers( - &more_headers.iter().map(crate::types::HashedBlockHeader::from).collect::>(), - ) - .await - .unwrap(); + // Batch complete: nothing wanted, one batch ready to take. + assert!(pipeline.is_idle()); + assert_eq!(pipeline.completed_batches.len(), 1); + assert_eq!(pipeline.filters_received, 100); + assert_eq!(pipeline.highest_received, 99); - let sent = pipeline.send_pending(&sender, &storage).await.unwrap(); - assert_eq!(sent, 1); - assert_eq!(pipeline.coordinator.active_count(), 2); - assert_eq!(pipeline.coordinator.pending_count(), 0); + let completed = pipeline.take_completed_batches(); + assert_eq!(completed.len(), 1); + assert!(pipeline.completed_batches.is_empty()); } } diff --git a/dash-spv/src/sync/filters/sync_manager.rs b/dash-spv/src/sync/filters/sync_manager.rs index ecce171fa..3867bfd72 100644 --- a/dash-spv/src/sync/filters/sync_manager.rs +++ b/dash-spv/src/sync/filters/sync_manager.rs @@ -1,5 +1,5 @@ use crate::error::{SyncError, SyncResult}; -use crate::network::{Message, MessageType, RequestSender}; +use crate::network::{MessageType, NetworkManager, RequestKey}; use crate::storage::{BlockHeaderStorage, FilterHeaderStorage, FilterStorage}; use crate::sync::sync_manager::ensure_not_started; use crate::sync::{ @@ -8,6 +8,8 @@ use crate::sync::{ use async_trait::async_trait; use dashcore::network::message::NetworkMessage; use key_wallet_manager::WalletInterface; +use std::net::SocketAddr; +use std::sync::Arc; #[async_trait] impl< @@ -37,17 +39,19 @@ impl< &[MessageType::CFilter] } - /// Keep `active_batches`, the block-match tracker, pending verified - /// batches, and the filter pipeline's per-batch trackers. Move in-flight - /// `getcfilters` slots back to pending so the next `send_pending` reissues - /// them to the new peer immediately. Without this preservation, a re-scan - /// after reconnect would re-track the same block hashes and leak - /// `pending_blocks` counters that never reach zero. - fn on_disconnect(&mut self) { - self.filter_pipeline.requeue_in_flight(); - } + /// Keep `active_batches`, the block-match tracker, pending verified batches, + /// and the filter pipeline's per-batch trackers across the disconnect. + /// In-flight `getcfilters` slots are re-queued by the network manager itself, + /// and the pipeline's wanted set is preserved so the next `send_pending` + /// reissues them to the new peer. Without this preservation, a re-scan after + /// reconnect would re-track the same block hashes and leak `pending_blocks` + /// counters that never reach zero. + fn on_disconnect(&mut self) {} - async fn start_sync(&mut self, requests: &RequestSender) -> SyncResult> { + async fn start_sync( + &mut self, + network: &Arc, + ) -> SyncResult> { ensure_not_started(self.state(), self.identifier())?; // Resume in-progress work preserved across a disconnect cycle. @@ -56,7 +60,7 @@ impl< // insert a fresh batch at `scan_start` and clobber the existing one, // leaking its `pending_blocks` counter forever. if !self.active_batches.is_empty() { - self.filter_pipeline.send_pending(requests, &*self.header_storage.read().await).await?; + self.filter_pipeline.send_pending(network, &*self.header_storage.read().await).await?; self.set_state(SyncState::Syncing); return Ok(vec![]); } @@ -76,7 +80,7 @@ impl< let mut events = vec![SyncEvent::SyncStart { identifier: self.identifier(), }]; - events.extend(self.start_download(requests).await?); + events.extend(self.start_download(network).await?); return Ok(events); } @@ -86,7 +90,7 @@ impl< // above this must not emit a SyncStart. if stored_filters_tip > 0 && stored_filters_tip == self.progress.committed_height() { self.progress.update_filter_header_tip_height(stored_filters_tip); - return self.start_download(requests).await; + return self.start_download(network).await; } // No stored filters to process - wait for FilterHeadersSyncComplete events @@ -96,10 +100,11 @@ impl< async fn handle_message( &mut self, - msg: Message, - requests: &RequestSender, + peer: SocketAddr, + msg: NetworkMessage, + network: &Arc, ) -> SyncResult> { - let NetworkMessage::CFilter(cfilter) = msg.inner() else { + let NetworkMessage::CFilter(cfilter) = &msg else { return Ok(vec![]); }; @@ -110,7 +115,7 @@ impl< let Some(h) = height else { tracing::warn!( block_hash = %cfilter.block_hash, - peer = %msg.peer_address(), + peer = %peer, "Received CFilter for unknown block hash, rejecting as invalid" ); // TODO: should we penalize the peer a bit? @@ -121,12 +126,21 @@ impl< }; // Buffer filter in pipeline - self.filter_pipeline.receive_with_data(h, cfilter.block_hash, &cfilter.filter); + let batch_completed = + self.filter_pipeline.receive_with_data(h, cfilter.block_hash, &cfilter.filter); + + // A completed batch == one `getcfilters` request fully answered: free that + // peer's in-flight unit (the reader skips per-`cfilter` decrements) and stop + // the network manager tracking the batch for timeout. + if let Some(batch_start) = batch_completed { + network.request_completed(peer, 1).await; + network.request_answered(RequestKey::CFilters(batch_start)).await; + } - // Send more requests if there are free slots - let header_storage = self.header_storage.read().await; - self.filter_pipeline.send_pending(requests, &*header_storage).await?; - drop(header_storage); + // No `send_pending` here: the whole wanted set is already declared to the + // broker, which paces it out as capacity frees. Re-declaring per received + // `cfilter` would re-scan every wanted batch on the hot path. The tick + // re-declares to pick up newly-available batches. Ok(self.store_and_match_batches().await?) } @@ -134,20 +148,20 @@ impl< async fn handle_sync_event( &mut self, event: &SyncEvent, - requests: &RequestSender, + network: &Arc, ) -> SyncResult> { match event { SyncEvent::FilterHeadersSyncComplete { tip_height, } => { - return self.handle_new_filter_headers(*tip_height, requests).await; + return self.handle_new_filter_headers(*tip_height, network).await; } SyncEvent::FilterHeadersStored { tip_height, .. } => { - return self.handle_new_filter_headers(*tip_height, requests).await; + return self.handle_new_filter_headers(*tip_height, network).await; } // React to BlockProcessed events from the BlocksManager @@ -196,7 +210,7 @@ impl< Ok(vec![]) } - async fn tick(&mut self, requests: &RequestSender) -> SyncResult> { + async fn tick(&mut self, network: &Arc) -> SyncResult> { // Detect a wallet that was added behind our scan progress and rescan // from its `synced_height`. Reset committed_height to the lowest // synced_height across the stale wallets only, so already-synced @@ -208,20 +222,32 @@ impl< let behind = wallet_read.wallets_behind(committed); let stale_min_synced = behind.iter().map(|id| wallet_read.wallet_synced_height(id)).min(); + let birth_height = wallet_read.earliest_required_height().await; drop(wallet_read); if let Some(stale_min_synced) = stale_min_synced { - tracing::info!( - "Wallet synced_height {} fell below filter committed_height {}, restarting scan", - stale_min_synced, - committed - ); - self.reset_for_rescan(); - self.progress.update_committed_height(stale_min_synced); - return self.start_download(requests).await; + // Where a restart would actually resume: `start_download` floors the + // scan at the wallets' birth heights and at the stored headers' start. + let scan_floor = birth_height + .max(self.header_storage.read().await.get_start_height().await.unwrap_or(0)); + let restart_at = stale_min_synced.saturating_add(1).max(scan_floor); + // Restart only if that reaches below the committed frontier. Comparing + // the raw `synced_height` instead would fire on every tick for a wallet + // that reports 0 yet needs nothing below the floor, wiping the in-flight + // filter batches before any can complete. + if restart_at <= committed { + tracing::info!( + "Wallet synced_height {} fell below filter committed_height {}, restarting scan at {}", + stale_min_synced, + committed, + restart_at + ); + self.reset_for_rescan(); + self.progress.update_committed_height(stale_min_synced); + return self.start_download(network).await; + } } } - // TODO: Get rid of the send pending in here? Or decouple it from the header storage? // Run tick when Syncing OR when Synced with pending work (new blocks arriving) let has_pending_work = !self.active_batches.is_empty(); let should_tick = match self.state() { @@ -233,12 +259,10 @@ impl< return Ok(vec![]); } - // Handle timeouts - self.filter_pipeline.handle_timeouts(); - - // Send pending requests (decoupled from processing) + // Timeouts/retry are the network manager's job now (the broker re-injects); + // just (re-)declare pending requests (decoupled from processing). let header_storage = self.header_storage.read().await; - self.filter_pipeline.send_pending(requests, &*header_storage).await?; + self.filter_pipeline.send_pending(network, &*header_storage).await?; drop(header_storage); // Store completed batches and do speculative matching diff --git a/dash-spv/src/sync/instantsend/manager.rs b/dash-spv/src/sync/instantsend/manager.rs index eb0ed5509..0d12cc08f 100644 --- a/dash-spv/src/sync/instantsend/manager.rs +++ b/dash-spv/src/sync/instantsend/manager.rs @@ -313,7 +313,7 @@ mod tests { let manager = create_test_manager(); assert_eq!(manager.identifier(), ManagerIdentifier::InstantSend); assert_eq!(manager.state(), SyncState::WaitForEvents); - assert_eq!(manager.wanted_message_types(), vec![MessageType::ISLock, MessageType::Inv]); + assert_eq!(manager.wanted_message_types(), [MessageType::IsDLock, MessageType::Inv]); } /// Buffered `MasternodeStateUpdated` events delivered during @@ -322,9 +322,9 @@ mod tests { /// `MasternodesManager` re-emits the event after reconnect. #[tokio::test] async fn test_handle_sync_event_drops_masternode_state_updated_in_waiting_for_connections() { - use crate::network::RequestSender; + use crate::network::NetworkManager; use crate::sync::SyncEvent; - use tokio::sync::mpsc::unbounded_channel; + use crate::test_utils::MockNetworkManager; let mut manager = create_test_manager(); manager.set_state(SyncState::WaitingForConnections); @@ -333,8 +333,8 @@ mod tests { height: 100, qr_info_result: None, }; - let (tx, _rx) = unbounded_channel(); - let events = manager.handle_sync_event(&event, &RequestSender::new(tx)).await.unwrap(); + let network: Arc = Arc::new(MockNetworkManager::new()); + let events = manager.handle_sync_event(&event, &network).await.unwrap(); assert!(events.is_empty()); assert_eq!(manager.state(), SyncState::WaitingForConnections); diff --git a/dash-spv/src/sync/instantsend/sync_manager.rs b/dash-spv/src/sync/instantsend/sync_manager.rs index e9ee4e96e..9136af604 100644 --- a/dash-spv/src/sync/instantsend/sync_manager.rs +++ b/dash-spv/src/sync/instantsend/sync_manager.rs @@ -1,11 +1,13 @@ use crate::error::SyncResult; -use crate::network::{Message, MessageType, RequestSender}; +use crate::network::{MessageType, NetworkManager}; use crate::sync::{ InstantSendManager, ManagerIdentifier, SyncEvent, SyncManager, SyncManagerProgress, SyncState, }; use async_trait::async_trait; use dashcore::network::message::NetworkMessage; use dashcore::network::message_blockdata::Inventory; +use std::net::SocketAddr; +use std::sync::Arc; #[async_trait] impl SyncManager for InstantSendManager { @@ -22,7 +24,7 @@ impl SyncManager for InstantSendManager { } fn wanted_message_types(&self) -> &'static [MessageType] { - &[MessageType::ISLock, MessageType::Inv] + &[MessageType::IsDLock, MessageType::Inv] } fn on_disconnect(&mut self) { @@ -31,10 +33,11 @@ impl SyncManager for InstantSendManager { async fn handle_message( &mut self, - msg: Message, - requests: &RequestSender, + peer: SocketAddr, + msg: NetworkMessage, + network: &Arc, ) -> SyncResult> { - match msg.inner() { + match &msg { NetworkMessage::ISLock(instantlock) => self.process_instantlock(instantlock).await, NetworkMessage::Inv(inv) => { // Check for InstantSendLock inventory items @@ -49,7 +52,7 @@ impl SyncManager for InstantSendManager { "Received {} InstantSendLock announcements, requesting via getdata", islocks_to_request.len() ); - requests.request_inventory(islocks_to_request, msg.peer_address())?; + network.send_to(peer, NetworkMessage::GetData(islocks_to_request)).await; } Ok(vec![]) } @@ -60,7 +63,7 @@ impl SyncManager for InstantSendManager { async fn handle_sync_event( &mut self, event: &SyncEvent, - _requests: &RequestSender, + _network: &Arc, ) -> SyncResult> { // Drop buffered events that arrive between `stop_sync` and the next // `start_sync`. `pending_instantlocks` is cleared on disconnect, and @@ -100,7 +103,7 @@ impl SyncManager for InstantSendManager { Ok(vec![]) } - async fn tick(&mut self, _requests: &RequestSender) -> SyncResult> { + async fn tick(&mut self, _network: &Arc) -> SyncResult> { // Prune old entries periodically self.prune_old_entries(); Ok(vec![]) diff --git a/dash-spv/src/sync/masternodes/manager.rs b/dash-spv/src/sync/masternodes/manager.rs index b31101586..425c8e6b1 100644 --- a/dash-spv/src/sync/masternodes/manager.rs +++ b/dash-spv/src/sync/masternodes/manager.rs @@ -13,10 +13,11 @@ use tokio::sync::RwLock; use super::pipeline::MnListDiffPipeline; use crate::error::{SyncError, SyncResult}; -use crate::network::RequestSender; +use crate::network::NetworkManager; use crate::storage::BlockHeaderStorage; use crate::sync::{MasternodesProgress, SyncEvent, SyncManager, SyncState}; -use dashcore::network::message_qrinfo::QRInfo; +use dashcore::network::message::NetworkMessage; +use dashcore::network::message_qrinfo::{GetQRInfo, QRInfo}; use dashcore::BlockHash; use std::collections::BTreeSet; @@ -392,7 +393,7 @@ impl MasternodesManager { /// lightweight completion path when the response drains the pipeline. pub(super) async fn send_tip_mnlistdiff_update( &mut self, - requests: &RequestSender, + network: &Arc, ) -> SyncResult> { let new_tip_hash = { let storage = self.header_storage.read().await; @@ -414,7 +415,7 @@ impl MasternodesManager { self.sync_state.pipeline_mode = PipelineMode::Incremental; self.sync_state.mnlistdiff_pipeline.queue_requests(vec![(base_hash, new_tip_hash)]); - self.sync_state.mnlistdiff_pipeline.send_pending(requests)?; + self.sync_state.mnlistdiff_pipeline.send_pending(network).await?; Ok(vec![]) } @@ -439,7 +440,7 @@ impl MasternodesManager { /// would have done had the intermediate events not been dropped. pub(super) async fn complete_pipeline( &mut self, - requests: &RequestSender, + network: &Arc, ) -> SyncResult> { match std::mem::take(&mut self.sync_state.pipeline_mode) { PipelineMode::QuorumValidation { @@ -457,7 +458,7 @@ impl MasternodesManager { ); self.sync_state.qrinfo_retry_count = 0; self.sync_state.clear_pending(); - match self.send_qrinfo_for_tip(requests).await { + match self.send_qrinfo_for_tip(network).await { Ok(extra) => events.extend(extra), Err(e) => tracing::warn!( error = %e, @@ -509,7 +510,7 @@ impl MasternodesManager { /// Called when BlockHeaderSyncComplete is received, ensuring we have all headers. pub(super) async fn send_qrinfo_for_tip( &mut self, - requests: &RequestSender, + network: &Arc, ) -> SyncResult> { // Get info from storage let (tip_height, tip_block_hash) = { @@ -541,11 +542,16 @@ impl MasternodesManager { tip_height, base_hashes.len() ); - // Send before mutating state. If the request errors (e.g. no peers - // connected during a reconnect race), the `?` propagates and we leave - // `WaitingForConnections` intact instead of stranding the manager in - // `Syncing` with `qrinfo_in_flight = None`, which `tick` cannot recover. - requests.request_qr_info(base_hashes, tip_block_hash, true)?; + // Fire the QRInfo. The broker does not track qrinfo (no `RequestKey`), + // so the manager owns its timeout/retry via `qrinfo_in_flight`; there is + // no `request_answered` for it. `send` is fire-and-forget and infallible. + network + .send(NetworkMessage::GetQRInfo(GetQRInfo { + base_block_hashes: base_hashes, + block_request_hash: tip_block_hash, + extra_share: true, + })) + .await; self.progress.add_qr_infos_requested(1); self.sync_state.record_qrinfo_attempt(tip_height); self.sync_state.start_waiting_for_qrinfo(tip_block_hash); @@ -618,15 +624,12 @@ impl std::fmt::Debug for MasternodesManager { #[cfg(test)] mod tests { use super::*; - use crate::network::{MessageType, NetworkRequest}; + use crate::network::MessageType; use crate::storage::{DiskStorageManager, PersistentBlockHeaderStorage, StorageManager}; use crate::sync::sync_manager::SyncManager; use crate::sync::{ManagerIdentifier, SyncManagerProgress}; - use dashcore::block::Header; use dashcore::hashes::Hash; - use dashcore::network::message::NetworkMessage; use dashcore::sml::masternode_list::MasternodeList; - use tokio::sync::mpsc; type TestMasternodesManager = MasternodesManager; @@ -640,54 +643,12 @@ mod tests { create_test_manager_for(dashcore::Network::Testnet).await } - /// Build a regtest manager whose engine has a single list at `tip` and - /// whose block header storage is populated with dummy headers up to - /// `tip`, in `Synced` state with `pipeline_mode = Incremental` and - /// `block_header_tip_height = tip`. Storage must be populated so that - /// `send_qrinfo_for_tip` finds a tip and reaches the network dispatch; - /// otherwise it short-circuits at `storage.get_tip()` and the catch-up - /// path can't be observed at the network layer. Returns the manager, a - /// `RequestSender`, and the matching receiver so the caller binds it - /// (the channel closes when the receiver drops). - async fn make_synced_incremental_manager( - tip: u32, - ) -> (TestMasternodesManager, RequestSender, mpsc::UnboundedReceiver) { - let storage = DiskStorageManager::with_temp_dir().await.unwrap(); - let block_headers = storage.block_headers(); - block_headers - .write() - .await - .store_headers( - &Header::dummy_batch(0..tip + 1) - .iter() - .map(crate::types::HashedBlockHeader::from) - .collect::>(), - ) - .await - .unwrap(); - let engine = engine_with_lists(&[(tip, 1)]); - let mut manager = MasternodesManager::new( - block_headers, - Arc::new(RwLock::new(engine)), - dashcore::Network::Regtest, - ) - .await; - manager.set_state(SyncState::Synced); - manager.sync_state.pipeline_mode = PipelineMode::Incremental; - manager.progress.update_block_header_tip_height(tip); - let (tx, rx) = mpsc::unbounded_channel(); - (manager, RequestSender::new(tx), rx) - } - #[tokio::test] async fn test_masternode_manager_new() { let manager = create_test_manager().await; assert_eq!(manager.identifier(), ManagerIdentifier::Masternode); assert_eq!(manager.state(), SyncState::WaitingForConnections); - assert_eq!( - manager.wanted_message_types(), - vec![MessageType::MnListDiff, MessageType::QRInfo] - ); + assert_eq!(manager.wanted_message_types(), [MessageType::MnListDiff, MessageType::QrInfo]); } #[tokio::test] @@ -927,6 +888,44 @@ mod tests { assert_eq!(manager.progress.current_height(), 0); } + /// Build a `Synced` manager already in the `Incremental` pipeline mode with + /// `tip` block headers stored and a single masternode list at `tip`, plus a + /// [`MockNetworkManager`] to inspect what requests it fires. + async fn make_synced_incremental_manager( + tip: u32, + ) -> (TestMasternodesManager, Arc, Arc) + { + use dashcore::Header; + + let storage = DiskStorageManager::with_temp_dir().await.unwrap(); + let block_headers = storage.block_headers(); + block_headers + .write() + .await + .store_headers( + &Header::dummy_batch(0..tip + 1) + .iter() + .map(crate::types::HashedBlockHeader::from) + .collect::>(), + ) + .await + .unwrap(); + let engine = engine_with_lists(&[(tip, 1)]); + let mut manager = MasternodesManager::new( + block_headers, + Arc::new(RwLock::new(engine)), + dashcore::Network::Regtest, + ) + .await; + manager.set_state(SyncState::Synced); + manager.sync_state.pipeline_mode = PipelineMode::Incremental; + manager.progress.update_block_header_tip_height(tip); + + let mock = Arc::new(crate::test_utils::MockNetworkManager::new()); + let network: Arc = mock.clone(); + (manager, network, mock) + } + /// `complete_pipeline` after `Incremental` re-evaluates the cycle gate at /// the latest tip and fires a catch-up QRInfo when the gate picks /// `QuorumValidation`. When a batch of headers lands while a prior @@ -940,9 +939,9 @@ mod tests { /// `rotation_cycles` from 0 to 1. #[tokio::test] async fn test_complete_incremental_fires_catch_up_when_window_missed() { - let (mut manager, requests, mut rx) = make_synced_incremental_manager(70).await; + let (mut manager, network, mock) = make_synced_incremental_manager(70).await; - manager.complete_pipeline(&requests).await.expect("complete_pipeline succeeds"); + manager.complete_pipeline(&network).await.expect("complete_pipeline succeeds"); assert_eq!( manager.sync_state.current_cycle_height, @@ -963,34 +962,37 @@ mod tests { manager.sync_state.qrinfo_in_flight.is_some(), "the catch-up branch must mark a QRInfo as in flight" ); - let queued = rx.try_recv().expect("a NetworkRequest must be queued by the catch-up"); + let sent = mock.sent_messages(); assert!( - matches!(queued, NetworkRequest::SendMessage(NetworkMessage::GetQRInfo(_))), - "the queued request must be a `GetQRInfo`, got {:?}", - queued + sent.iter().any(|m| matches!(m, NetworkMessage::GetQRInfo(_))), + "the catch-up must send a `GetQRInfo`, got {:?}", + sent ); } - /// `send_qrinfo_for_tip` must not strand the manager in `Syncing` when - /// the network send fails. A buffered `BlockHeaderSyncComplete` consumed - /// during `WaitingForConnections` reaches `send_qrinfo_for_tip` while no - /// peers are connected. If state transitions before the failing send, - /// `tick` cannot recover because it gates on `qrinfo_in_flight.is_some()`. + /// `send_qrinfo_for_tip` fires the QRInfo request and moves the manager out + /// of `WaitingForConnections`. The old `dev` test asserted the send-failure + /// path preserved state, but `NetworkManager::send` is now an infallible + /// fire-and-forget declaration to the broker, so there is no failing-send + /// branch left to exercise. Adapted to assert the normal-send outcome: a + /// `GetQRInfo` is emitted, `qrinfo_in_flight` is set, and the state advances + /// to `Syncing` (the manager is never stranded in `WaitingForConnections`). #[tokio::test] - async fn test_send_qrinfo_for_tip_preserves_state_when_send_fails() { - let (mut manager, requests, rx) = make_synced_incremental_manager(70).await; + async fn test_send_qrinfo_for_tip_fires_and_transitions_from_waiting() { + let (mut manager, network, mock) = make_synced_incremental_manager(70).await; manager.set_state(SyncState::WaitingForConnections); - drop(rx); - let err = manager - .send_qrinfo_for_tip(&requests) - .await - .expect_err("send must fail when the receiver is dropped"); - assert!(matches!(err, SyncError::Network(_)), "expected Network error, got {:?}", err); + manager.send_qrinfo_for_tip(&network).await.expect("send_qrinfo_for_tip succeeds"); - assert_eq!(manager.state(), SyncState::WaitingForConnections); - assert!(manager.sync_state.qrinfo_in_flight.is_none()); - assert_eq!(manager.progress.qr_infos_requested(), 0); + assert_eq!(manager.state(), SyncState::Syncing); + assert!(manager.sync_state.qrinfo_in_flight.is_some()); + assert_eq!(manager.progress.qr_infos_requested(), 1); + let sent = mock.sent_messages(); + assert!( + sent.iter().any(|m| matches!(m, NetworkMessage::GetQRInfo(_))), + "send_qrinfo_for_tip must send a `GetQRInfo`, got {:?}", + sent + ); } /// When the cycle gate picks `Incremental` after an `Incremental` @@ -1000,9 +1002,9 @@ mod tests { /// through to `Incremental` and no QRInfo fires. #[tokio::test] async fn test_complete_incremental_does_not_fire_when_gate_picks_incremental() { - let (mut manager, requests, _rx) = make_synced_incremental_manager(50).await; + let (mut manager, network, mock) = make_synced_incremental_manager(50).await; - manager.complete_pipeline(&requests).await.expect("complete_pipeline succeeds"); + manager.complete_pipeline(&network).await.expect("complete_pipeline succeeds"); assert!( manager.sync_state.qrinfo_in_flight.is_none(), @@ -1013,5 +1015,9 @@ mod tests { 0, "no QRInfo must be requested when the gate picks Incremental" ); + assert!( + mock.sent_messages().is_empty(), + "no request must be sent when the gate picks Incremental" + ); } } diff --git a/dash-spv/src/sync/masternodes/pipeline.rs b/dash-spv/src/sync/masternodes/pipeline.rs index 9a67149ec..c93f9f0f0 100644 --- a/dash-spv/src/sync/masternodes/pipeline.rs +++ b/dash-spv/src/sync/masternodes/pipeline.rs @@ -1,57 +1,34 @@ //! MnListDiff pipeline implementation. //! -//! Handles pipelined download of MnListDiff messages for quorum validation. -//! Uses DownloadCoordinator for request tracking with timeout and retry logic. +//! Declares wanted MnListDiff requests (keyed by target block hash) to the +//! network manager (the broker). The broker owns pacing, de-duplication, +//! timeouts and retries — this pipeline keeps no in-flight queue of its own and +//! simply tracks which `(base, target)` diffs are still wanted. use std::collections::HashMap; -use std::time::Duration; +use std::sync::Arc; use crate::error::SyncResult; -use crate::network::RequestSender; -use crate::sync::download_coordinator::{DownloadConfig, DownloadCoordinator}; -use dashcore::network::message_sml::MnListDiff; +use crate::network::NetworkManager; +use dashcore::network::message::NetworkMessage; +use dashcore::network::message_sml::{GetMnListDiff, MnListDiff}; use dashcore::BlockHash; -/// Maximum concurrent MnListDiff requests. -const MAX_CONCURRENT_MNLISTDIFF: usize = 20; - -/// Timeout for MnListDiff requests. -const MNLISTDIFF_TIMEOUT: Duration = Duration::from_secs(15); - /// Pipeline for downloading MnListDiff messages for quorum validation. /// -/// Uses `DownloadCoordinator` for request tracking (keyed by target block_hash), -/// with a HashMap to store the base hash for each request. -#[derive(Debug)] +/// Holds no request queue of its own: the `base_hashes` map (target -> base) is +/// the "wanted" set. A diff is wanted for exactly as long as it sits in this map; +/// `receive` removes it. The broker de-duplicates, paces, times out and retries. +#[derive(Debug, Default)] pub(super) struct MnListDiffPipeline { - /// Core coordinator tracks requests by target block_hash. - coordinator: DownloadCoordinator, - /// Maps target_hash -> base_hash for each request. + /// Wanted requests: target_hash -> base_hash. Doubles as the "is this diff + /// wanted?" set for validating arrivals. base_hashes: HashMap, } -impl Default for MnListDiffPipeline { - fn default() -> Self { - Self::new() - } -} - impl MnListDiffPipeline { - /// Create a new MnListDiff pipeline. - pub(super) fn new() -> Self { - Self { - coordinator: DownloadCoordinator::new( - DownloadConfig::default() - .with_max_concurrent(MAX_CONCURRENT_MNLISTDIFF) - .with_timeout(MNLISTDIFF_TIMEOUT), - ), - base_hashes: HashMap::new(), - } - } - /// Clear all state. pub(super) fn clear(&mut self) { - self.coordinator.clear(); self.base_hashes.clear(); } @@ -60,7 +37,6 @@ impl MnListDiffPipeline { /// Each request is a (base_hash, target_hash) pair. pub(super) fn queue_requests(&mut self, requests: Vec<(BlockHash, BlockHash)>) { for (base_hash, target_hash) in requests { - self.coordinator.enqueue([target_hash]); self.base_hashes.insert(target_hash, base_hash); } @@ -69,94 +45,67 @@ impl MnListDiffPipeline { } } - /// Send pending requests. + /// Declare every wanted MnListDiff to the network manager. /// - /// Returns the number of requests sent. - pub(super) fn send_pending(&mut self, requests: &RequestSender) -> SyncResult<()> { - let count = self.coordinator.available_to_send(); - if count == 0 { + /// Fired freely (on queue, on tick): the broker de-duplicates, so re-declaring + /// a diff already queued or on the wire is a no-op, and it owns pacing and + /// retry. Re-declaring each tick is the safety net if a peer drops. + pub(super) async fn send_pending( + &mut self, + network: &Arc, + ) -> SyncResult<()> { + if self.base_hashes.is_empty() { return Ok(()); } - let target_hashes = self.coordinator.take_pending(count); - - for target_hash in target_hashes { - let Some(&base_hash) = self.base_hashes.get(&target_hash) else { - tracing::warn!("Missing base hash for target {}, skipping", target_hash); - continue; - }; - - requests.request_mnlist_diff(base_hash, target_hash)?; - self.coordinator.mark_sent(&[target_hash]); - + // Collect first so no borrow of `self` is held across the awaits. + let requests: Vec<(BlockHash, BlockHash)> = + self.base_hashes.iter().map(|(target, base)| (*base, *target)).collect(); + + for (base_block_hash, block_hash) in requests { + network + .send(NetworkMessage::GetMnListD(GetMnListDiff { + base_block_hash, + block_hash, + })) + .await; tracing::trace!( - "Sent GetMnListDiff: base={}, target={} ({} active, {} pending)", - base_hash, - target_hash, - self.coordinator.active_count(), - self.coordinator.pending_count() + "Declared GetMnListDiff: base={}, target={}", + base_block_hash, + block_hash ); } Ok(()) } - /// Check if response matches an in-flight request. + /// Check if response matches a still-wanted request. pub(super) fn match_response(&self, diff: &MnListDiff) -> bool { - self.coordinator.is_in_flight(&diff.block_hash) + self.base_hashes.contains_key(&diff.block_hash) } - /// Receive a MnListDiff response. + /// Receive a MnListDiff response, removing it from the wanted set. /// /// Returns true if the diff was expected, false if unexpected. pub(super) fn receive(&mut self, diff: &MnListDiff) -> bool { let target_hash = diff.block_hash; - if !self.coordinator.receive(&target_hash) { + if self.base_hashes.remove(&target_hash).is_none() { return false; } - self.base_hashes.remove(&target_hash); - tracing::debug!( "Received MnListDiff for {} ({} remaining)", target_hash, - self.coordinator.remaining() + self.base_hashes.len() ); true } - /// Requeue a received MnListDiff for retry. - /// - /// Removes from in-flight tracking and pushes back to the front of the - /// pending queue. - pub(super) fn requeue(&mut self, diff: &MnListDiff) { - let target_hash = diff.block_hash; - - // Remove from in-flight - self.coordinator.receive(&target_hash); - - // Re-enqueue for retry - self.coordinator.enqueue_retry(target_hash); - tracing::debug!("Requeued MnListDiff for {} for retry", diff.block_hash); - } - - /// Handle timeouts, re-queuing timed out requests. - pub(super) fn handle_timeouts(&mut self) { - for target_hash in self.coordinator.check_timeouts() { - self.coordinator.enqueue_retry(target_hash); - } - } - /// Check if pipeline has no pending work. pub(super) fn is_complete(&self) -> bool { - self.coordinator.is_empty() - } - - /// Get the number of in-flight requests. - pub(super) fn active_count(&self) -> usize { - self.coordinator.active_count() + self.base_hashes.is_empty() } } @@ -205,14 +154,13 @@ mod tests { #[test] fn test_pipeline_new() { - let pipeline = MnListDiffPipeline::new(); + let pipeline = MnListDiffPipeline::default(); assert!(pipeline.is_complete()); - assert_eq!(pipeline.active_count(), 0); } #[test] fn test_queue_requests() { - let mut pipeline = MnListDiffPipeline::new(); + let mut pipeline = MnListDiffPipeline::default(); let base1 = BlockHash::from_byte_array([0x01; 32]); let target1 = BlockHash::from_byte_array([0x02; 32]); @@ -222,7 +170,6 @@ mod tests { pipeline.queue_requests(vec![(base1, target1), (base2, target2)]); assert!(!pipeline.is_complete()); - assert_eq!(pipeline.coordinator.pending_count(), 2); assert_eq!(pipeline.base_hashes.len(), 2); assert_eq!(pipeline.base_hashes.get(&target1), Some(&base1)); assert_eq!(pipeline.base_hashes.get(&target2), Some(&base2)); @@ -230,18 +177,14 @@ mod tests { #[test] fn test_match_response() { - let mut pipeline = MnListDiffPipeline::new(); + let mut pipeline = MnListDiffPipeline::default(); let base = BlockHash::from_byte_array([0x01; 32]); let target = BlockHash::from_byte_array([0x02; 32]); pipeline.queue_requests(vec![(base, target)]); - // Take and mark as sent - let items = pipeline.coordinator.take_pending(1); - pipeline.coordinator.mark_sent(&items); - - // Create a test diff + // A queued (wanted) diff matches. let diff = create_test_diff(base, target); assert!(pipeline.match_response(&diff)); @@ -252,17 +195,13 @@ mod tests { #[test] fn test_receive() { - let mut pipeline = MnListDiffPipeline::new(); + let mut pipeline = MnListDiffPipeline::default(); let base = BlockHash::from_byte_array([0x01; 32]); let target = BlockHash::from_byte_array([0x02; 32]); pipeline.queue_requests(vec![(base, target)]); - // Take and mark as sent - let items = pipeline.coordinator.take_pending(1); - pipeline.coordinator.mark_sent(&items); - let diff = create_test_diff(base, target); assert!(pipeline.receive(&diff)); assert!(pipeline.is_complete()); @@ -271,7 +210,7 @@ mod tests { #[test] fn test_receive_unexpected() { - let mut pipeline = MnListDiffPipeline::new(); + let mut pipeline = MnListDiffPipeline::default(); let diff = create_test_diff( BlockHash::from_byte_array([0x01; 32]), @@ -283,86 +222,33 @@ mod tests { } #[test] - fn test_clear() { - let mut pipeline = MnListDiffPipeline::new(); + fn test_receive_duplicate() { + let mut pipeline = MnListDiffPipeline::default(); let base = BlockHash::from_byte_array([0x01; 32]); let target = BlockHash::from_byte_array([0x02; 32]); pipeline.queue_requests(vec![(base, target)]); - pipeline.clear(); + let diff = create_test_diff(base, target); + // First receive removes it from the wanted set. + assert!(pipeline.receive(&diff)); + // Duplicate receive: no longer wanted. + assert!(!pipeline.receive(&diff)); assert!(pipeline.is_complete()); - assert!(pipeline.base_hashes.is_empty()); } #[test] - fn test_handle_timeouts() { - use std::time::Duration; - - let mut pipeline = MnListDiffPipeline { - coordinator: DownloadCoordinator::new( - DownloadConfig::default().with_timeout(Duration::from_millis(1)), - ), - base_hashes: HashMap::new(), - }; - - let base = BlockHash::from_byte_array([0x01; 32]); - let target = BlockHash::from_byte_array([0x02; 32]); - - pipeline.base_hashes.insert(target, base); - pipeline.coordinator.mark_sent(&[target]); - - std::thread::sleep(Duration::from_millis(5)); - - // Timeout re-queues the request, base_hashes preserved - pipeline.handle_timeouts(); - assert_eq!(pipeline.coordinator.pending_count(), 1); - assert!(pipeline.base_hashes.contains_key(&target)); - } - - #[test] - fn test_requeue_puts_back_in_pending() { - let mut pipeline = MnListDiffPipeline::new(); + fn test_clear() { + let mut pipeline = MnListDiffPipeline::default(); let base = BlockHash::from_byte_array([0x01; 32]); let target = BlockHash::from_byte_array([0x02; 32]); pipeline.queue_requests(vec![(base, target)]); + pipeline.clear(); - // Take and mark as sent (simulates sending the request) - let items = pipeline.coordinator.take_pending(1); - pipeline.coordinator.mark_sent(&items); - assert_eq!(pipeline.active_count(), 1); - assert_eq!(pipeline.coordinator.pending_count(), 0); - - let diff = create_test_diff(base, target); - - // Requeue should move from in-flight back to pending - pipeline.requeue(&diff); - assert_eq!(pipeline.active_count(), 0); - assert_eq!(pipeline.coordinator.pending_count(), 1); - // base_hash mapping should be preserved for the retry - assert!(pipeline.base_hashes.contains_key(&target)); - // Pipeline should not be considered complete - assert!(!pipeline.is_complete()); - } - - #[test] - fn test_requeue_always_succeeds() { - let mut pipeline = MnListDiffPipeline::new(); - - let base = BlockHash::from_byte_array([0x01; 32]); - let target = BlockHash::from_byte_array([0x02; 32]); - - pipeline.base_hashes.insert(target, base); - pipeline.coordinator.mark_sent(&[target]); - - let diff = create_test_diff(base, target); - - // Requeue always succeeds - pipeline.requeue(&diff); - assert!(pipeline.base_hashes.contains_key(&target)); - assert_eq!(pipeline.coordinator.pending_count(), 1); + assert!(pipeline.is_complete()); + assert!(pipeline.base_hashes.is_empty()); } } diff --git a/dash-spv/src/sync/masternodes/sync_manager.rs b/dash-spv/src/sync/masternodes/sync_manager.rs index 69c10209a..8e3445bd4 100644 --- a/dash-spv/src/sync/masternodes/sync_manager.rs +++ b/dash-spv/src/sync/masternodes/sync_manager.rs @@ -1,6 +1,6 @@ use super::manager::PipelineMode; use crate::error::SyncResult; -use crate::network::{Message, MessageType, RequestSender}; +use crate::network::{MessageType, NetworkManager, RequestKey}; use crate::storage::BlockHeaderStorage; use crate::sync::{ ManagerIdentifier, MasternodesManager, SyncEvent, SyncManager, SyncManagerProgress, SyncState, @@ -13,6 +13,8 @@ use dashcore::sml::masternode_list_engine::{MasternodeListEngine, WORK_DIFF_DEPT use dashcore::{BlockHash, QuorumHash}; use dashcore_hashes::Hash; use std::collections::{BTreeSet, HashSet}; +use std::net::SocketAddr; +use std::sync::Arc; use std::time::Duration; /// Per-attempt timeout schedule for QRInfo, indexed by the in-flight attempt's @@ -222,7 +224,7 @@ impl SyncManager for MasternodesManager { } fn wanted_message_types(&self) -> &'static [MessageType] { - &[MessageType::MnListDiff, MessageType::QRInfo] + &[MessageType::MnListDiff, MessageType::QrInfo] } fn on_disconnect(&mut self) { @@ -233,10 +235,11 @@ impl SyncManager for MasternodesManager { async fn handle_message( &mut self, - msg: Message, - requests: &RequestSender, + _peer: SocketAddr, + msg: NetworkMessage, + network: &Arc, ) -> SyncResult> { - match msg.inner() { + match &msg { NetworkMessage::QRInfo(qr_info) => { if !self.sync_state.should_process_qrinfo(qr_info) { return Ok(vec![]); @@ -252,7 +255,7 @@ impl SyncManager for MasternodesManager { tracing::info!("Fed {} block heights to engine", fed); // Feed QRInfo to engine first to populate masternode lists - let qr_info_result = match engine.feed_qr_info(qr_info.clone(), true, true) { + let qr_info_result = match engine.feed_qr_info((*qr_info).clone(), true, true) { Ok(qr_info_result) => qr_info_result, Err(e) => { tracing::error!("QRInfo feed into engine failed: {}", e); @@ -315,13 +318,13 @@ impl SyncManager for MasternodesManager { qr_info_result, }; self.sync_state.mnlistdiff_pipeline.queue_requests(request_pairs); - self.sync_state.mnlistdiff_pipeline.send_pending(requests)?; + self.sync_state.mnlistdiff_pipeline.send_pending(network).await?; self.progress.bump_last_activity(); // If no pending requests, complete if !self.sync_state.has_pending_requests() { - return self.complete_pipeline(requests).await; + return self.complete_pipeline(network).await; } } @@ -341,21 +344,22 @@ impl SyncManager for MasternodesManager { Ok(Some(h)) => h, Ok(None) => { tracing::warn!( - "Height not found for MnListDiff block {}, requeuing for retry", + "Height not found for MnListDiff block {}, leaving wanted for retry", diff.block_hash ); - self.sync_state.mnlistdiff_pipeline.requeue(diff); - self.sync_state.mnlistdiff_pipeline.send_pending(requests)?; + // Leave it in the wanted set (do not answer the broker) so + // the broker's timeout/retry re-sends it; re-declare as a + // safety net. + self.sync_state.mnlistdiff_pipeline.send_pending(network).await?; return Ok(vec![]); } Err(e) => { tracing::warn!( - "Failed to get height for MnListDiff block {}: {}, requeuing for retry", + "Failed to get height for MnListDiff block {}: {}, leaving wanted for retry", diff.block_hash, e ); - self.sync_state.mnlistdiff_pipeline.requeue(diff); - self.sync_state.mnlistdiff_pipeline.send_pending(requests)?; + self.sync_state.mnlistdiff_pipeline.send_pending(network).await?; return Ok(vec![]); } }; @@ -366,7 +370,7 @@ impl SyncManager for MasternodesManager { engine.feed_block_height(target_height, diff.block_hash); let apply_ok = - match engine.apply_diff(diff.clone(), Some(target_height), false, None) { + match engine.apply_diff((*diff).clone(), Some(target_height), false, None) { Ok(_) => { self.sync_state.known_mn_list_heights.insert(target_height); tracing::debug!("Applied MnListDiff at height {}", target_height); @@ -385,7 +389,9 @@ impl SyncManager for MasternodesManager { self.progress.add_diffs_processed(1); self.sync_state.mnlistdiff_pipeline.receive(diff); - self.sync_state.mnlistdiff_pipeline.send_pending(requests)?; + // Response correlated: tell the broker to stop tracking this + // request for timeout/retry. + network.request_answered(RequestKey::MnListDiff(diff.block_hash)).await; // Check if all responses received if self.sync_state.mnlistdiff_pipeline.is_complete() { @@ -399,7 +405,7 @@ impl SyncManager for MasternodesManager { return Ok(vec![]); } tracing::info!("All MnListDiff responses received"); - return self.complete_pipeline(requests).await; + return self.complete_pipeline(network).await; } } @@ -412,7 +418,7 @@ impl SyncManager for MasternodesManager { async fn handle_sync_event( &mut self, event: &SyncEvent, - requests: &RequestSender, + network: &Arc, ) -> SyncResult> { // Track block header tip height as headers come in if let SyncEvent::BlockHeadersStored { @@ -462,7 +468,7 @@ impl SyncManager for MasternodesManager { ); self.sync_state.qrinfo_retry_count = 0; self.sync_state.clear_pending(); - return self.send_qrinfo_for_tip(requests).await; + return self.send_qrinfo_for_tip(network).await; } PipelineMode::Incremental => { tracing::debug!( @@ -470,7 +476,7 @@ impl SyncManager for MasternodesManager { tip_height, self.progress.current_height() ); - return self.send_tip_mnlistdiff_update(requests).await; + return self.send_tip_mnlistdiff_update(network).await; } } } @@ -544,10 +550,10 @@ impl SyncManager for MasternodesManager { } self.sync_state.qrinfo_retry_count = 0; self.sync_state.clear_pending(); - return self.send_qrinfo_for_tip(requests).await; + return self.send_qrinfo_for_tip(network).await; } PipelineMode::Incremental => { - return self.send_tip_mnlistdiff_update(requests).await; + return self.send_tip_mnlistdiff_update(network).await; } } } @@ -557,14 +563,14 @@ impl SyncManager for MasternodesManager { ); self.sync_state.qrinfo_retry_count = 0; self.sync_state.clear_pending(); - return self.send_qrinfo_for_tip(requests).await; + return self.send_qrinfo_for_tip(network).await; } } Ok(vec![]) } - async fn tick(&mut self, requests: &RequestSender) -> SyncResult> { + async fn tick(&mut self, network: &Arc) -> SyncResult> { // Handle ticks for both Syncing (initial) and Synced (incremental updates) if !matches!(self.state(), SyncState::Syncing | SyncState::Synced) { return Ok(vec![]); @@ -585,18 +591,19 @@ impl SyncManager for MasternodesManager { if self.sync_state.qrinfo_in_flight.is_none() { self.sync_state.qrinfo_retry_count = 0; self.sync_state.clear_pending(); - return self.send_qrinfo_for_tip(requests).await; + return self.send_qrinfo_for_tip(network).await; } } PipelineMode::Incremental => { - return self.send_tip_mnlistdiff_update(requests).await; + return self.send_tip_mnlistdiff_update(network).await; } } } return Ok(vec![]); } - // Check for QRInfo timeout + // Check for QRInfo timeout. The broker does not track qrinfo, so the + // manager owns its timeout/retry schedule here. if let Some(in_flight) = self.sync_state.qrinfo_in_flight { let timeout = qrinfo_timeout_for(self.sync_state.qrinfo_retry_count); if in_flight.wait_start.elapsed() > timeout { @@ -608,33 +615,19 @@ impl SyncManager for MasternodesManager { ); self.sync_state.qrinfo_retry_count += 1; self.sync_state.clear_pending(); - return self.send_qrinfo_for_tip(requests).await; + return self.send_qrinfo_for_tip(network).await; } else { tracing::warn!( "QRInfo timeout after {} retries, skipping masternode sync", MAX_RETRY_ATTEMPTS ); self.sync_state.clear_pending(); - return self.complete_pipeline(requests).await; + return self.complete_pipeline(network).await; } } return Ok(vec![]); } - // Check for MnListDiff timeouts via pipeline - if self.sync_state.mnlistdiff_pipeline.active_count() > 0 { - self.sync_state.mnlistdiff_pipeline.handle_timeouts(); - - // Send any re-queued requests - self.sync_state.mnlistdiff_pipeline.send_pending(requests)?; - - // Check if complete after handling timeouts - if self.sync_state.mnlistdiff_pipeline.is_complete() { - tracing::info!("MnListDiff pipeline complete"); - return self.complete_pipeline(requests).await; - } - } - Ok(vec![]) } diff --git a/dash-spv/src/sync/mempool/manager.rs b/dash-spv/src/sync/mempool/manager.rs index dd822b0fa..3cc700b59 100644 --- a/dash-spv/src/sync/mempool/manager.rs +++ b/dash-spv/src/sync/mempool/manager.rs @@ -22,7 +22,7 @@ use super::filter::build_wallet_bloom_filter; use super::BLOOM_FALSE_POSITIVE_RATE; use crate::client::config::MempoolStrategy; use crate::error::SyncResult; -use crate::network::RequestSender; +use crate::network::NetworkManager; use crate::sync::mempool::MempoolProgress; use crate::sync::SyncEvent; use crate::types::UnconfirmedTransaction; @@ -111,30 +111,36 @@ impl MempoolManager { pub(super) async fn activate_peer( &mut self, peer: SocketAddr, - requests: &RequestSender, + network: &Arc, ) -> SyncResult<()> { tracing::info!("Activating mempool on peer {} (strategy: {:?})", peer, self.strategy); + // Addressed to THIS peer, not handed to the router: relay is per-peer state on the + // remote node, so a `filterclear`/`filterload` that lands on a different peer than + // the one we mean to activate simply leaves this one mute. match self.strategy { MempoolStrategy::BloomFilter => { - self.load_bloom_filter(peer, requests).await?; + self.load_bloom_filter(peer, network).await?; } MempoolStrategy::FetchAll => { - requests.send_filter_clear(peer)?; + network.send_to(peer, NetworkMessage::FilterClear).await; } } - requests.request_mempool(peer)?; + network.send_to(peer, NetworkMessage::MemPool).await; self.peers.insert(peer, Some(VecDeque::new())); Ok(()) } /// Activate mempool relay on all connected but not-yet-activated peers. - pub(super) async fn activate_all_peers(&mut self, requests: &RequestSender) -> SyncResult<()> { + pub(super) async fn activate_all_peers( + &mut self, + network: &Arc, + ) -> SyncResult<()> { let inactive: Vec = self.peers.iter().filter(|(_, v)| v.is_none()).map(|(k, _)| *k).collect(); for peer in inactive { - self.activate_peer(peer, requests).await?; + self.activate_peer(peer, network).await?; } Ok(()) } @@ -143,7 +149,7 @@ impl MempoolManager { async fn load_bloom_filter( &mut self, peer: SocketAddr, - requests: &RequestSender, + network: &Arc, ) -> SyncResult<()> { let wallet = self.wallet.read().await; let addresses = wallet.monitored_addresses(); @@ -170,13 +176,16 @@ impl MempoolManager { filter_load.filter.len() ); - requests.send_filter_load(filter_load, peer)?; + network.send_to(peer, NetworkMessage::FilterLoad(filter_load)).await; Ok(()) } /// Rebuild the bloom filter on all activated peers. - pub(super) async fn rebuild_filter(&mut self, requests: &RequestSender) -> SyncResult<()> { + pub(super) async fn rebuild_filter( + &mut self, + network: &Arc, + ) -> SyncResult<()> { if self.strategy != MempoolStrategy::BloomFilter { return Ok(()); } @@ -189,9 +198,9 @@ impl MempoolManager { } for peer in activated { - requests.send_filter_clear(peer)?; - self.load_bloom_filter(peer, requests).await?; - requests.request_mempool(peer)?; + network.send_to(peer, NetworkMessage::FilterClear).await; + self.load_bloom_filter(peer, network).await?; + network.send_to(peer, NetworkMessage::MemPool).await; } Ok(()) @@ -205,7 +214,7 @@ impl MempoolManager { &mut self, inv: &[Inventory], peer: SocketAddr, - requests: &RequestSender, + network: &Arc, ) -> SyncResult> { // Acceptance detection must run before any early return: an inv for // one of our own broadcasts from a peer we did NOT send it to proves @@ -244,7 +253,7 @@ impl MempoolManager { if enqueued > 0 { tracing::debug!("Enqueued {} mempool txids for download", enqueued); - self.send_queued(requests).await?; + self.send_queued(network).await?; } Ok(events) @@ -299,7 +308,10 @@ impl MempoolManager { /// /// Deduplicates at send time against `pending_requests` and `mempool_state` /// in case a transaction was received between enqueue and send. - pub(super) async fn send_queued(&mut self, requests: &RequestSender) -> SyncResult<()> { + pub(super) async fn send_queued( + &mut self, + network: &Arc, + ) -> SyncResult<()> { let mut available = MAX_IN_FLIGHT.saturating_sub(self.pending_requests.len()); let has_queued = self.peers.values().any(|v| v.as_ref().is_some_and(|q| !q.is_empty())); if available == 0 || !has_queued { @@ -349,7 +361,10 @@ impl MempoolManager { peer, total_queued, ); - requests.request_inventory(inventory, peer)?; + // Ask the peer that ANNOUNCED these txids, not whichever the router favours: + // a mempool transaction only exists on the nodes that have it, and the queue + // was built per-peer precisely so each one is asked for what it offered. + network.send_to(peer, NetworkMessage::GetData(inventory)).await; } Ok(()) } @@ -364,7 +379,7 @@ impl MempoolManager { &mut self, tx: Transaction, peer: SocketAddr, - requests: &RequestSender, + network: &Arc, ) -> SyncResult> { let txid = tx.txid(); self.pending_requests.remove(&txid); @@ -373,7 +388,7 @@ impl MempoolManager { // Send self-originated transactions to the network regardless of // wallet relevance — the caller explicitly asked to broadcast. if is_local { - self.start_broadcast(&tx, requests); + self.start_broadcast(&tx, network).await; } // Skip if already tracked (e.g., locally broadcast then received from a peer) @@ -422,7 +437,11 @@ impl MempoolManager { /// /// Idempotent per txid: repeated local dispatches of the same transaction /// do not resend (the rebroadcast timer handles resends). - pub(super) fn start_broadcast(&mut self, tx: &Transaction, requests: &RequestSender) { + pub(super) async fn start_broadcast( + &mut self, + tx: &Transaction, + network: &Arc, + ) { let txid = tx.txid(); if self.broadcasts.contains_key(&txid) { return; @@ -442,7 +461,9 @@ impl MempoolManager { if state.holdout.contains(peer) { continue; } - if requests.send_transaction(tx.clone(), *peer).is_ok() { + // Addressed sends, not router sends: the holdout only means something + // if we control exactly which peers received the transaction. + if network.send_to(*peer, NetworkMessage::Tx(tx.clone())).await { state.sent_to.insert(*peer); } } @@ -615,14 +636,14 @@ impl MempoolManager { /// - `Accepted`/`Uncertain`: plain broadcast to all peers (the holdout no /// longer matters; a late echo from a new peer can still upgrade /// `Uncertain` to `Accepted`). - pub(super) async fn rebroadcast_if_due(&mut self, requests: &RequestSender) { - self.rebroadcast_if_due_at(requests, Instant::now()).await + pub(super) async fn rebroadcast_if_due(&mut self, network: &Arc) { + self.rebroadcast_if_due_at(network, Instant::now()).await } /// `now`-injected variant of [`Self::rebroadcast_if_due`]. Tests project `now` /// forward instead of subtracting from `Instant::now()`, which underflows on /// Windows when the QPC-based monotonic clock has a small value at boot. - async fn rebroadcast_if_due_at(&mut self, requests: &RequestSender, now: Instant) { + async fn rebroadcast_if_due_at(&mut self, network: &Arc, now: Instant) { let current_peers: Vec = self.peers.keys().copied().collect(); let mut count: usize = 0; for (txid, state) in &mut self.broadcasts { @@ -658,18 +679,21 @@ impl MempoolManager { // recipients are gone); relay through everyone rather // than letting the transaction stall. if !current_peers.is_empty() { - let _ = requests.broadcast(NetworkMessage::Tx(state.transaction.clone())); + network.broadcast(NetworkMessage::Tx(state.transaction.clone())); state.sent_to.extend(current_peers.iter().copied()); } } else { for peer in targets { - if requests.send_transaction(state.transaction.clone(), peer).is_ok() { + if network + .send_to(peer, NetworkMessage::Tx(state.transaction.clone())) + .await + { state.sent_to.insert(peer); } } } } else { - let _ = requests.broadcast(NetworkMessage::Tx(state.transaction.clone())); + network.broadcast(NetworkMessage::Tx(state.transaction.clone())); } tracing::debug!("Rebroadcast unconfirmed transaction {}", txid); state.last_broadcast = now; @@ -766,20 +790,16 @@ impl fmt::Debug for MempoolManager { .finish() } } - #[cfg(test)] mod tests { use super::*; - use crate::network::NetworkRequest; use dashcore::hashes::Hash; - use dashcore::network::message::NetworkMessage; - use dashcore::{Address, BlockHash, Network, ScriptBuf, Transaction}; + use dashcore::{Address, BlockHash, Network, ScriptBuf}; use key_wallet::transaction_checking::TransactionContext; use key_wallet_manager::test_utils::MockWallet; use crate::sync::SyncState; - use crate::test_utils::test_socket_address; - use tokio::sync::mpsc; + use crate::test_utils::{test_socket_address, MockNetworkManager}; fn dummy_instant_lock(txid: Txid) -> InstantLock { InstantLock { @@ -796,12 +816,8 @@ mod tests { } } - fn create_test_manager( - ) -> (MempoolManager, RequestSender, mpsc::UnboundedReceiver) { + fn create_test_manager() -> MempoolManager { let wallet = Arc::new(RwLock::new(MockWallet::new())); - let (tx, rx) = mpsc::unbounded_channel::(); - let requests = RequestSender::new(tx); - let mut manager = MempoolManager::new( wallet, MempoolStrategy::FetchAll, @@ -810,63 +826,117 @@ mod tests { BroadcastConfig::default(), ); manager.progress.set_state(SyncState::Synced); + manager + } - (manager, requests, rx) + /// Create a manager with BloomFilter strategy where the wallet reports + /// mempool transactions as relevant. BloomFilter strategy skips local + /// address pre-filtering, relying on the wallet for definitive checks. + fn create_relevant_manager() -> (MempoolManager, Arc>) { + let mut mock = MockWallet::new(); + mock.set_mempool_relevant(true); + let wallet = Arc::new(RwLock::new(mock)); + let manager = MempoolManager::new( + wallet.clone(), + MempoolStrategy::BloomFilter, + 1000, + 0, + BroadcastConfig::default(), + ); + (manager, wallet) } - fn create_bloom_manager( - ) -> (MempoolManager, RequestSender, mpsc::UnboundedReceiver) { + /// Create a BloomFilter-strategy manager with an empty (default) wallet. + fn create_bloom_manager() -> MempoolManager { let wallet = Arc::new(RwLock::new(MockWallet::new())); - let (tx, rx) = mpsc::unbounded_channel::(); - let requests = RequestSender::new(tx); + MempoolManager::new( + wallet, + MempoolStrategy::BloomFilter, + 1000, + 0, + BroadcastConfig::default(), + ) + } - let manager = MempoolManager::new( + /// Create a BloomFilter-strategy manager whose wallet monitors `addresses`. + fn create_bloom_manager_with_addresses(addresses: Vec
) -> MempoolManager { + let mut mock = MockWallet::new(); + mock.set_addresses(addresses); + let wallet = Arc::new(RwLock::new(mock)); + MempoolManager::new( wallet, MempoolStrategy::BloomFilter, 1000, 0, BroadcastConfig::default(), - ); + ) + } + + /// Create a test P2PKH address from a byte pattern. + fn test_address(byte: u8) -> Address { + // Build OP_DUP OP_HASH160 <20-byte-hash> OP_EQUALVERIFY OP_CHECKSIG + let mut script_bytes = vec![0x76, 0xa9, 0x14]; // OP_DUP OP_HASH160 PUSH20 + script_bytes.extend_from_slice(&[byte; 20]); + script_bytes.push(0x88); // OP_EQUALVERIFY + script_bytes.push(0xac); // OP_CHECKSIG + let script = ScriptBuf::from(script_bytes); + Address::from_script(&script, Network::Testnet).unwrap() + } - (manager, requests, rx) + /// Build a mock network manager and a trait-object handle to pass to + /// manager methods. Returns `(mock, network)` where `mock` is used for + /// assertions and `network` is passed by reference into the manager. + fn mock_network() -> (Arc, Arc) { + let mock = Arc::new(MockNetworkManager::new()); + let network: Arc = mock.clone(); + (mock, network) } #[tokio::test] async fn test_activation_fetch_all() { let peer = test_socket_address(1); - let (mut manager, requests, mut rx) = create_test_manager(); - manager.activate_peer(peer, &requests).await.unwrap(); - - // FetchAll activation sends filterclear then mempool to the chosen peer - let msg1 = rx.recv().await.unwrap(); - assert!( - matches!(msg1, NetworkRequest::SendMessageToPeer(NetworkMessage::FilterClear, p) if p == peer) - ); - let msg2 = rx.recv().await.unwrap(); - assert!( - matches!(msg2, NetworkRequest::SendMessageToPeer(NetworkMessage::MemPool, p) if p == peer) - ); + let mut manager = create_test_manager(); + let (mock, network) = mock_network(); + manager.activate_peer(peer, &network).await.unwrap(); + + // FetchAll activation sends filterclear then mempool to the chosen peer. + let sent = mock.sent_to_messages(); + assert_eq!(sent.len(), 2); + assert_eq!(sent[0].0, peer); + assert!(matches!(sent[0].1, NetworkMessage::FilterClear)); + assert_eq!(sent[1].0, peer); + assert!(matches!(sent[1].1, NetworkMessage::MemPool)); assert!(matches!(manager.peers.get(&peer), Some(Some(_)))); } #[tokio::test] async fn test_activation_bloom_filter_skips_empty_wallet() { - let (mut manager, requests, mut rx) = create_bloom_manager(); - manager.activate_peer(test_socket_address(1), &requests).await.unwrap(); - - // No addresses in mock wallet, so only MemPool should be sent (no FilterLoad) - let mut found_filter_load = false; - while let Ok(msg) = rx.try_recv() { - if matches!(msg, NetworkRequest::SendMessageToPeer(NetworkMessage::FilterLoad(_), _)) { - found_filter_load = true; - } - } + let mut manager = create_bloom_manager(); + let (mock, network) = mock_network(); + manager.activate_peer(test_socket_address(1), &network).await.unwrap(); + + // No addresses in mock wallet, so only MemPool should be sent (no FilterLoad). + let found_filter_load = + mock.sent_to_messages().iter().any(|(_, m)| matches!(m, NetworkMessage::FilterLoad(_))); assert!(!found_filter_load, "should not send FilterLoad for empty wallet"); } + #[tokio::test] + async fn test_bloom_filter_loaded_with_addresses() { + let addr = test_address(0xab); + let mut manager = create_bloom_manager_with_addresses(vec![addr]); + let (mock, network) = mock_network(); + manager.activate_peer(test_socket_address(1), &network).await.unwrap(); + + let found_filter_load = + mock.sent_to_messages().iter().any(|(_, m)| matches!(m, NetworkMessage::FilterLoad(_))); + assert!(found_filter_load, "expected FilterLoad for wallet with addresses"); + } + #[tokio::test] async fn test_handle_inv_deduplication() { - let (mut manager, requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); + let (_mock, network) = mock_network(); let peer = test_socket_address(1); manager.peers.insert(peer, Some(VecDeque::new())); @@ -874,12 +944,12 @@ mod tests { let inv = vec![Inventory::Transaction(txid)]; // First call should add to pending - let events = manager.handle_inv(&inv, peer, &requests).await.unwrap(); + let events = manager.handle_inv(&inv, peer, &network).await.unwrap(); assert!(events.is_empty()); assert!(manager.pending_requests.contains_key(&txid)); // Second call with same txid should be filtered out - let events = manager.handle_inv(&inv, peer, &requests).await.unwrap(); + let events = manager.handle_inv(&inv, peer, &network).await.unwrap(); assert!(events.is_empty()); assert_eq!(manager.pending_requests.len(), 1); } @@ -887,9 +957,6 @@ mod tests { #[tokio::test] async fn test_handle_inv_capacity_limit() { let wallet = Arc::new(RwLock::new(MockWallet::new())); - let (tx, _rx) = mpsc::unbounded_channel::(); - let requests = RequestSender::new(tx); - let mut manager = MempoolManager::new( wallet, MempoolStrategy::FetchAll, @@ -897,6 +964,7 @@ mod tests { 0, BroadcastConfig::default(), ); + let (_mock, network) = mock_network(); let peer = test_socket_address(1); manager.peers.insert(peer, Some(VecDeque::new())); @@ -919,7 +987,7 @@ mod tests { // New transactions should be filtered out let new_txid = Txid::from_byte_array([99u8; 32]); let inv = vec![Inventory::Transaction(new_txid)]; - let events = manager.handle_inv(&inv, peer, &requests).await.unwrap(); + let events = manager.handle_inv(&inv, peer, &network).await.unwrap(); assert!(events.is_empty()); assert!(!manager.pending_requests.contains_key(&new_txid)); } @@ -927,9 +995,6 @@ mod tests { #[tokio::test] async fn test_handle_inv_pending_requests_limit() { let wallet = Arc::new(RwLock::new(MockWallet::new())); - let (tx, _rx) = mpsc::unbounded_channel::(); - let requests = RequestSender::new(tx); - let mut manager = MempoolManager::new( wallet, MempoolStrategy::FetchAll, @@ -938,76 +1003,27 @@ mod tests { BroadcastConfig::default(), ); manager.progress.set_state(SyncState::Synced); + let (_mock, network) = mock_network(); let peer = test_socket_address(1); manager.peers.insert(peer, Some(VecDeque::new())); // Fill pending requests to capacity let inv1: Vec = (0..2).map(|i| Inventory::Transaction(Txid::from_byte_array([i; 32]))).collect(); - manager.handle_inv(&inv1, peer, &requests).await.unwrap(); + manager.handle_inv(&inv1, peer, &network).await.unwrap(); assert_eq!(manager.pending_requests.len(), 2); // Additional requests should be rejected when pending is at capacity let extra_txid = Txid::from_byte_array([99; 32]); let inv2 = vec![Inventory::Transaction(extra_txid)]; - manager.handle_inv(&inv2, peer, &requests).await.unwrap(); + manager.handle_inv(&inv2, peer, &network).await.unwrap(); assert!(!manager.pending_requests.contains_key(&extra_txid)); } - #[test] - fn test_prune_pending_requests_timeout() { - let wallet = Arc::new(RwLock::new(MockWallet::new())); - let (tx, _rx) = mpsc::unbounded_channel::(); - let _requests = RequestSender::new(tx); - - let mut manager = MempoolManager::new( - wallet, - MempoolStrategy::FetchAll, - 1000, - 0, - BroadcastConfig::default(), - ); - - let fresh_txid = Txid::from_byte_array([1; 32]); - let stale_txid = Txid::from_byte_array([2; 32]); - - manager.pending_requests.insert(fresh_txid, Instant::now()); - manager - .pending_requests - .insert(stale_txid, Instant::now() - PENDING_REQUEST_TIMEOUT - Duration::from_secs(1)); - - manager.prune_pending_requests(); - - assert!(manager.pending_requests.contains_key(&fresh_txid)); - assert!(!manager.pending_requests.contains_key(&stale_txid)); - } - - #[tokio::test] - async fn test_handle_tx_irrelevant() { - let (mut manager, requests, _rx) = create_test_manager(); - - let tx = Transaction { - version: 1, - lock_time: 0, - input: vec![], - output: vec![], - special_transaction_payload: None, - }; - let txid = tx.txid(); - - let events = manager.handle_tx(tx, test_socket_address(1), &requests).await.unwrap(); - // MockWallet returns is_relevant=false by default - assert!(events.is_empty()); - assert_eq!(manager.progress.received(), 1); - - // Irrelevant tx should not be stored - assert!(!manager.transactions.contains_key(&txid)); - assert_eq!(manager.progress.relevant(), 0); - } - #[tokio::test] async fn test_handle_inv_non_transaction_filtered() { - let (mut manager, requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); + let (_mock, network) = mock_network(); let peer = test_socket_address(1); manager.peers.insert(peer, Some(VecDeque::new())); @@ -1016,390 +1032,587 @@ mod tests { Inventory::Transaction(Txid::from_byte_array([1u8; 32])), ]; - let events = manager.handle_inv(&inv, peer, &requests).await.unwrap(); + let events = manager.handle_inv(&inv, peer, &network).await.unwrap(); assert!(events.is_empty()); // Only the transaction should be tracked, not the block assert_eq!(manager.pending_requests.len(), 1); } - #[test] - fn test_prune_expired() { - let (mut manager, _requests, _rx) = create_test_manager(); + #[tokio::test] + async fn test_handle_inv_dedup_against_queue() { + let mut manager = create_test_manager(); + let (_mock, network) = mock_network(); + let peer = test_socket_address(1); + manager.peers.insert(peer, Some(VecDeque::new())); - let fresh_tx = Transaction { - version: 1, - lock_time: 0, - input: vec![], - output: vec![], - special_transaction_payload: None, - }; - let fresh_txid = fresh_tx.txid(); + // Fill pending to capacity so items go to queue + for i in 0..MAX_IN_FLIGHT as u16 { + let mut bytes = [0u8; 32]; + bytes[0..2].copy_from_slice(&i.to_le_bytes()); + manager.pending_requests.insert(Txid::from_byte_array(bytes), Instant::now()); + } - let expired_tx = Transaction { - version: 1, - lock_time: 99, - input: vec![], - output: vec![], - special_transaction_payload: None, - }; - let expired_txid = expired_tx.txid(); - let test_timeout = Duration::from_secs(2); + let txid = Txid::from_byte_array([0xff; 32]); + let inv = vec![Inventory::Transaction(txid)]; - manager.transactions.insert( - fresh_txid, - UnconfirmedTransaction::new(fresh_tx, Amount::from_sat(0), false, false, Vec::new(), 0), + // First call enqueues + manager.handle_inv(&inv, peer, &network).await.unwrap(); + assert_eq!( + manager.peers.values().filter_map(|v| v.as_ref()).map(|q| q.len()).sum::(), + 1 ); - let mut expired_utx = UnconfirmedTransaction::new( - expired_tx, - Amount::from_sat(0), - false, - false, - Vec::new(), - 0, + + // Second call with same txid should be deduped + manager.handle_inv(&inv, peer, &network).await.unwrap(); + assert_eq!( + manager.peers.values().filter_map(|v| v.as_ref()).map(|q| q.len()).sum::(), + 1 ); - expired_utx.first_seen = Instant::now() - test_timeout - Duration::from_secs(1); - manager.transactions.insert(expired_txid, expired_utx); + } - manager.prune_expired(test_timeout); + #[tokio::test] + async fn test_in_flight_limit() { + let mut manager = create_test_manager(); + let (_mock, network) = mock_network(); + let peer = test_socket_address(1); + manager.peers.insert(peer, Some(VecDeque::new())); - assert_eq!(manager.transactions.len(), 1); - assert!(manager.transactions.contains_key(&fresh_txid)); - assert!(!manager.transactions.contains_key(&expired_txid)); - assert_eq!(manager.progress.removed(), 1); + // Send 200 INVs — only MAX_IN_FLIGHT should go to pending, rest queued + let inv: Vec = (0..200u16) + .map(|i| { + let mut bytes = [0u8; 32]; + bytes[0..2].copy_from_slice(&i.to_le_bytes()); + Inventory::Transaction(Txid::from_byte_array(bytes)) + }) + .collect(); + + manager.handle_inv(&inv, peer, &network).await.unwrap(); + assert_eq!(manager.pending_requests.len(), MAX_IN_FLIGHT); + assert_eq!( + manager.peers.values().filter_map(|v| v.as_ref()).map(|q| q.len()).sum::(), + 100 + ); } - /// Create a manager with BloomFilter strategy where the wallet reports - /// mempool transactions as relevant. BloomFilter strategy skips local - /// address pre-filtering, relying on the wallet for definitive checks. - fn create_relevant_manager( - ) -> (MempoolManager, RequestSender, Arc>) { - let mut mock = MockWallet::new(); - mock.set_mempool_relevant(true); - let wallet = Arc::new(RwLock::new(mock)); - let (tx, _rx) = mpsc::unbounded_channel::(); - let requests = RequestSender::new(tx); + #[tokio::test] + async fn test_send_queued_drains_after_response() { + let mut manager = create_test_manager(); + let (_mock, network) = mock_network(); + let peer = test_socket_address(1); + manager.peers.insert(peer, Some(VecDeque::new())); - let manager = MempoolManager::new( - wallet.clone(), - MempoolStrategy::BloomFilter, - 1000, - 0, - BroadcastConfig::default(), + // Fill with 150 INVs + let inv: Vec = (0..150u16) + .map(|i| { + let mut bytes = [0u8; 32]; + bytes[0..2].copy_from_slice(&i.to_le_bytes()); + Inventory::Transaction(Txid::from_byte_array(bytes)) + }) + .collect(); + + manager.handle_inv(&inv, peer, &network).await.unwrap(); + assert_eq!(manager.pending_requests.len(), MAX_IN_FLIGHT); + assert_eq!( + manager.peers.values().filter_map(|v| v.as_ref()).map(|q| q.len()).sum::(), + 50 ); - (manager, requests, wallet) + // Simulate receiving 10 responses (freeing 10 slots) + let pending_txids: Vec = manager.pending_requests.keys().take(10).copied().collect(); + for txid in &pending_txids { + manager.pending_requests.remove(txid); + } + assert_eq!(manager.pending_requests.len(), 90); + + // send_queued should fill the freed slots + manager.send_queued(&network).await.unwrap(); + assert_eq!(manager.pending_requests.len(), MAX_IN_FLIGHT); + assert_eq!( + manager.peers.values().filter_map(|v| v.as_ref()).map(|q| q.len()).sum::(), + 40 + ); } #[tokio::test] - async fn test_handle_tx_relevant_stores_transaction() { - let (mut manager, requests, _wallet) = create_relevant_manager(); + async fn test_send_queued_skips_already_received() { + let mut manager = create_test_manager(); + let (_mock, network) = mock_network(); + let peer = test_socket_address(1); + // Create a real transaction and get its actual txid let tx = Transaction { version: 1, - lock_time: 0, + lock_time: 0xaa, input: vec![], output: vec![], special_transaction_payload: None, }; let txid = tx.txid(); - let events = manager.handle_tx(tx, test_socket_address(1), &requests).await.unwrap(); - assert!(events.is_empty()); + // Enqueue the txid on an activated peer + manager.peers.insert(peer, Some(VecDeque::from([txid]))); - // Verify transaction was stored - assert!(manager.transactions.contains_key(&txid)); - assert_eq!(manager.progress.received(), 1); - assert_eq!(manager.progress.relevant(), 1); - assert_eq!(manager.progress.tracked(), 1); + // Simulate the transaction arriving before send + manager.transactions.insert( + txid, + UnconfirmedTransaction::new(tx, Amount::from_sat(0), false, false, Vec::new(), 0), + ); - // Processing the same transaction again should be a no-op (dedup guard) - let tx2 = Transaction { - version: 1, - lock_time: 0, - input: vec![], - output: vec![], - special_transaction_payload: None, - }; - let events = manager.handle_tx(tx2, test_socket_address(1), &requests).await.unwrap(); - assert!(events.is_empty()); + manager.send_queued(&network).await.unwrap(); + // Txid should have been skipped, not added to pending + assert!(manager.pending_requests.is_empty()); + assert!(manager.peers.values().filter_map(|v| v.as_ref()).all(|q| q.is_empty())); + } - assert_eq!(manager.transactions.len(), 1); - // Progress counters should not have incremented - assert_eq!(manager.progress.received(), 1); - assert_eq!(manager.progress.relevant(), 1); + #[tokio::test] + async fn test_send_queued_noop_at_capacity() { + let mut manager = create_test_manager(); + let (_mock, network) = mock_network(); + + // Fill pending to MAX_IN_FLIGHT + for i in 0..MAX_IN_FLIGHT as u16 { + let mut bytes = [0u8; 32]; + bytes[0..2].copy_from_slice(&i.to_le_bytes()); + manager.pending_requests.insert(Txid::from_byte_array(bytes), Instant::now()); + } + + // Add something to the queue on an activated peer + manager.peers.insert( + test_socket_address(1), + Some(VecDeque::from([Txid::from_byte_array([0xff; 32])])), + ); + + manager.send_queued(&network).await.unwrap(); + // Queue should remain unchanged (one peer with one txid) + assert_eq!( + manager.peers.values().filter_map(|v| v.as_ref()).map(|q| q.len()).sum::(), + 1 + ); + assert_eq!(manager.pending_requests.len(), MAX_IN_FLIGHT); } #[tokio::test] - async fn test_handle_tx_local_records_send() { - let (mut manager, requests, _wallet) = create_relevant_manager(); + async fn test_seen_txids_deduplication_window() { + let mut manager = create_test_manager(); + let (_mock, network) = mock_network(); + let peer = test_socket_address(1); + manager.peers.insert(peer, Some(VecDeque::new())); - let tx = Transaction { - version: 2, - lock_time: 0, - input: vec![], - output: vec![], - special_transaction_payload: None, - }; - let txid = tx.txid(); + let txid = Txid::from_byte_array([1u8; 32]); + let inv = vec![Inventory::Transaction(txid)]; - // Use the unspecified address to simulate a locally broadcast transaction - let local_addr = SocketAddr::from(([0, 0, 0, 0], 0)); - manager.handle_tx(tx, local_addr, &requests).await.unwrap(); + // A fresh seen_txids entry should cause handle_inv to skip the txid + manager.seen_txids.insert(txid, Instant::now()); + manager.handle_inv(&inv, peer, &network).await.unwrap(); + assert!(manager.pending_requests.is_empty(), "seen txid should be skipped"); - assert!(manager.transactions.contains_key(&txid)); + // An expired entry should allow the txid to be accepted again + manager.seen_txids.insert(txid, Instant::now() - SEEN_TXID_EXPIRY - Duration::from_secs(1)); + manager.handle_inv(&inv, peer, &network).await.unwrap(); assert!( - manager.broadcasts.contains_key(&txid), - "locally dispatched transaction should be tracked as a broadcast" + manager.pending_requests.contains_key(&txid), + "expired seen txid should be accepted" ); } #[tokio::test] - async fn test_handle_tx_remote_does_not_record_send() { - let (mut manager, requests, _wallet) = create_relevant_manager(); + async fn test_rebuild_filter_clears_and_reloads() { + let addr = test_address(0xab); + let mut manager = create_bloom_manager_with_addresses(vec![addr]); + let (mock, network) = mock_network(); + let peer = test_socket_address(1); - let tx = Transaction { - version: 3, + manager.activate_peer(peer, &network).await.unwrap(); + + // Drain activation messages + mock.clear_sent(); + + manager.rebuild_filter(&network).await.unwrap(); + + // Verify message sequence: FilterClear, FilterLoad, MemPool + let sent = mock.sent_to_messages(); + assert_eq!(sent.len(), 3); + assert!(matches!(sent[0].1, NetworkMessage::FilterClear)); + assert!(matches!(sent[1].1, NetworkMessage::FilterLoad(_))); + assert!(matches!(sent[2].1, NetworkMessage::MemPool)); + } + + #[tokio::test] + async fn test_rebuild_filter_no_activated_peers_noop() { + let mut manager = create_bloom_manager(); + let (mock, network) = mock_network(); + // No activation, so no activated peers + assert!(manager.peers.values().all(|v| v.is_none())); + + manager.rebuild_filter(&network).await.unwrap(); + assert!(mock.sent_to_messages().is_empty()); + } + + fn test_transaction(version: u16) -> Transaction { + Transaction { + version, lock_time: 0, input: vec![], output: vec![], special_transaction_payload: None, - }; + } + } + + #[tokio::test] + async fn test_rebroadcast_sends_old_pending_broadcasts() { + let mut manager = create_test_manager(); + let (mock, network) = mock_network(); + let peer = test_socket_address(1); + manager.peers.insert(peer, Some(VecDeque::new())); + + let tx = test_transaction(10); let txid = tx.txid(); - manager.handle_tx(tx, test_socket_address(1), &requests).await.unwrap(); + let t0 = Instant::now(); + let later = t0 + REBROADCAST_INTERVAL + Duration::from_secs(1); - assert!(manager.transactions.contains_key(&txid)); + let mut state = TxBroadcastState::new(tx, t0); + state.sent_to.insert(peer); + manager.broadcasts.insert(txid, state); + + manager.rebroadcast_if_due_at(&network, later).await; + + // Pending entries are resent via targeted sends (respecting the holdout) + let sends = mock.sent_to_messages(); + assert_eq!(sends.len(), 1, "expected a rebroadcast message"); + assert_eq!(sends[0].0, peer); assert!( - !manager.broadcasts.contains_key(&txid), - "peer-received transaction should not be tracked as a broadcast" + matches!(sends[0].1, NetworkMessage::Tx(_)), + "expected targeted Tx, got {:?}", + sends[0].1 + ); + + // Timestamp should be reset to `later`, so a second call at the same instant + // must not rebroadcast. + mock.clear_sent(); + manager.rebroadcast_if_due_at(&network, later).await; + assert!( + mock.sent_to_messages().is_empty(), + "should not rebroadcast immediately after reset" ); } #[tokio::test] - async fn test_handle_tx_clears_pending_request() { - let (mut manager, requests, _wallet) = create_relevant_manager(); + async fn test_rebroadcast_uncertain_uses_full_broadcast() { + let mut manager = create_test_manager(); + let (mock, network) = mock_network(); + let peer = test_socket_address(1); + manager.peers.insert(peer, Some(VecDeque::new())); - let tx = Transaction { - version: 1, - lock_time: 0, - input: vec![], - output: vec![], - special_transaction_payload: None, - }; + let tx = test_transaction(12); let txid = tx.txid(); - // Simulate that we requested this transaction - manager.pending_requests.insert(txid, Instant::now()); - assert!(manager.pending_requests.contains_key(&txid)); + let t0 = Instant::now(); + let later = t0 + REBROADCAST_INTERVAL + Duration::from_secs(1); - manager.handle_tx(tx, test_socket_address(1), &requests).await.unwrap(); - // Pending request should be cleared regardless of relevance - assert!(!manager.pending_requests.contains_key(&txid)); + let mut state = TxBroadcastState::new(tx, t0); + state.sent_to.insert(peer); + state.status = BroadcastStatus::Uncertain; + manager.broadcasts.insert(txid, state); - // Since the manager uses BloomFilter strategy (relevant mock), tx should be stored - assert!(manager.transactions.contains_key(&txid)); + manager.rebroadcast_if_due_at(&network, later).await; + + let broadcasts = mock.broadcast_messages(); + assert_eq!(broadcasts.len(), 1, "expected a rebroadcast message"); + assert!( + matches!(broadcasts[0], NetworkMessage::Tx(_)), + "expected broadcast Tx, got {:?}", + broadcasts[0] + ); } - fn create_bloom_manager_with_addresses( - addresses: Vec
, - ) -> (MempoolManager, RequestSender, mpsc::UnboundedReceiver) { - let mut mock = MockWallet::new(); - mock.set_addresses(addresses); - let wallet = Arc::new(RwLock::new(mock)); - let (tx, rx) = mpsc::unbounded_channel::(); - let requests = RequestSender::new(tx); + #[tokio::test] + async fn test_rebroadcast_skips_recent_transactions() { + let mut manager = create_test_manager(); + let (mock, network) = mock_network(); + let peer = test_socket_address(1); + manager.peers.insert(peer, Some(VecDeque::new())); - let manager = MempoolManager::new( - wallet, - MempoolStrategy::BloomFilter, - 1000, - 0, - BroadcastConfig::default(), - ); + let tx = test_transaction(11); + let txid = tx.txid(); - (manager, requests, rx) - } + // Add a broadcast that was just sent (within the rebroadcast interval) + let mut state = TxBroadcastState::new(tx, Instant::now()); + state.sent_to.insert(peer); + manager.broadcasts.insert(txid, state); - /// Create a test P2PKH address from a byte pattern. - fn test_address(byte: u8) -> Address { - // Build OP_DUP OP_HASH160 <20-byte-hash> OP_EQUALVERIFY OP_CHECKSIG - let mut script_bytes = vec![0x76, 0xa9, 0x14]; // OP_DUP OP_HASH160 PUSH20 - script_bytes.extend_from_slice(&[byte; 20]); - script_bytes.push(0x88); // OP_EQUALVERIFY - script_bytes.push(0xac); // OP_CHECKSIG - let script = ScriptBuf::from(script_bytes); - Address::from_script(&script, Network::Testnet).unwrap() + manager.rebroadcast_if_due(&network).await; + + assert!( + mock.sent_to_messages().is_empty() && mock.broadcast_messages().is_empty(), + "recently sent transactions should not be rebroadcast" + ); } - #[tokio::test] - async fn test_bloom_filter_loaded_with_addresses() { - let addr = test_address(0xab); + #[test] + fn test_prune_pending_requests_timeout() { + let mut manager = create_test_manager(); - let (mut manager, requests, mut rx) = create_bloom_manager_with_addresses(vec![addr]); - manager.activate_peer(test_socket_address(1), &requests).await.unwrap(); + let fresh_txid = Txid::from_byte_array([1; 32]); + let stale_txid = Txid::from_byte_array([2; 32]); - let mut found_filter_load = false; - while let Ok(msg) = rx.try_recv() { - if matches!(msg, NetworkRequest::SendMessageToPeer(NetworkMessage::FilterLoad(_), _)) { - found_filter_load = true; - } - } - assert!(found_filter_load, "expected FilterLoad for wallet with addresses"); + manager.pending_requests.insert(fresh_txid, Instant::now()); + manager + .pending_requests + .insert(stale_txid, Instant::now() - PENDING_REQUEST_TIMEOUT - Duration::from_secs(1)); + + manager.prune_pending_requests(); + + assert!(manager.pending_requests.contains_key(&fresh_txid)); + assert!(!manager.pending_requests.contains_key(&stale_txid)); } #[tokio::test] - async fn test_mark_instant_send_emits_status_change() { - let (mut manager, _requests, _rx) = create_test_manager(); + async fn test_handle_tx_irrelevant() { + let mut manager = create_test_manager(); + let (_mock, network) = mock_network(); let tx = Transaction { version: 1, - lock_time: 42, + lock_time: 0, input: vec![], output: vec![], special_transaction_payload: None, }; let txid = tx.txid(); - manager.broadcasts.insert(txid, TxBroadcastState::new(tx.clone(), Instant::now())); - manager.transactions.insert( - txid, - UnconfirmedTransaction::new(tx, Amount::from_sat(0), false, false, Vec::new(), 0), - ); - let events = manager.process_instant_send(dummy_instant_lock(txid)).await; + let events = manager.handle_tx(tx, test_socket_address(1), &network).await.unwrap(); + // MockWallet returns is_relevant=false by default + assert!(events.is_empty()); + assert_eq!(manager.progress.received(), 1); - // Verify IS flag, broadcast promotion, and tracking cleanup - assert!(manager.transactions.get(&txid).unwrap().is_instant_send); - assert!( - !manager.broadcasts.contains_key(&txid), - "IS-locked transaction should no longer be tracked as a broadcast" + // Irrelevant tx should not be stored + assert!(!manager.transactions.contains_key(&txid)); + assert_eq!(manager.progress.relevant(), 0); + } + + #[test] + fn test_prune_expired() { + let mut manager = create_test_manager(); + + let fresh_tx = Transaction { + version: 1, + lock_time: 0, + input: vec![], + output: vec![], + special_transaction_payload: None, + }; + let fresh_txid = fresh_tx.txid(); + + let expired_tx = Transaction { + version: 1, + lock_time: 99, + input: vec![], + output: vec![], + special_transaction_payload: None, + }; + let expired_txid = expired_tx.txid(); + let test_timeout = Duration::from_secs(2); + + manager.transactions.insert( + fresh_txid, + UnconfirmedTransaction::new(fresh_tx, Amount::from_sat(0), false, false, Vec::new(), 0), ); - assert!(matches!( - events.as_slice(), - [SyncEvent::TransactionBroadcastResult { - result: BroadcastResult::Accepted { .. }, - .. - }] - )); + let mut expired_utx = UnconfirmedTransaction::new( + expired_tx, + Amount::from_sat(0), + false, + false, + Vec::new(), + 0, + ); + expired_utx.first_seen = Instant::now() - test_timeout - Duration::from_secs(1); + manager.transactions.insert(expired_txid, expired_utx); - let wallet = manager.wallet.read().await; - let status_changes = wallet.status_changes(); - let changes = status_changes.lock().await; - assert_eq!(changes.len(), 1); - assert_eq!(changes[0].0, txid); - assert!(matches!(changes[0].1, TransactionContext::InstantSend(_))); + manager.prune_expired(test_timeout); + + assert_eq!(manager.transactions.len(), 1); + assert!(manager.transactions.contains_key(&fresh_txid)); + assert!(!manager.transactions.contains_key(&expired_txid)); + assert_eq!(manager.progress.removed(), 1); } #[tokio::test] - async fn test_mark_instant_send_stores_pending_for_unknown() { - let (mut manager, _requests, _rx) = create_test_manager(); + async fn test_handle_tx_relevant_stores_transaction() { + let (mut manager, _wallet) = create_relevant_manager(); + let (_mock, network) = mock_network(); - let unknown_txid = Txid::from_byte_array([0xbb; 32]); - manager.process_instant_send(dummy_instant_lock(unknown_txid)).await; + let tx = Transaction { + version: 1, + lock_time: 0, + input: vec![], + output: vec![], + special_transaction_payload: None, + }; + let txid = tx.txid(); - // No immediate wallet notification - let wallet = manager.wallet.read().await; - let status_changes = wallet.status_changes(); - let changes = status_changes.lock().await; - assert!(changes.is_empty()); + let events = manager.handle_tx(tx, test_socket_address(1), &network).await.unwrap(); + assert!(events.is_empty()); - // But the txid is remembered for when the transaction arrives - assert!(manager.pending_is_locks.contains_key(&unknown_txid)); + // Verify transaction was stored + assert!(manager.transactions.contains_key(&txid)); + assert_eq!(manager.progress.received(), 1); + assert_eq!(manager.progress.relevant(), 1); + assert_eq!(manager.progress.tracked(), 1); + + // Processing the same transaction again should be a no-op (dedup guard) + let tx2 = Transaction { + version: 1, + lock_time: 0, + input: vec![], + output: vec![], + special_transaction_payload: None, + }; + let events = manager.handle_tx(tx2, test_socket_address(1), &network).await.unwrap(); + assert!(events.is_empty()); + + assert_eq!(manager.transactions.len(), 1); + // Progress counters should not have incremented + assert_eq!(manager.progress.received(), 1); + assert_eq!(manager.progress.relevant(), 1); } #[tokio::test] - async fn test_in_flight_limit() { - let (mut manager, requests, _rx) = create_test_manager(); - let peer = test_socket_address(1); - manager.peers.insert(peer, Some(VecDeque::new())); + async fn test_handle_tx_local_records_send() { + let (mut manager, _wallet) = create_relevant_manager(); + let (_mock, network) = mock_network(); - // Send 200 INVs — only MAX_IN_FLIGHT should go to pending, rest queued - let inv: Vec = (0..200u16) - .map(|i| { - let mut bytes = [0u8; 32]; - bytes[0..2].copy_from_slice(&i.to_le_bytes()); - Inventory::Transaction(Txid::from_byte_array(bytes)) - }) - .collect(); + let tx = Transaction { + version: 2, + lock_time: 0, + input: vec![], + output: vec![], + special_transaction_payload: None, + }; + let txid = tx.txid(); - manager.handle_inv(&inv, peer, &requests).await.unwrap(); - assert_eq!(manager.pending_requests.len(), MAX_IN_FLIGHT); - assert_eq!( - manager.peers.values().filter_map(|v| v.as_ref()).map(|q| q.len()).sum::(), - 100 + // Use the unspecified address to simulate a locally broadcast transaction + let local_addr = SocketAddr::from(([0, 0, 0, 0], 0)); + manager.handle_tx(tx, local_addr, &network).await.unwrap(); + + assert!(manager.transactions.contains_key(&txid)); + assert!( + manager.broadcasts.contains_key(&txid), + "locally dispatched transaction should be tracked as a broadcast" ); } #[tokio::test] - async fn test_send_queued_drains_after_response() { - let (mut manager, requests, _rx) = create_test_manager(); - let peer = test_socket_address(1); - manager.peers.insert(peer, Some(VecDeque::new())); - - // Fill with 150 INVs - let inv: Vec = (0..150u16) - .map(|i| { - let mut bytes = [0u8; 32]; - bytes[0..2].copy_from_slice(&i.to_le_bytes()); - Inventory::Transaction(Txid::from_byte_array(bytes)) - }) - .collect(); + async fn test_handle_tx_remote_does_not_record_send() { + let (mut manager, _wallet) = create_relevant_manager(); + let (_mock, network) = mock_network(); - manager.handle_inv(&inv, peer, &requests).await.unwrap(); - assert_eq!(manager.pending_requests.len(), MAX_IN_FLIGHT); - assert_eq!( - manager.peers.values().filter_map(|v| v.as_ref()).map(|q| q.len()).sum::(), - 50 - ); + let tx = Transaction { + version: 3, + lock_time: 0, + input: vec![], + output: vec![], + special_transaction_payload: None, + }; + let txid = tx.txid(); - // Simulate receiving 10 responses (freeing 10 slots) - let pending_txids: Vec = manager.pending_requests.keys().take(10).copied().collect(); - for txid in &pending_txids { - manager.pending_requests.remove(txid); - } - assert_eq!(manager.pending_requests.len(), 90); + manager.handle_tx(tx, test_socket_address(1), &network).await.unwrap(); - // send_queued should fill the freed slots - manager.send_queued(&requests).await.unwrap(); - assert_eq!(manager.pending_requests.len(), MAX_IN_FLIGHT); - assert_eq!( - manager.peers.values().filter_map(|v| v.as_ref()).map(|q| q.len()).sum::(), - 40 + assert!(manager.transactions.contains_key(&txid)); + assert!( + !manager.broadcasts.contains_key(&txid), + "peer-received transaction should not be tracked as a broadcast" ); } #[tokio::test] - async fn test_send_queued_skips_already_received() { - let (mut manager, requests, _rx) = create_test_manager(); - let peer = test_socket_address(1); + async fn test_handle_tx_clears_pending_request() { + let (mut manager, _wallet) = create_relevant_manager(); + let (_mock, network) = mock_network(); - // Create a real transaction and get its actual txid let tx = Transaction { version: 1, - lock_time: 0xaa, + lock_time: 0, input: vec![], output: vec![], special_transaction_payload: None, }; let txid = tx.txid(); - // Enqueue the txid on an activated peer - manager.peers.insert(peer, Some(VecDeque::from([txid]))); + // Simulate that we requested this transaction + manager.pending_requests.insert(txid, Instant::now()); + assert!(manager.pending_requests.contains_key(&txid)); - // Simulate the transaction arriving before send + manager.handle_tx(tx, test_socket_address(1), &network).await.unwrap(); + // Pending request should be cleared regardless of relevance + assert!(!manager.pending_requests.contains_key(&txid)); + + // Since the manager uses BloomFilter strategy (relevant mock), tx should be stored + assert!(manager.transactions.contains_key(&txid)); + } + + #[tokio::test] + async fn test_mark_instant_send_emits_status_change() { + let mut manager = create_test_manager(); + + let tx = Transaction { + version: 1, + lock_time: 42, + input: vec![], + output: vec![], + special_transaction_payload: None, + }; + let txid = tx.txid(); + manager.broadcasts.insert(txid, TxBroadcastState::new(tx.clone(), Instant::now())); manager.transactions.insert( txid, UnconfirmedTransaction::new(tx, Amount::from_sat(0), false, false, Vec::new(), 0), ); - manager.send_queued(&requests).await.unwrap(); - // Txid should have been skipped, not added to pending - assert!(manager.pending_requests.is_empty()); - assert!(manager.peers.values().filter_map(|v| v.as_ref()).all(|q| q.is_empty())); + let events = manager.process_instant_send(dummy_instant_lock(txid)).await; + + // Verify IS flag, broadcast promotion, and tracking cleanup + assert!(manager.transactions.get(&txid).unwrap().is_instant_send); + assert!( + !manager.broadcasts.contains_key(&txid), + "IS-locked transaction should no longer be tracked as a broadcast" + ); + assert!(matches!( + events.as_slice(), + [SyncEvent::TransactionBroadcastResult { + result: BroadcastResult::Accepted { .. }, + .. + }] + )); + + let wallet = manager.wallet.read().await; + let status_changes = wallet.status_changes(); + let changes = status_changes.lock().await; + assert_eq!(changes.len(), 1); + assert_eq!(changes[0].0, txid); + assert!(matches!(changes[0].1, TransactionContext::InstantSend(_))); + } + + #[tokio::test] + async fn test_mark_instant_send_stores_pending_for_unknown() { + let mut manager = create_test_manager(); + + let unknown_txid = Txid::from_byte_array([0xbb; 32]); + manager.process_instant_send(dummy_instant_lock(unknown_txid)).await; + + // No immediate wallet notification + let wallet = manager.wallet.read().await; + let status_changes = wallet.status_changes(); + let changes = status_changes.lock().await; + assert!(changes.is_empty()); + + // But the txid is remembered for when the transaction arrives + assert!(manager.pending_is_locks.contains_key(&unknown_txid)); } #[test] fn test_clear_pending_clears_queue() { - let (mut manager, _requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); manager.pending_requests.insert(Txid::from_byte_array([1; 32]), Instant::now()); manager @@ -1415,35 +1628,10 @@ mod tests { assert!(manager.pending_is_locks.is_empty()); } - #[tokio::test] - async fn test_send_queued_noop_at_capacity() { - let (mut manager, requests, _rx) = create_test_manager(); - - // Fill pending to MAX_IN_FLIGHT - for i in 0..MAX_IN_FLIGHT as u16 { - let mut bytes = [0u8; 32]; - bytes[0..2].copy_from_slice(&i.to_le_bytes()); - manager.pending_requests.insert(Txid::from_byte_array(bytes), Instant::now()); - } - - // Add something to the queue on an activated peer - manager.peers.insert( - test_socket_address(1), - Some(VecDeque::from([Txid::from_byte_array([0xff; 32])])), - ); - - manager.send_queued(&requests).await.unwrap(); - // Queue should remain unchanged (one peer with one txid) - assert_eq!( - manager.peers.values().filter_map(|v| v.as_ref()).map(|q| q.len()).sum::(), - 1 - ); - assert_eq!(manager.pending_requests.len(), MAX_IN_FLIGHT); - } - #[tokio::test] async fn test_instant_send_before_transaction() { - let (mut manager, requests, wallet) = create_relevant_manager(); + let (mut manager, wallet) = create_relevant_manager(); + let (_mock, network) = mock_network(); let tx = Transaction { version: 1, @@ -1459,7 +1647,7 @@ mod tests { assert!(manager.pending_is_locks.contains_key(&txid)); // Transaction arrives - manager.handle_tx(tx, test_socket_address(1), &requests).await.unwrap(); + manager.handle_tx(tx, test_socket_address(1), &network).await.unwrap(); // Pending IS lock consumed assert!(manager.pending_is_locks.is_empty()); @@ -1481,7 +1669,8 @@ mod tests { #[tokio::test] async fn test_instant_send_before_irrelevant_transaction() { - let (mut manager, requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); + let (_mock, network) = mock_network(); let tx = Transaction { version: 1, @@ -1497,7 +1686,7 @@ mod tests { assert!(manager.pending_is_locks.contains_key(&txid)); // Transaction arrives but wallet says it's not relevant - manager.handle_tx(tx, test_socket_address(1), &requests).await.unwrap(); + manager.handle_tx(tx, test_socket_address(1), &network).await.unwrap(); // Pending IS lock cleaned up (no leak) assert!(manager.pending_is_locks.is_empty()); @@ -1508,7 +1697,7 @@ mod tests { #[tokio::test] async fn test_pending_is_locks_capacity_limit() { - let (mut manager, _requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); // Fill pending IS locks to capacity for i in 0..MAX_PENDING_IS_LOCKS { @@ -1528,7 +1717,7 @@ mod tests { #[test] fn test_prune_expired_removes_is_lock_for_expired_tx() { - let (mut manager, _requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); let tx = Transaction { version: 1, @@ -1570,7 +1759,7 @@ mod tests { #[test] fn test_prune_expired_removes_stale_pending_is_locks() { - let (mut manager, _requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); let test_timeout = Duration::from_secs(2); @@ -1602,65 +1791,10 @@ mod tests { ); } - #[tokio::test] - async fn test_handle_inv_dedup_against_queue() { - let (mut manager, requests, _rx) = create_test_manager(); - let peer = test_socket_address(1); - manager.peers.insert(peer, Some(VecDeque::new())); - - // Fill pending to capacity so items go to queue - for i in 0..MAX_IN_FLIGHT as u16 { - let mut bytes = [0u8; 32]; - bytes[0..2].copy_from_slice(&i.to_le_bytes()); - manager.pending_requests.insert(Txid::from_byte_array(bytes), Instant::now()); - } - - let txid = Txid::from_byte_array([0xff; 32]); - let inv = vec![Inventory::Transaction(txid)]; - - // First call enqueues - manager.handle_inv(&inv, peer, &requests).await.unwrap(); - assert_eq!( - manager.peers.values().filter_map(|v| v.as_ref()).map(|q| q.len()).sum::(), - 1 - ); - - // Second call with same txid should be deduped - manager.handle_inv(&inv, peer, &requests).await.unwrap(); - assert_eq!( - manager.peers.values().filter_map(|v| v.as_ref()).map(|q| q.len()).sum::(), - 1 - ); - } - - #[tokio::test] - async fn test_bloom_filter_load_failure_propagates() { - let addr = test_address(0xab); - let mut mock = MockWallet::new(); - mock.set_addresses(vec![addr]); - let wallet = Arc::new(RwLock::new(mock)); - let (tx, rx) = mpsc::unbounded_channel::(); - let requests = RequestSender::new(tx); - - let mut manager = MempoolManager::new( - wallet, - MempoolStrategy::BloomFilter, - 1000, - 0, - BroadcastConfig::default(), - ); - - // Drop receiver so send_filter_load fails - drop(rx); - - // activate() should propagate the error - let result = manager.activate_peer(test_socket_address(1), &requests).await; - assert!(result.is_err()); - } - #[tokio::test] async fn test_handle_tx_relevant_populates_wallet_effect_fields() { - let (mut manager, requests, wallet) = create_relevant_manager(); + let (mut manager, wallet) = create_relevant_manager(); + let (_mock, network) = mock_network(); let tx = Transaction { version: 1, @@ -1679,7 +1813,7 @@ mod tests { w.set_mempool_addresses(vec![addr.clone()]); } - manager.handle_tx(tx, test_socket_address(1), &requests).await.unwrap(); + manager.handle_tx(tx, test_socket_address(1), &network).await.unwrap(); let stored = manager.transactions.get(&txid).unwrap(); assert_eq!(stored.net_amount, 50000); @@ -1691,7 +1825,8 @@ mod tests { #[tokio::test] async fn test_handle_tx_outgoing_transaction() { - let (mut manager, requests, wallet) = create_relevant_manager(); + let (mut manager, wallet) = create_relevant_manager(); + let (_mock, network) = mock_network(); let tx = Transaction { version: 1, @@ -1707,7 +1842,7 @@ mod tests { w.set_mempool_net_amount(-30000); } - manager.handle_tx(tx, test_socket_address(1), &requests).await.unwrap(); + manager.handle_tx(tx, test_socket_address(1), &network).await.unwrap(); let stored = manager.transactions.get(&txid).unwrap(); assert_eq!(stored.net_amount, -30000); @@ -1718,7 +1853,7 @@ mod tests { #[test] fn test_peer_connected_creates_entry() { - let (mut manager, _requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); let peer = test_socket_address(1); assert!(!manager.peers.contains_key(&peer)); @@ -1729,7 +1864,7 @@ mod tests { #[test] fn test_peer_disconnected_redistributes_queue() { - let (mut manager, _requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); let peer1 = test_socket_address(1); let peer2 = test_socket_address(2); @@ -1750,7 +1885,7 @@ mod tests { #[test] fn test_peer_disconnected_no_peers_drops_queue() { - let (mut manager, _requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); let peer = test_socket_address(1); manager.peers.insert(peer, Some(VecDeque::from([Txid::from_byte_array([1; 32])]))); @@ -1762,7 +1897,7 @@ mod tests { #[test] fn test_prune_pending_requeues_to_activated_peer() { - let (mut manager, _requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); let peer = test_socket_address(1); manager.peers.insert(peer, Some(VecDeque::new())); @@ -1779,7 +1914,7 @@ mod tests { #[test] fn test_prune_pending_drops_when_no_peers() { - let (mut manager, _requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); let txid = Txid::from_byte_array([1; 32]); manager @@ -1794,7 +1929,7 @@ mod tests { #[test] fn test_remove_confirmed_removes_txids() { - let (mut manager, _requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); let mut txids = Vec::new(); for i in 0..3u32 { @@ -1843,7 +1978,7 @@ mod tests { #[test] fn test_remove_confirmed_unknown_txids_noop() { - let (mut manager, _requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); let unknown = vec![Txid::from_byte_array([0xaa; 32]), Txid::from_byte_array([0xbb; 32])]; @@ -1853,155 +1988,9 @@ mod tests { assert_eq!(manager.progress.removed(), 0); } - #[tokio::test] - async fn test_rebuild_filter_clears_and_reloads() { - let addr = test_address(0xab); - let (mut manager, requests, mut rx) = create_bloom_manager_with_addresses(vec![addr]); - let peer = test_socket_address(1); - - manager.activate_peer(peer, &requests).await.unwrap(); - - // Drain activation messages - while rx.try_recv().is_ok() {} - - manager.rebuild_filter(&requests).await.unwrap(); - - // Verify message sequence: FilterClear, FilterLoad, MemPool - let msg1 = rx.try_recv().unwrap(); - assert!(matches!(msg1, NetworkRequest::SendMessageToPeer(NetworkMessage::FilterClear, _))); - let msg2 = rx.try_recv().unwrap(); - assert!(matches!( - msg2, - NetworkRequest::SendMessageToPeer(NetworkMessage::FilterLoad(_), _) - )); - let msg3 = rx.try_recv().unwrap(); - assert!(matches!(msg3, NetworkRequest::SendMessageToPeer(NetworkMessage::MemPool, _))); - } - - #[tokio::test] - async fn test_rebuild_filter_no_activated_peers_noop() { - let (mut manager, requests, mut rx) = create_bloom_manager(); - // No activation, so no activated peers - assert!(manager.peers.values().all(|v| v.is_none())); - - manager.rebuild_filter(&requests).await.unwrap(); - assert!(rx.try_recv().is_err()); - } - - #[tokio::test] - async fn test_seen_txids_deduplication_window() { - let (mut manager, requests, _rx) = create_test_manager(); - let peer = test_socket_address(1); - manager.peers.insert(peer, Some(VecDeque::new())); - - let txid = Txid::from_byte_array([1u8; 32]); - let inv = vec![Inventory::Transaction(txid)]; - - // A fresh seen_txids entry should cause handle_inv to skip the txid - manager.seen_txids.insert(txid, Instant::now()); - manager.handle_inv(&inv, peer, &requests).await.unwrap(); - assert!(manager.pending_requests.is_empty(), "seen txid should be skipped"); - - // An expired entry should allow the txid to be accepted again - manager.seen_txids.insert(txid, Instant::now() - SEEN_TXID_EXPIRY - Duration::from_secs(1)); - manager.handle_inv(&inv, peer, &requests).await.unwrap(); - assert!( - manager.pending_requests.contains_key(&txid), - "expired seen txid should be accepted" - ); - } - - fn test_transaction(version: u16) -> Transaction { - Transaction { - version, - lock_time: 0, - input: vec![], - output: vec![], - special_transaction_payload: None, - } - } - - #[tokio::test] - async fn test_rebroadcast_sends_old_pending_broadcasts() { - let (mut manager, requests, mut rx) = create_test_manager(); - let peer = test_socket_address(1); - manager.peers.insert(peer, Some(VecDeque::new())); - - let tx = test_transaction(10); - let txid = tx.txid(); - - let t0 = Instant::now(); - let later = t0 + REBROADCAST_INTERVAL + Duration::from_secs(1); - - let mut state = TxBroadcastState::new(tx, t0); - state.sent_to.insert(peer); - manager.broadcasts.insert(txid, state); - - manager.rebroadcast_if_due_at(&requests, later).await; - - // Pending entries are resent via targeted sends (respecting the holdout) - let msg = rx.try_recv().expect("expected a rebroadcast message"); - assert!( - matches!(msg, NetworkRequest::SendMessageToPeer(NetworkMessage::Tx(_), p) if p == peer), - "expected SendMessageToPeer(Tx), got {:?}", - msg - ); - - // Timestamp should be reset to `later`, so a second call at the same instant - // must not rebroadcast. - manager.rebroadcast_if_due_at(&requests, later).await; - assert!(rx.try_recv().is_err(), "should not rebroadcast immediately after reset"); - } - - #[tokio::test] - async fn test_rebroadcast_uncertain_uses_full_broadcast() { - let (mut manager, requests, mut rx) = create_test_manager(); - let peer = test_socket_address(1); - manager.peers.insert(peer, Some(VecDeque::new())); - - let tx = test_transaction(12); - let txid = tx.txid(); - - let t0 = Instant::now(); - let later = t0 + REBROADCAST_INTERVAL + Duration::from_secs(1); - - let mut state = TxBroadcastState::new(tx, t0); - state.sent_to.insert(peer); - state.status = BroadcastStatus::Uncertain; - manager.broadcasts.insert(txid, state); - - manager.rebroadcast_if_due_at(&requests, later).await; - - let msg = rx.try_recv().expect("expected a rebroadcast message"); - assert!( - matches!(msg, NetworkRequest::BroadcastMessage(NetworkMessage::Tx(_))), - "expected BroadcastMessage(Tx), got {:?}", - msg - ); - } - - #[tokio::test] - async fn test_rebroadcast_skips_recent_transactions() { - let (mut manager, requests, mut rx) = create_test_manager(); - let peer = test_socket_address(1); - manager.peers.insert(peer, Some(VecDeque::new())); - - let tx = test_transaction(11); - let txid = tx.txid(); - - // Add a broadcast that was just sent (within the rebroadcast interval) - let mut state = TxBroadcastState::new(tx, Instant::now()); - state.sent_to.insert(peer); - manager.broadcasts.insert(txid, state); - - manager.rebroadcast_if_due(&requests).await; - - assert!(rx.try_recv().is_err(), "recently sent transactions should not be rebroadcast"); - } - #[test] fn test_peer_disconnect_keeps_other_peers_intact() { - let (mut manager, _requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); let peer1 = test_socket_address(1); let peer2 = test_socket_address(2); @@ -2022,14 +2011,12 @@ mod tests { const LOCAL_SENTINEL: SocketAddr = SocketAddr::new(std::net::IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED), 0); - fn drain_tx_sends(rx: &mut mpsc::UnboundedReceiver) -> Vec { - let mut sends = Vec::new(); - while let Ok(msg) = rx.try_recv() { - if let NetworkRequest::SendMessageToPeer(NetworkMessage::Tx(_), peer) = msg { - sends.push(peer); - } - } - sends + /// Peers that received a `tx` via a targeted send, in order. + fn tx_send_targets(mock: &MockNetworkManager) -> Vec { + mock.sent_to_messages() + .into_iter() + .filter_map(|(peer, msg)| matches!(msg, NetworkMessage::Tx(_)).then_some(peer)) + .collect() } fn accepted_event_count(events: &[SyncEvent]) -> usize { @@ -2049,17 +2036,17 @@ mod tests { #[tokio::test] async fn test_local_tx_sends_to_half_of_peers() { - let (mut manager, requests, mut rx) = create_test_manager(); + let mut manager = create_test_manager(); + let (mock, network) = mock_network(); for i in 1..=4 { manager.peers.insert(test_socket_address(i), None); } let tx = test_transaction(20); let txid = tx.txid(); - manager.handle_tx(tx, LOCAL_SENTINEL, &requests).await.unwrap(); + manager.handle_tx(tx, LOCAL_SENTINEL, &network).await.unwrap(); - let sends = drain_tx_sends(&mut rx); - assert_eq!(sends.len(), 2, "should send to half of 4 peers"); + assert_eq!(tx_send_targets(&mock).len(), 2, "should send to half of 4 peers"); let state = manager.broadcasts.get(&txid).expect("broadcast tracked"); assert_eq!(state.sent_to.len(), 2); @@ -2068,29 +2055,32 @@ mod tests { assert_eq!(state.status, BroadcastStatus::Pending); // A second local dispatch of the same tx must not resend + mock.clear_sent(); let tx = test_transaction(20); - manager.handle_tx(tx, LOCAL_SENTINEL, &requests).await.unwrap(); - assert!(drain_tx_sends(&mut rx).is_empty(), "idempotent per txid"); + manager.handle_tx(tx, LOCAL_SENTINEL, &network).await.unwrap(); + assert!(tx_send_targets(&mock).is_empty(), "idempotent per txid"); } #[tokio::test] async fn test_local_tx_single_peer_no_holdout() { - let (mut manager, requests, mut rx) = create_test_manager(); + let mut manager = create_test_manager(); + let (mock, network) = mock_network(); let peer = test_socket_address(1); manager.peers.insert(peer, None); let tx = test_transaction(21); let txid = tx.txid(); - manager.handle_tx(tx, LOCAL_SENTINEL, &requests).await.unwrap(); + manager.handle_tx(tx, LOCAL_SENTINEL, &network).await.unwrap(); - assert_eq!(drain_tx_sends(&mut rx), vec![peer]); + assert_eq!(tx_send_targets(&mock), vec![peer]); let state = manager.broadcasts.get(&txid).unwrap(); assert!(state.holdout.is_empty(), "single peer leaves nobody to hold out"); } #[tokio::test] async fn test_echo_from_holdout_peer_accepts_once() { - let (mut manager, requests, mut rx) = create_test_manager(); + let mut manager = create_test_manager(); + let (_mock, network) = mock_network(); let recipient = test_socket_address(1); let holdout = test_socket_address(2); manager.peers.insert(recipient, None); @@ -2102,30 +2092,28 @@ mod tests { state.sent_to.insert(recipient); state.holdout.insert(holdout); manager.broadcasts.insert(txid, state); - drain_tx_sends(&mut rx); let inv = vec![Inventory::Transaction(txid)]; // Echo from the recipient peer carries no information - let events = manager.handle_inv(&inv, recipient, &requests).await.unwrap(); + let events = manager.handle_inv(&inv, recipient, &network).await.unwrap(); assert_eq!(accepted_event_count(&events), 0); assert_eq!(manager.broadcasts[&txid].status, BroadcastStatus::Pending); // Echo from the holdout peer proves propagation - let events = manager.handle_inv(&inv, holdout, &requests).await.unwrap(); + let events = manager.handle_inv(&inv, holdout, &network).await.unwrap(); assert_eq!(accepted_event_count(&events), 1); assert_eq!(manager.broadcasts[&txid].status, BroadcastStatus::Accepted); // A repeat announcement must not emit a second event - let events = manager.handle_inv(&inv, holdout, &requests).await.unwrap(); + let events = manager.handle_inv(&inv, holdout, &network).await.unwrap(); assert_eq!(accepted_event_count(&events), 0); } #[tokio::test] async fn test_echo_detected_when_mempool_full() { let wallet = Arc::new(RwLock::new(MockWallet::new())); - let (tx_chan, _rx) = mpsc::unbounded_channel::(); - let requests = RequestSender::new(tx_chan); + let (_mock, network) = mock_network(); let mut manager = MempoolManager::new( wallet, MempoolStrategy::FetchAll, @@ -2151,13 +2139,14 @@ mod tests { // The mempool-full early return must not swallow the acceptance echo let inv = vec![Inventory::Transaction(txid)]; - let events = manager.handle_inv(&inv, holdout, &requests).await.unwrap(); + let events = manager.handle_inv(&inv, holdout, &network).await.unwrap(); assert_eq!(accepted_event_count(&events), 1); } #[tokio::test] async fn test_timeout_uncertain_then_late_echo_upgrades() { - let (mut manager, requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); + let (_mock, network) = mock_network(); let holdout = test_socket_address(2); let tx = test_transaction(28); @@ -2187,33 +2176,35 @@ mod tests { // A late echo still upgrades the outcome to accepted let inv = vec![Inventory::Transaction(txid)]; - let events = manager.handle_inv(&inv, holdout, &requests).await.unwrap(); + let events = manager.handle_inv(&inv, holdout, &network).await.unwrap(); assert_eq!(accepted_event_count(&events), 1); assert_eq!(manager.broadcasts[&txid].status, BroadcastStatus::Accepted); } #[tokio::test] async fn test_zero_peers_at_broadcast_sends_on_next_tick() { - let (mut manager, requests, mut rx) = create_test_manager(); + let mut manager = create_test_manager(); + let (mock, network) = mock_network(); // No peers connected at broadcast time let tx = test_transaction(29); let txid = tx.txid(); - manager.handle_tx(tx, LOCAL_SENTINEL, &requests).await.unwrap(); - assert!(drain_tx_sends(&mut rx).is_empty()); + manager.handle_tx(tx, LOCAL_SENTINEL, &network).await.unwrap(); + assert!(tx_send_targets(&mock).is_empty()); assert!(manager.broadcasts[&txid].sent_to.is_empty()); // A peer connects; the never-sent broadcast is due immediately let peer = test_socket_address(1); manager.peers.insert(peer, None); - manager.rebroadcast_if_due(&requests).await; - assert_eq!(drain_tx_sends(&mut rx), vec![peer]); + manager.rebroadcast_if_due(&network).await; + assert_eq!(tx_send_targets(&mock), vec![peer]); assert!(manager.broadcasts[&txid].sent_to.contains(&peer)); } #[tokio::test] async fn test_holdout_sticky_across_rebroadcasts() { - let (mut manager, requests, mut rx) = create_test_manager(); + let mut manager = create_test_manager(); + let (mock, network) = mock_network(); let recipient = test_socket_address(1); let holdout = test_socket_address(2); manager.peers.insert(recipient, None); @@ -2228,16 +2219,17 @@ mod tests { manager.broadcasts.insert(txid, state); let later = t0 + REBROADCAST_INTERVAL + Duration::from_secs(1); - manager.rebroadcast_if_due_at(&requests, later).await; + manager.rebroadcast_if_due_at(&network, later).await; // Only the original recipient is resent to; the holdout stays withheld - assert_eq!(drain_tx_sends(&mut rx), vec![recipient]); + assert_eq!(tx_send_targets(&mock), vec![recipient]); assert_eq!(manager.broadcasts[&txid].holdout, [holdout].into_iter().collect()); } #[tokio::test] async fn test_holdout_repicked_when_all_holdouts_disconnect() { - let (mut manager, requests, mut rx) = create_test_manager(); + let mut manager = create_test_manager(); + let (mock, network) = mock_network(); let recipient = test_socket_address(1); let new_peer = test_socket_address(3); manager.peers.insert(recipient, None); @@ -2253,11 +2245,11 @@ mod tests { manager.broadcasts.insert(txid, state); let later = t0 + REBROADCAST_INTERVAL + Duration::from_secs(1); - manager.rebroadcast_if_due_at(&requests, later).await; + manager.rebroadcast_if_due_at(&network, later).await; // The never-sent connected peer becomes the replacement holdout, // so the resend goes only to the original recipient. - assert_eq!(drain_tx_sends(&mut rx), vec![recipient]); + assert_eq!(tx_send_targets(&mock), vec![recipient]); assert!(manager.broadcasts[&txid].holdout.contains(&new_peer)); assert!(!manager.broadcasts[&txid].sent_to.contains(&new_peer)); } @@ -2265,8 +2257,7 @@ mod tests { #[tokio::test] async fn test_acceptance_threshold_two_requires_two_peers() { let wallet = Arc::new(RwLock::new(MockWallet::new())); - let (tx_chan, _rx) = mpsc::unbounded_channel::(); - let requests = RequestSender::new(tx_chan); + let (_mock, network) = mock_network(); let mut manager = MempoolManager::new( wallet, MempoolStrategy::FetchAll, @@ -2288,10 +2279,10 @@ mod tests { manager.broadcasts.insert(txid, state); let inv = vec![Inventory::Transaction(txid)]; - let events = manager.handle_inv(&inv, holdout1, &requests).await.unwrap(); + let events = manager.handle_inv(&inv, holdout1, &network).await.unwrap(); assert_eq!(accepted_event_count(&events), 0, "one echo below threshold"); - let events = manager.handle_inv(&inv, holdout2, &requests).await.unwrap(); + let events = manager.handle_inv(&inv, holdout2, &network).await.unwrap(); assert_eq!(accepted_event_count(&events), 1, "second distinct peer meets threshold"); assert!(matches!( events.as_slice(), @@ -2306,7 +2297,7 @@ mod tests { #[tokio::test] async fn test_clear_pending_preserves_broadcasts() { - let (mut manager, _requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); let tx = test_transaction(33); let txid = tx.txid(); @@ -2322,7 +2313,7 @@ mod tests { #[test] fn test_prune_expired_removes_old_broadcasts() { - let (mut manager, _requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); let timeout = Duration::from_secs(2); let fresh = test_transaction(34); diff --git a/dash-spv/src/sync/mempool/sync_manager.rs b/dash-spv/src/sync/mempool/sync_manager.rs index 1e0e15569..2d1fb1be2 100644 --- a/dash-spv/src/sync/mempool/sync_manager.rs +++ b/dash-spv/src/sync/mempool/sync_manager.rs @@ -1,12 +1,14 @@ use super::manager::MEMPOOL_TX_EXPIRY; use crate::error::SyncResult; -use crate::network::{Message, MessageType, NetworkEvent, RequestSender}; +use crate::network::{MessageType, NetworkEvent, NetworkManager}; use crate::sync::{ ManagerIdentifier, MempoolManager, SyncEvent, SyncManager, SyncManagerProgress, SyncState, }; use async_trait::async_trait; use dashcore::network::message::NetworkMessage; use key_wallet_manager::WalletInterface; +use std::net::SocketAddr; +use std::sync::Arc; #[async_trait] impl SyncManager for MempoolManager { @@ -26,9 +28,12 @@ impl SyncManager for MempoolManager { &[MessageType::Inv, MessageType::Tx] } - async fn start_sync(&mut self, requests: &RequestSender) -> SyncResult> { + async fn start_sync( + &mut self, + network: &Arc, + ) -> SyncResult> { // After a full disconnect, re-activate mempool on all connected peers - self.activate_all_peers(requests).await?; + self.activate_all_peers(network).await?; let has_activated = self.peers.values().any(|v| v.is_some()); if has_activated { self.set_state(SyncState::Synced); @@ -39,20 +44,55 @@ impl SyncManager for MempoolManager { Ok(vec![]) } + /// Track the peer set as the network reports it. + /// + /// This is the only thing that seeds `self.peers`, and it works only because the + /// network manager connects in `start` — after the coordinator has subscribed us. + /// Everything else here (relay activation, and so every transaction and InstantSend + /// lock we ever see) hangs off it. + async fn handle_network_event( + &mut self, + event: &NetworkEvent, + network: &Arc, + ) -> SyncResult> { + match event { + NetworkEvent::PeerConnected(addr) => { + self.handle_peer_connected(*addr); + // If synced, activate the new peer immediately; otherwise + // `FiltersSyncComplete` (or `start_sync`) will. + if self.state() == SyncState::Synced + && self.peers.get(addr).is_some_and(|v| v.is_none()) + { + tracing::info!("Activating mempool on newly connected peer {}", addr); + self.activate_peer(*addr, network).await?; + } + Ok(vec![]) + } + NetworkEvent::PeerDisconnected(addr) => { + // Hands this peer's queued txids to another activated one rather than + // dropping them. + self.handle_peer_disconnected(*addr); + Ok(vec![]) + } + _ => { + crate::sync::sync_manager::default_handle_network_event(self, event, network).await + } + } + } + fn on_disconnect(&mut self) { self.clear_pending(); } async fn handle_message( &mut self, - msg: Message, - requests: &RequestSender, + peer: SocketAddr, + msg: NetworkMessage, + network: &Arc, ) -> SyncResult> { - match msg.inner() { - NetworkMessage::Inv(inv) => self.handle_inv(inv, msg.peer_address(), requests).await, - NetworkMessage::Tx(tx) => { - self.handle_tx(tx.clone(), msg.peer_address(), requests).await - } + match msg { + NetworkMessage::Tx(tx) => self.handle_tx(tx, peer, network).await, + NetworkMessage::Inv(inv) => self.handle_inv(&inv, peer, network).await, _ => Ok(vec![]), } } @@ -60,7 +100,7 @@ impl SyncManager for MempoolManager { async fn handle_sync_event( &mut self, event: &SyncEvent, - requests: &RequestSender, + network: &Arc, ) -> SyncResult> { match event { // Activate as soon as filter sync completes — the wallet's address @@ -69,7 +109,7 @@ impl SyncManager for MempoolManager { .. } => { if self.state() != SyncState::Synced { - self.activate_all_peers(requests).await?; + self.activate_all_peers(network).await?; let has_activated = self.peers.values().any(|v| v.is_some()); if has_activated { self.set_state(SyncState::Synced); @@ -102,11 +142,11 @@ impl SyncManager for MempoolManager { } } - async fn tick(&mut self, requests: &RequestSender) -> SyncResult> { + async fn tick(&mut self, network: &Arc) -> SyncResult> { // Broadcast bookkeeping runs regardless of sync state: broadcasts can // be initiated (and time out) before the mempool phase is synced. let events = self.expire_broadcasts(); - self.rebroadcast_if_due(requests).await; + self.rebroadcast_if_due(network).await; if self.state() != SyncState::Synced { return Ok(events); @@ -119,7 +159,7 @@ impl SyncManager for MempoolManager { self.prune_pending_requests(); // Send queued getdata requests now that slots may have freed up - self.send_queued(requests).await?; + self.send_queued(network).await?; // Rebuild bloom filter if the wallet's monitored set has changed. // @@ -134,54 +174,13 @@ impl SyncManager for MempoolManager { let current_revision = self.wallet.read().await.monitor_revision(); if current_revision != self.last_monitor_revision { tracing::info!("Wallet monitor revision changed, rebuilding bloom filter"); - self.rebuild_filter(requests).await?; + self.rebuild_filter(network).await?; self.last_monitor_revision = current_revision; } Ok(events) } - async fn handle_network_event( - &mut self, - event: &NetworkEvent, - requests: &RequestSender, - ) -> SyncResult> { - match event { - NetworkEvent::PeerConnected { - address, - } => { - self.handle_peer_connected(*address); - // If synced, activate the new peer immediately - if self.state() == SyncState::Synced - && self.peers.get(address).is_some_and(|v| v.is_none()) - { - tracing::info!("Activating mempool on newly connected peer {}", address); - self.activate_peer(*address, requests).await?; - } - } - NetworkEvent::PeerDisconnected { - address, - } => { - self.handle_peer_disconnected(*address); - } - NetworkEvent::PeersUpdated { - connected_count, - best_height, - .. - } => { - if let Some(best_height) = best_height { - self.update_target_height(*best_height); - } - if *connected_count == 0 { - self.stop_sync(); - } else if self.state() == SyncState::WaitingForConnections { - return self.start_sync(requests).await; - } - } - } - Ok(vec![]) - } - fn progress(&self) -> SyncManagerProgress { SyncManagerProgress::Mempool(self.progress.clone()) } @@ -191,35 +190,48 @@ impl SyncManager for MempoolManager { mod tests { use super::*; use crate::client::config::MempoolStrategy; - use crate::network::NetworkRequest; use crate::sync::BroadcastConfig; - use crate::test_utils::test_socket_address; + use crate::test_utils::{test_socket_address, MockNetworkManager}; use dashcore::hashes::Hash; use key_wallet_manager::test_utils::MockWallet; use std::collections::{BTreeMap, BTreeSet}; - use std::sync::Arc; - use tokio::sync::{mpsc, RwLock}; + use tokio::sync::RwLock; - fn create_test_manager( - ) -> (MempoolManager, RequestSender, mpsc::UnboundedReceiver) { + fn create_test_manager() -> MempoolManager { let wallet = Arc::new(RwLock::new(MockWallet::new())); - let (tx, rx) = mpsc::unbounded_channel::(); - let requests = RequestSender::new(tx); + MempoolManager::new(wallet, MempoolStrategy::FetchAll, 1000, 0, BroadcastConfig::default()) + } - let manager = MempoolManager::new( - wallet, - MempoolStrategy::FetchAll, - 1000, - 0, - BroadcastConfig::default(), - ); + /// Build a mock network manager and a trait-object handle to pass to + /// manager methods. Returns `(mock, network)` where `mock` is used for + /// assertions and `network` is passed by reference into the manager. + fn mock_network() -> (Arc, Arc) { + let mock = Arc::new(MockNetworkManager::new()); + let network: Arc = mock.clone(); + (mock, network) + } + + /// Did the manager send a `filterload` to any peer? + fn sent_filter_load(mock: &MockNetworkManager) -> bool { + mock.sent_to_messages().iter().any(|(_, m)| matches!(m, NetworkMessage::FilterLoad(_))) + } + + /// Did the manager send any bloom-filter message to any peer? + fn sent_any_filter_message(mock: &MockNetworkManager) -> bool { + mock.sent_to_messages() + .iter() + .any(|(_, m)| matches!(m, NetworkMessage::FilterLoad(_) | NetworkMessage::FilterClear)) + } - (manager, requests, rx) + fn filters_synced() -> SyncEvent { + SyncEvent::FiltersSyncComplete { + tip_height: 1000, + } } #[test] fn test_sync_manager_trait_basics() { - let (mut manager, _, _rx) = create_test_manager(); + let mut manager = create_test_manager(); assert_eq!(manager.identifier(), ManagerIdentifier::Mempool); assert_eq!(manager.state(), SyncState::WaitForEvents); @@ -237,15 +249,12 @@ mod tests { #[tokio::test] async fn test_filters_sync_complete_activates() { - let (mut manager, requests, _rx) = create_test_manager(); - let peer = crate::test_utils::test_socket_address(1); + let mut manager = create_test_manager(); + let (_mock, network) = mock_network(); + let peer = test_socket_address(1); manager.handle_peer_connected(peer); - let event = SyncEvent::FiltersSyncComplete { - tip_height: 1000, - }; - - let events = manager.handle_sync_event(&event, &requests).await.unwrap(); + let events = manager.handle_sync_event(&filters_synced(), &network).await.unwrap(); assert!(events.is_empty()); assert_eq!(manager.state(), SyncState::Synced); assert!(matches!(manager.peers.get(&peer), Some(Some(_)))); @@ -253,35 +262,31 @@ mod tests { #[tokio::test] async fn test_filters_sync_complete_subsequent_is_noop() { - let (mut manager, requests, _rx) = create_test_manager(); - manager.handle_peer_connected(crate::test_utils::test_socket_address(1)); + let mut manager = create_test_manager(); + let (_mock, network) = mock_network(); + manager.handle_peer_connected(test_socket_address(1)); // Activate first - let event0 = SyncEvent::FiltersSyncComplete { - tip_height: 1000, - }; - manager.handle_sync_event(&event0, &requests).await.unwrap(); + manager.handle_sync_event(&filters_synced(), &network).await.unwrap(); // Subsequent filter sync completions should not change state let event1 = SyncEvent::FiltersSyncComplete { tip_height: 1001, }; - let events = manager.handle_sync_event(&event1, &requests).await.unwrap(); + let events = manager.handle_sync_event(&event1, &network).await.unwrap(); assert!(events.is_empty()); assert_eq!(manager.state(), SyncState::Synced); } #[tokio::test] async fn test_reactivation_after_disconnect() { - let (mut manager, requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); + let (_mock, network) = mock_network(); let peer = test_socket_address(1); manager.handle_peer_connected(peer); // Initial activation - let event = SyncEvent::FiltersSyncComplete { - tip_height: 1000, - }; - let events = manager.handle_sync_event(&event, &requests).await.unwrap(); + let events = manager.handle_sync_event(&filters_synced(), &network).await.unwrap(); assert!(events.is_empty()); assert_eq!(manager.state(), SyncState::Synced); @@ -292,60 +297,53 @@ mod tests { let event = SyncEvent::FiltersSyncComplete { tip_height: 1001, }; - let events = manager.handle_sync_event(&event, &requests).await.unwrap(); + let events = manager.handle_sync_event(&event, &network).await.unwrap(); assert!(events.is_empty()); assert_eq!(manager.state(), SyncState::Synced); } #[tokio::test] async fn test_peer_connect_activates_when_synced() { - let (mut manager, requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); + let (_mock, network) = mock_network(); let peer1 = test_socket_address(1); manager.handle_peer_connected(peer1); // Activate via SyncComplete - let event = SyncEvent::FiltersSyncComplete { - tip_height: 1000, - }; - manager.handle_sync_event(&event, &requests).await.unwrap(); + manager.handle_sync_event(&filters_synced(), &network).await.unwrap(); assert!(matches!(manager.peers.get(&peer1), Some(Some(_)))); // New peer connects while synced => should activate immediately let peer2 = test_socket_address(2); - let connect = NetworkEvent::PeerConnected { - address: peer2, - }; - let events = manager.handle_network_event(&connect, &requests).await.unwrap(); + let events = manager + .handle_network_event(&NetworkEvent::PeerConnected(peer2), &network) + .await + .unwrap(); assert!(events.is_empty()); assert!(matches!(manager.peers.get(&peer2), Some(Some(_)))); } #[tokio::test] async fn test_network_event_peer_connect_disconnect() { - let (mut manager, requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); + let (_mock, network) = mock_network(); let peer1 = test_socket_address(1); let peer2 = test_socket_address(2); // Connecting peers should return empty events (not synced yet) - let connect1 = NetworkEvent::PeerConnected { - address: peer1, - }; - let events = manager.handle_network_event(&connect1, &requests).await.unwrap(); + let connect1 = NetworkEvent::PeerConnected(peer1); + let events = manager.handle_network_event(&connect1, &network).await.unwrap(); assert!(events.is_empty()); assert!(manager.peers.contains_key(&peer1)); - let connect2 = NetworkEvent::PeerConnected { - address: peer2, - }; - let events = manager.handle_network_event(&connect2, &requests).await.unwrap(); + let connect2 = NetworkEvent::PeerConnected(peer2); + let events = manager.handle_network_event(&connect2, &network).await.unwrap(); assert!(events.is_empty()); assert_eq!(manager.peers.len(), 2); - let disconnect1 = NetworkEvent::PeerDisconnected { - address: peer1, - }; - let events = manager.handle_network_event(&disconnect1, &requests).await.unwrap(); + let disconnect1 = NetworkEvent::PeerDisconnected(peer1); + let events = manager.handle_network_event(&disconnect1, &network).await.unwrap(); assert!(events.is_empty()); // Still have peer2 available @@ -353,21 +351,19 @@ mod tests { assert_eq!(manager.peers.len(), 1); // Disconnecting an already-disconnected peer should not error - let events = manager.handle_network_event(&disconnect1, &requests).await.unwrap(); + let events = manager.handle_network_event(&disconnect1, &network).await.unwrap(); assert!(events.is_empty()); } #[tokio::test] async fn test_block_processed_removes_confirmed_txids() { - let (mut manager, requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); + let (_mock, network) = mock_network(); let peer = test_socket_address(1); manager.handle_peer_connected(peer); // Activate - let sync = SyncEvent::FiltersSyncComplete { - tip_height: 1000, - }; - manager.handle_sync_event(&sync, &requests).await.unwrap(); + manager.handle_sync_event(&filters_synced(), &network).await.unwrap(); // Add transactions to mempool let mut txids = Vec::new(); @@ -401,7 +397,8 @@ mod tests { new_scripts: BTreeMap::new(), confirmed_txids: txids.clone(), }; - let events = manager.handle_sync_event(&event, &requests).await.unwrap(); + // No broadcasts were tracked, so confirming them emits nothing. + let events = manager.handle_sync_event(&event, &network).await.unwrap(); assert!(events.is_empty()); assert!(manager.transactions.is_empty()); @@ -409,15 +406,13 @@ mod tests { #[tokio::test] async fn test_instant_lock_received_marks_transaction() { - let (mut manager, requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); + let (_mock, network) = mock_network(); let peer = test_socket_address(1); manager.handle_peer_connected(peer); // Activate - let sync = SyncEvent::FiltersSyncComplete { - tip_height: 1000, - }; - manager.handle_sync_event(&sync, &requests).await.unwrap(); + manager.handle_sync_event(&filters_synced(), &network).await.unwrap(); // Add a transaction to mempool let tx = dashcore::Transaction { @@ -448,7 +443,8 @@ mod tests { instant_lock: is_lock, validated: true, }; - let events = manager.handle_sync_event(&event, &requests).await.unwrap(); + // The transaction was not self-broadcast, so no broadcast result fires. + let events = manager.handle_sync_event(&event, &network).await.unwrap(); assert!(events.is_empty()); assert!(manager.transactions.get(&txid).unwrap().is_instant_send); @@ -456,33 +452,27 @@ mod tests { #[tokio::test] async fn test_peer_disconnect_removes_from_peers() { - let (mut manager, requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); + let (_mock, network) = mock_network(); let peer = test_socket_address(1); manager.handle_peer_connected(peer); // Activate - let sync = SyncEvent::FiltersSyncComplete { - tip_height: 1000, - }; - manager.handle_sync_event(&sync, &requests).await.unwrap(); + manager.handle_sync_event(&filters_synced(), &network).await.unwrap(); // Disconnect the only peer - let disconnect = NetworkEvent::PeerDisconnected { - address: peer, - }; - let events = manager.handle_network_event(&disconnect, &requests).await.unwrap(); + let disconnect = NetworkEvent::PeerDisconnected(peer); + let events = manager.handle_network_event(&disconnect, &network).await.unwrap(); assert!(events.is_empty()); assert!(manager.peers.is_empty()); } #[tokio::test] async fn test_sync_complete_no_peers_stays_inactive() { - let (mut manager, requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); + let (_mock, network) = mock_network(); - let event = SyncEvent::FiltersSyncComplete { - tip_height: 1000, - }; - let events = manager.handle_sync_event(&event, &requests).await.unwrap(); + let events = manager.handle_sync_event(&filters_synced(), &network).await.unwrap(); assert!(events.is_empty()); assert_eq!(manager.state(), SyncState::WaitForEvents); @@ -491,78 +481,65 @@ mod tests { #[tokio::test] async fn test_start_sync_no_peers_stays_waiting() { - let (mut manager, requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); + let (_mock, network) = mock_network(); // Simulate full disconnect setting state to WaitingForConnections manager.set_state(SyncState::WaitingForConnections); // start_sync with no peers should stay in WaitingForConnections - let events = manager.start_sync(&requests).await.unwrap(); + let events = manager.start_sync(&network).await.unwrap(); assert!(events.is_empty()); assert_eq!(manager.state(), SyncState::WaitingForConnections); } #[tokio::test] async fn test_disconnect_recovery_reactivates_on_reconnect() { - let (mut manager, requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); + let (_mock, network) = mock_network(); let peer = test_socket_address(1); manager.handle_peer_connected(peer); // Activate via SyncComplete - let event = SyncEvent::FiltersSyncComplete { - tip_height: 1000, - }; - manager.handle_sync_event(&event, &requests).await.unwrap(); + manager.handle_sync_event(&filters_synced(), &network).await.unwrap(); assert_eq!(manager.state(), SyncState::Synced); - // Disconnect peer - let disconnect = NetworkEvent::PeerDisconnected { - address: peer, - }; - manager.handle_network_event(&disconnect, &requests).await.unwrap(); - - // PeersUpdated with 0 triggers stop_sync - let update = NetworkEvent::PeersUpdated { - connected_count: 0, - addresses: vec![], - best_height: None, - }; - manager.handle_network_event(&update, &requests).await.unwrap(); - assert_eq!(manager.state(), SyncState::WaitingForConnections); + // Every peer drops: the coordinator resets the manager to + // WaitingForConnections after `on_disconnect`. + manager + .handle_network_event(&NetworkEvent::PeerDisconnected(peer), &network) + .await + .unwrap(); + manager.on_disconnect(); + manager.set_state(SyncState::WaitingForConnections); - // PeersUpdated with 1 but no peers tracked yet: stays WaitingForConnections + // PeersUpdated with no peers tracked yet: nothing to activate, so the + // manager stays waiting. let update = NetworkEvent::PeersUpdated { connected_count: 1, - addresses: vec![peer], - best_height: Some(1000), + best_height: 1000, }; - manager.handle_network_event(&update, &requests).await.unwrap(); + manager.handle_network_event(&update, &network).await.unwrap(); assert_eq!(manager.state(), SyncState::WaitingForConnections); // Peer reconnects and PeersUpdated fires again manager.handle_peer_connected(peer); - let update = NetworkEvent::PeersUpdated { - connected_count: 1, - addresses: vec![peer], - best_height: Some(1000), - }; - manager.handle_network_event(&update, &requests).await.unwrap(); + manager.handle_network_event(&update, &network).await.unwrap(); assert_eq!(manager.state(), SyncState::Synced); assert!(matches!(manager.peers.get(&peer), Some(Some(_)))); } #[tokio::test] async fn test_block_processed_confirmed_txids_does_not_eagerly_rebuild() { - let mut mock = MockWallet::new(); + let mut mock_wallet = MockWallet::new(); let script = dashcore::ScriptBuf::from_bytes(vec![ 0x76, 0xa9, 0x14, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0x88, 0xac, ]); let addr = dashcore::Address::from_script(&script, dashcore::Network::Testnet).unwrap(); - mock.set_addresses(vec![addr]); - let wallet = Arc::new(RwLock::new(mock)); - let (tx, mut rx) = mpsc::unbounded_channel::(); - let requests = RequestSender::new(tx); + mock_wallet.set_addresses(vec![addr]); + let wallet = Arc::new(RwLock::new(mock_wallet)); + let (mock, network) = mock_network(); let mut manager = MempoolManager::new( wallet, @@ -576,13 +553,10 @@ mod tests { manager.handle_peer_connected(peer); // Activate - let sync = SyncEvent::FiltersSyncComplete { - tip_height: 1000, - }; - manager.handle_sync_event(&sync, &requests).await.unwrap(); + manager.handle_sync_event(&filters_synced(), &network).await.unwrap(); - // Drain activation messages - while rx.try_recv().is_ok() {} + // Drop the activation messages + mock.clear_sent(); // BlockProcessed does not eagerly rebuild — the tick handles it via // the revision check. Verify no FilterLoad is sent from the event handler. @@ -593,24 +567,19 @@ mod tests { new_scripts: BTreeMap::new(), confirmed_txids: vec![dashcore::Txid::all_zeros()], }; - manager.handle_sync_event(&event, &requests).await.unwrap(); + manager.handle_sync_event(&event, &network).await.unwrap(); - let has_filter_load = std::iter::from_fn(|| rx.try_recv().ok()).any(|req| { - matches!(req, NetworkRequest::SendMessageToPeer(NetworkMessage::FilterLoad(_), _)) - }); - assert!(!has_filter_load, "BlockProcessed should not eagerly rebuild filter"); + assert!(!sent_filter_load(&mock), "BlockProcessed should not eagerly rebuild filter"); } #[tokio::test] async fn test_block_processed_no_changes_no_rebuild_flag() { - let (mut manager, requests, _rx) = create_test_manager(); + let mut manager = create_test_manager(); + let (_mock, network) = mock_network(); let peer = test_socket_address(1); manager.handle_peer_connected(peer); - let sync = SyncEvent::FiltersSyncComplete { - tip_height: 1000, - }; - manager.handle_sync_event(&sync, &requests).await.unwrap(); + manager.handle_sync_event(&filters_synced(), &network).await.unwrap(); // BlockProcessed with no confirmed txids and no new addresses let event = SyncEvent::BlockProcessed { @@ -620,7 +589,7 @@ mod tests { new_scripts: BTreeMap::new(), confirmed_txids: vec![], }; - manager.handle_sync_event(&event, &requests).await.unwrap(); + manager.handle_sync_event(&event, &network).await.unwrap(); } #[tokio::test] @@ -633,12 +602,11 @@ mod tests { dashcore::Address::from_script(&script, dashcore::Network::Testnet).unwrap() }; - let mut mock = MockWallet::new(); - mock.set_addresses(vec![addr.clone()]); - let initial_revision = mock.monitor_revision(); - let wallet = Arc::new(RwLock::new(mock)); - let (tx, mut rx) = mpsc::unbounded_channel::(); - let requests = RequestSender::new(tx); + let mut mock_wallet = MockWallet::new(); + mock_wallet.set_addresses(vec![addr.clone()]); + let initial_revision = mock_wallet.monitor_revision(); + let wallet = Arc::new(RwLock::new(mock_wallet)); + let (mock, network) = mock_network(); let mut manager = MempoolManager::new( wallet.clone(), @@ -652,18 +620,15 @@ mod tests { manager.handle_peer_connected(peer); // Activate — this snapshots the monitor revision - let sync = SyncEvent::FiltersSyncComplete { - tip_height: 1000, - }; - manager.handle_sync_event(&sync, &requests).await.unwrap(); + manager.handle_sync_event(&filters_synced(), &network).await.unwrap(); assert_eq!(manager.state(), SyncState::Synced); - // Drain activation messages - while rx.try_recv().is_ok() {} + // Drop the activation messages + mock.clear_sent(); // tick with unchanged revision should not rebuild - manager.tick(&requests).await.unwrap(); - assert!(rx.try_recv().is_err(), "no messages expected when revision unchanged"); + manager.tick(&network).await.unwrap(); + assert!(mock.sent_to_messages().is_empty(), "no messages expected when revision unchanged"); // Simulate wallet adding new addresses (bumps revision) { @@ -680,26 +645,22 @@ mod tests { } // tick should detect stale filter and rebuild - manager.tick(&requests).await.unwrap(); - - let mut found_filter_load = false; - while let Ok(msg) = rx.try_recv() { - if matches!(msg, NetworkRequest::SendMessageToPeer(NetworkMessage::FilterLoad(_), _)) { - found_filter_load = true; - } - } - assert!(found_filter_load, "expected FilterLoad after monitor revision change"); + manager.tick(&network).await.unwrap(); + assert!(sent_filter_load(&mock), "expected FilterLoad after monitor revision change"); // Subsequent tick should not rebuild again (revision was snapshotted) - manager.tick(&requests).await.unwrap(); - assert!(rx.try_recv().is_err(), "no messages expected after revision re-snapshot"); + mock.clear_sent(); + manager.tick(&network).await.unwrap(); + assert!( + mock.sent_to_messages().is_empty(), + "no messages expected after revision re-snapshot" + ); } #[tokio::test] async fn test_tick_skips_rebuild_for_fetch_all_strategy() { let wallet = Arc::new(RwLock::new(MockWallet::new())); - let (tx, mut rx) = mpsc::unbounded_channel::(); - let requests = RequestSender::new(tx); + let (mock, network) = mock_network(); let mut manager = MempoolManager::new( wallet.clone(), @@ -712,11 +673,8 @@ mod tests { let peer = test_socket_address(1); manager.handle_peer_connected(peer); - let sync = SyncEvent::FiltersSyncComplete { - tip_height: 1000, - }; - manager.handle_sync_event(&sync, &requests).await.unwrap(); - while rx.try_recv().is_ok() {} + manager.handle_sync_event(&filters_synced(), &network).await.unwrap(); + mock.clear_sent(); // Bump revision { @@ -725,18 +683,11 @@ mod tests { } // tick should not send any filter messages for FetchAll - manager.tick(&requests).await.unwrap(); - let mut found_filter = false; - while let Ok(msg) = rx.try_recv() { - if matches!( - msg, - NetworkRequest::SendMessageToPeer(NetworkMessage::FilterLoad(_), _) - | NetworkRequest::SendMessageToPeer(NetworkMessage::FilterClear, _) - ) { - found_filter = true; - } - } - assert!(!found_filter, "FetchAll should not send filter messages on revision change"); + manager.tick(&network).await.unwrap(); + assert!( + !sent_any_filter_message(&mock), + "FetchAll should not send filter messages on revision change" + ); } #[tokio::test] @@ -749,12 +700,11 @@ mod tests { dashcore::Address::from_script(&script, dashcore::Network::Testnet).unwrap() }; - let mut mock = MockWallet::new(); - mock.set_addresses(vec![addr]); - let initial_revision = mock.monitor_revision(); - let wallet = Arc::new(RwLock::new(mock)); - let (tx, mut rx) = mpsc::unbounded_channel::(); - let requests = RequestSender::new(tx); + let mut mock_wallet = MockWallet::new(); + mock_wallet.set_addresses(vec![addr]); + let initial_revision = mock_wallet.monitor_revision(); + let wallet = Arc::new(RwLock::new(mock_wallet)); + let (mock, network) = mock_network(); let mut manager = MempoolManager::new( wallet.clone(), @@ -767,11 +717,8 @@ mod tests { let peer = test_socket_address(1); manager.handle_peer_connected(peer); - let sync = SyncEvent::FiltersSyncComplete { - tip_height: 1000, - }; - manager.handle_sync_event(&sync, &requests).await.unwrap(); - while rx.try_recv().is_ok() {} + manager.handle_sync_event(&filters_synced(), &network).await.unwrap(); + mock.clear_sent(); // Simulate UTXO set change (new outpoint added) { @@ -783,28 +730,23 @@ mod tests { } // tick should detect the revision change and rebuild - manager.tick(&requests).await.unwrap(); - - let found_filter_load = std::iter::from_fn(|| rx.try_recv().ok()).any(|msg| { - matches!(msg, NetworkRequest::SendMessageToPeer(NetworkMessage::FilterLoad(_), _)) - }); - assert!(found_filter_load, "expected FilterLoad after outpoint change"); + manager.tick(&network).await.unwrap(); + assert!(sent_filter_load(&mock), "expected FilterLoad after outpoint change"); } #[tokio::test] async fn test_handle_tx_does_not_eagerly_rebuild_filter() { - let mut mock = MockWallet::new(); - mock.set_mempool_relevant(true); + let mut mock_wallet = MockWallet::new(); + mock_wallet.set_mempool_relevant(true); let script = dashcore::ScriptBuf::from_bytes(vec![ 0x76, 0xa9, 0x14, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0x88, 0xac, ]); let addr = dashcore::Address::from_script(&script, dashcore::Network::Testnet).unwrap(); - mock.set_addresses(vec![addr]); - let initial_revision = mock.monitor_revision(); - let wallet = Arc::new(RwLock::new(mock)); - let (tx_chan, mut rx) = mpsc::unbounded_channel::(); - let requests = RequestSender::new(tx_chan); + mock_wallet.set_addresses(vec![addr]); + let initial_revision = mock_wallet.monitor_revision(); + let wallet = Arc::new(RwLock::new(mock_wallet)); + let (mock, network) = mock_network(); let mut manager = MempoolManager::new( wallet.clone(), @@ -817,11 +759,8 @@ mod tests { let peer = test_socket_address(1); manager.handle_peer_connected(peer); - let sync = SyncEvent::FiltersSyncComplete { - tip_height: 1000, - }; - manager.handle_sync_event(&sync, &requests).await.unwrap(); - while rx.try_recv().is_ok() {} + manager.handle_sync_event(&filters_synced(), &network).await.unwrap(); + mock.clear_sent(); // handle_tx with a relevant transaction should NOT eagerly rebuild let tx = dashcore::Transaction { @@ -831,12 +770,9 @@ mod tests { output: vec![], special_transaction_payload: None, }; - manager.handle_tx(tx, test_socket_address(1), &requests).await.unwrap(); + manager.handle_tx(tx, test_socket_address(1), &network).await.unwrap(); - let has_filter_load = std::iter::from_fn(|| rx.try_recv().ok()).any(|msg| { - matches!(msg, NetworkRequest::SendMessageToPeer(NetworkMessage::FilterLoad(_), _)) - }); - assert!(!has_filter_load, "handle_tx should not eagerly rebuild filter"); + assert!(!sent_filter_load(&mock), "handle_tx should not eagerly rebuild filter"); // But the next tick should catch it if the wallet revision changed // (MockWallet bumps revision when set_mempool_relevant triggers processing) @@ -844,11 +780,8 @@ mod tests { let mut w = wallet.write().await; w.set_addresses(vec![dashcore::Address::dummy(dashcore::Network::Testnet, 0)]); } - manager.tick(&requests).await.unwrap(); + manager.tick(&network).await.unwrap(); - let found_filter_load = std::iter::from_fn(|| rx.try_recv().ok()).any(|msg| { - matches!(msg, NetworkRequest::SendMessageToPeer(NetworkMessage::FilterLoad(_), _)) - }); - assert!(found_filter_load, "tick should rebuild after revision change"); + assert!(sent_filter_load(&mock), "tick should rebuild after revision change"); } } diff --git a/dash-spv/src/sync/mod.rs b/dash-spv/src/sync/mod.rs index f059ebc09..525a0440d 100644 --- a/dash-spv/src/sync/mod.rs +++ b/dash-spv/src/sync/mod.rs @@ -3,7 +3,6 @@ mod block_headers; mod blocks; mod chainlock; -pub(super) mod download_coordinator; mod events; mod filter_headers; mod filters; diff --git a/dash-spv/src/sync/sync_coordinator.rs b/dash-spv/src/sync/sync_coordinator.rs index 7f9bc245d..6d4d48fa9 100644 --- a/dash-spv/src/sync/sync_coordinator.rs +++ b/dash-spv/src/sync/sync_coordinator.rs @@ -12,6 +12,8 @@ use tokio::task::JoinSet; use tokio_stream::wrappers::WatchStream; use tokio_util::sync::CancellationToken; +use std::sync::Arc; + use crate::error::SyncResult; use crate::network::NetworkManager; use crate::storage::{ @@ -35,9 +37,8 @@ macro_rules! spawn_manager { if let Some(manager) = $manager { let identifier = manager.identifier(); let wanted_message_types = manager.wanted_message_types(); - let requests = $network.request_sender(); - let message_receiver = $network.message_receiver(wanted_message_types).await; - let network_event_rx = $network.subscribe_network_events(); + let message_receiver = $network.subscribe(wanted_message_types).await; + let network_event_rx = $network.events(); let (progress_sender, progress_receiver) = watch::channel(manager.progress()); tracing::info!( @@ -50,7 +51,7 @@ macro_rules! spawn_manager { message_receiver, sync_event_sender: $self.sync_event_sender.clone(), network_event_receiver: network_event_rx, - requests, + network: $network.clone(), shutdown: $self.shutdown.clone(), progress_sender, }; @@ -192,12 +193,9 @@ where /// Each manager receives: /// - A message stream filtered by its subscribed types /// - An event bus subscription for inter-manager events - /// - A request sender for outgoing network messages + /// - A handle to the network manager, to declare requests on /// - A shutdown token for graceful termination - pub async fn start(&mut self, network: &mut N) -> SyncResult<()> - where - N: NetworkManager, - { + pub async fn start(&mut self, network: &Arc) -> SyncResult<()> { if !self.tasks.is_empty() { return Err(SyncError::InvalidState("SyncCoordinator already started".to_string())); } diff --git a/dash-spv/src/sync/sync_manager.rs b/dash-spv/src/sync/sync_manager.rs index 997e9f505..7c0443d9a 100644 --- a/dash-spv/src/sync/sync_manager.rs +++ b/dash-spv/src/sync/sync_manager.rs @@ -1,11 +1,14 @@ use crate::error::SyncResult; -use crate::network::{Message, MessageType, NetworkEvent, RequestSender}; +use crate::network::{MessageType, NetworkEvent, NetworkManager}; use crate::sync::{ BlockHeadersProgress, BlocksProgress, ChainLockProgress, FilterHeadersProgress, FiltersProgress, InstantSendProgress, ManagerIdentifier, MasternodesProgress, MempoolProgress, SyncEvent, SyncState, }; use async_trait::async_trait; +use dashcore::network::message::NetworkMessage; +use std::net::SocketAddr; +use std::sync::Arc; use crate::SyncError; @@ -48,11 +51,13 @@ impl SyncManagerProgress { } } +pub type Inbound = (SocketAddr, Arc); + pub struct SyncManagerTaskContext { - pub(super) message_receiver: UnboundedReceiver, + pub(super) message_receiver: UnboundedReceiver, pub(super) sync_event_sender: broadcast::Sender, pub(super) network_event_receiver: broadcast::Receiver, - pub(super) requests: RequestSender, + pub(super) network: Arc, pub(super) shutdown: CancellationToken, pub(super) progress_sender: watch::Sender, } @@ -68,6 +73,51 @@ impl SyncManagerTaskContext { } } +// Display for the network event so the sync loop / broadcast monitor can log +// it (they require `Display`). Kept here to avoid modifying the network module. +impl std::fmt::Display for NetworkEvent { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + NetworkEvent::PeersUpdated { + connected_count, + .. + } => write!(f, "PeersUpdated({connected_count} peers)"), + NetworkEvent::PeerConnected(addr) => write!(f, "PeerConnected({addr})"), + NetworkEvent::PeerDisconnected(addr) => write!(f, "PeerDisconnected({addr})"), + } + } +} + +/// The default [`SyncManager::handle_network_event`] body, callable from an override. +/// +/// A manager that only cares about some `NetworkEvent` variants overrides +/// `handle_network_event` and delegates the rest here. Rust gives no way to call a +/// trait's default body from an override, so the shared logic lives in this free +/// function — the trait's default method just calls it. +pub(super) async fn default_handle_network_event( + manager: &mut M, + event: &NetworkEvent, + network: &Arc, +) -> SyncResult> { + // `PeersUpdated` is the cue to kick off the initial requests: the network manager + // connects in `start`, after every manager has subscribed, so this is the first thing + // we hear from it. Individual peer disconnects are recovered by per-request + // timeout+retry, so they don't stop sync. + if let NetworkEvent::PeersUpdated { + .. + } = event + { + // Seed every manager's target from the peers' advertised tip so the + // height shows up right away (matches the pre-network `best_height`). + manager.update_target_height(network.tip()); + if manager.state() == SyncState::WaitingForConnections { + tracing::info!("{} - peers available, starting sync", manager.identifier()); + return manager.start_sync(network).await; + } + } + Ok(vec![]) +} + /// Guard that verifies a manager has not already been started. pub(super) fn ensure_not_started( state: SyncState, @@ -105,7 +155,10 @@ pub trait SyncManager: Send + Sync + std::fmt::Debug { /// Called after initialization to trigger the initial sync requests. /// For example, BlockHeadersManager sends its first getheaders request here. /// The default implementation is for reactive managers that just wait for events. - async fn start_sync(&mut self, _requests: &RequestSender) -> SyncResult> { + async fn start_sync( + &mut self, + _network: &Arc, + ) -> SyncResult> { ensure_not_started(self.state(), self.identifier())?; self.set_state(SyncState::WaitForEvents); Ok(vec![SyncEvent::SyncStart { @@ -127,10 +180,6 @@ pub trait SyncManager: Send + Sync + std::fmt::Debug { /// derivable from durable storage (block headers, filter headers, the /// masternode engine) or from preserved per-batch bookkeeping should /// survive so reconnect resumes instead of restarting. - /// - /// `BlocksManager` and `FiltersManager` go further and requeue their - /// in-flight network slots so the next `send_pending` reissues them - /// immediately to the new peer. fn on_disconnect(&mut self); /// Handle an incoming network message. @@ -138,8 +187,9 @@ pub trait SyncManager: Send + Sync + std::fmt::Debug { /// Returns events to emit to other managers. async fn handle_message( &mut self, - msg: Message, - requests: &RequestSender, + peer: SocketAddr, + msg: NetworkMessage, + network: &Arc, ) -> SyncResult>; /// Handle a sync event from another manager. @@ -150,50 +200,31 @@ pub trait SyncManager: Send + Sync + std::fmt::Debug { async fn handle_sync_event( &mut self, event: &SyncEvent, - requests: &RequestSender, + network: &Arc, ) -> SyncResult>; - /// Periodic tick for timeouts, retries, and proactive work. + /// Periodic tick for work that only a clock can trigger. /// - /// Called regularly by the coordinator (e.g., every 100ms). - /// Use this for: - /// - Timeout detection and retry logic - /// - Proactive request sending - /// - State cleanup - async fn tick(&mut self, requests: &RequestSender) -> SyncResult>; + /// Called every 100ms by the coordinator. Reserved for genuinely time-based + /// work — expiry, retry schedules the broker does not own, polling — and for + /// processing that cannot be driven from an arrival. It is *not* the place to + /// re-declare requests: once declared, the broker owns their timeout, retry and + /// peer hot-swap, so re-offering them costs a registry lock per item and buys + /// nothing. Managers with no such work leave this alone. + async fn tick(&mut self, _network: &Arc) -> SyncResult> { + Ok(vec![]) + } /// Handle a network event (peer connection changes). /// - /// Default implementation handles state transitions for WaitingForConnections. - /// Managers can override to customize behavior. + /// The default body handles state transitions for `WaitingForConnections`. + /// Managers can override this to customize behavior. async fn handle_network_event( &mut self, event: &NetworkEvent, - requests: &RequestSender, + network: &Arc, ) -> SyncResult> { - // Default: transition from WaitingForConnections to Syncing when peers connect - if let NetworkEvent::PeersUpdated { - connected_count, - best_height, - .. - } = event - { - if let Some(best_height) = best_height { - self.update_target_height(*best_height); - } - if *connected_count == 0 { - tracing::info!("{} - no peers available, stopping sync", self.identifier()); - self.stop_sync(); - } else if *connected_count > 0 && self.state() == SyncState::WaitingForConnections { - tracing::info!( - "{} - peers available ({}), starting sync", - self.identifier(), - connected_count - ); - return self.start_sync(requests).await; - } - } - Ok(vec![]) + default_handle_network_event(self, event, network).await } /// Retrieves the current progress of the Manager. @@ -234,10 +265,16 @@ pub trait SyncManager: Send + Sync + std::fmt::Debug { break; } // Process incoming network messages - Some(message) = context.message_receiver.recv() => { + Some((peer, message)) = context.message_receiver.recv() => { tracing::trace!("{} received message: {}", identifier, message.cmd()); + // The pump gives its last subscriber the sole reference, so for the + // message types only one manager watches — block, cfilter, cfheaders, + // headers — this takes the payload without copying it. A genuinely + // shared message (`inv` goes to several managers) falls back to a clone. + let message = + Arc::try_unwrap(message).unwrap_or_else(|shared| (*shared).clone()); let progress_before = self.progress(); - match self.handle_message(message, &context.requests).await { + match self.handle_message(peer, message, &context.network).await { Ok(events) => { if !events.is_empty() { for event in &events { @@ -263,7 +300,7 @@ pub trait SyncManager: Send + Sync + std::fmt::Debug { Ok(event) => { tracing::trace!("{} received event: {}", identifier, event); let progress_before = self.progress(); - match self.handle_sync_event(&event, &context.requests).await { + match self.handle_sync_event(&event, &context.network).await { Ok(events) => { if !events.is_empty() { for e in &events { @@ -278,6 +315,13 @@ pub trait SyncManager: Send + Sync + std::fmt::Debug { } } } + Err(broadcast::error::RecvError::Lagged(n)) => { + // Sync-event bus overflowed for this manager; skipped `n` + // events. Keep running rather than killing the task: a dead + // manager is a stalled sync. Nothing recovers the skipped + // events, so the bus is sized (10k) not to overflow. + tracing::warn!("{} lagged sync events, skipped {}", identifier, n); + } Err(error) => { tracing::error!("{} sync event error: {}", identifier, error); break; @@ -290,7 +334,7 @@ pub trait SyncManager: Send + Sync + std::fmt::Debug { Ok(event) => { tracing::debug!("{} received network event: {}", identifier, event); let progress_before = self.progress(); - match self.handle_network_event(&event, &context.requests).await { + match self.handle_network_event(&event, &context.network).await { Ok(events) => { if !events.is_empty() { for e in &events { @@ -305,6 +349,12 @@ pub trait SyncManager: Send + Sync + std::fmt::Debug { } } } + Err(broadcast::error::RecvError::Lagged(n)) => { + // Network-event bus overflowed. The bus carries only + // low-volume events (peer churn), so lagging 4096 behind is + // not realistic; keep running rather than kill the manager. + tracing::warn!("{} lagged network events, skipped {}", identifier, n); + } Err(error) => { tracing::error!("{} network event error: {}", identifier, error); break; @@ -314,7 +364,7 @@ pub trait SyncManager: Send + Sync + std::fmt::Debug { // Periodic tick for timeouts and housekeeping _ = tick_interval.tick() => { let progress_before = self.progress(); - match self.tick(&context.requests).await { + match self.tick(&context.network).await { Ok(events) => { if !events.is_empty() { context.emit_sync_events(events); @@ -337,20 +387,16 @@ pub trait SyncManager: Send + Sync + std::fmt::Debug { #[cfg(test)] mod tests { use super::*; - use crate::network::NetworkRequest; use crate::sync::BlockHeadersProgress; - use crate::sync::SyncState; - use async_trait::async_trait; + use crate::test_utils::MockNetworkManager; use std::sync::atomic::{AtomicU32, Ordering}; - use std::sync::Arc; - use tokio::sync::{broadcast, mpsc}; + use tokio::sync::mpsc; - /// Mock manager for testing the task runner. + /// Minimal manager that only counts the callbacks the run loop makes, so a + /// test can observe that the loop is alive and that it stops on cancel. struct MockManager { identifier: ManagerIdentifier, state: SyncState, - message_count: Arc, - event_count: Arc, tick_count: Arc, } @@ -382,23 +428,22 @@ mod tests { async fn handle_message( &mut self, - _msg: Message, - _requests: &RequestSender, + _peer: SocketAddr, + _msg: NetworkMessage, + _network: &Arc, ) -> SyncResult> { - self.message_count.fetch_add(1, Ordering::Relaxed); Ok(vec![]) } async fn handle_sync_event( &mut self, _event: &SyncEvent, - _requests: &RequestSender, + _network: &Arc, ) -> SyncResult> { - self.event_count.fetch_add(1, Ordering::Relaxed); Ok(vec![]) } - async fn tick(&mut self, _requests: &RequestSender) -> SyncResult> { + async fn tick(&mut self, _network: &Arc) -> SyncResult> { self.tick_count.fetch_add(1, Ordering::Relaxed); Ok(vec![]) } @@ -410,26 +455,23 @@ mod tests { } } + /// The shared run loop must keep ticking while it lives and return its + /// identifier once the shutdown token is cancelled — a manager task that + /// ignored the token would hang `SyncCoordinator::shutdown` forever. #[tokio::test] async fn test_manager_task_shutdown() { - let message_count = Arc::new(AtomicU32::new(0)); - let event_count = Arc::new(AtomicU32::new(0)); let tick_count = Arc::new(AtomicU32::new(0)); let manager = MockManager { identifier: ManagerIdentifier::BlockHeader, state: SyncState::WaitForEvents, - message_count: message_count.clone(), - event_count: event_count.clone(), tick_count: tick_count.clone(), }; - // Create channels - let (_, message_receiver) = mpsc::unbounded_channel(); + let (_msg_tx, message_receiver) = mpsc::unbounded_channel(); let sync_event_sender = broadcast::Sender::::new(100); let network_event_sender = broadcast::Sender::::new(100); - let (req_tx, _req_rx) = mpsc::unbounded_channel::(); - let requests = RequestSender::new(req_tx); + let network: Arc = Arc::new(MockNetworkManager::new()); let shutdown = CancellationToken::new(); let (progress_sender, _progress_rx) = watch::channel(manager.progress()); @@ -437,28 +479,23 @@ mod tests { message_receiver, sync_event_sender, network_event_receiver: network_event_sender.subscribe(), - requests, + network, shutdown: shutdown.clone(), progress_sender, }; - // Spawn the task using trait's run method let handle = tokio::spawn(async move { manager.run(context).await }); - // Let it run for a bit + // Let the 100ms tick fire a few times. tokio::time::sleep(Duration::from_millis(250)).await; - - // Signal shutdown shutdown.cancel(); - // Wait for task to complete let result = handle.await.unwrap(); - assert!(result.is_ok()); - - // Verify the returned identifier matches assert_eq!(result.unwrap(), ManagerIdentifier::BlockHeader); - - // Verify tick was called multiple times - assert!(tick_count.load(Ordering::Relaxed) > 0); + assert!( + tick_count.load(Ordering::Relaxed) >= 2, + "the run loop should have ticked while alive, got {}", + tick_count.load(Ordering::Relaxed) + ); } } diff --git a/dash-spv/src/test_utils/network.rs b/dash-spv/src/test_utils/network.rs index 5ac883944..dab645787 100644 --- a/dash-spv/src/test_utils/network.rs +++ b/dash-spv/src/test_utils/network.rs @@ -1,193 +1,186 @@ -use crate::error::{NetworkError, NetworkResult}; -use crate::network::peer::Peer; -use crate::network::{ - Message, MessageDispatcher, MessageType, NetworkEvent, NetworkManager, NetworkRequest, - RequestSender, -}; +//! A lightweight in-memory [`NetworkManager`] for unit tests. +//! +//! It records everything the sync layer sends (so a test can assert on the +//! requests a manager/pipeline issued), lets a test inject inbound messages and +//! peer events, and exposes the advertised tip / connected-peer count. No sockets, +//! no background tasks, no DNS. + +// The recorders below need a *sync* mutex: `NetworkManager::broadcast` is +// non-async by design (the real implementation spawns and returns), so a +// `tokio::sync::Mutex` cannot be locked there without either making the trait +// method async or spawning — and spawning would race the assertion that follows +// `broadcast()` in a test. Every guard here is a short-lived temporary that never +// crosses an `.await`, so the crate-wide ban on `std::sync::Mutex` (which exists +// to stop guards being held across await points) does not apply to this module. +#![allow(clippy::disallowed_types)] + +use std::net::SocketAddr; +use std::sync::atomic::{AtomicU32, Ordering}; +use std::sync::Mutex; + use async_trait::async_trait; -use dashcore::{ - block::Header as BlockHeader, network::message::NetworkMessage, - network::message_blockdata::GetHeadersMessage, BlockHash, Network, -}; -use dashcore_hashes::Hash; -use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4}; -use std::time::Duration; +use dashcore::network::message::NetworkMessage; use tokio::sync::broadcast; use tokio::sync::mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender}; -use tokio::sync::Mutex; +use crate::network::{Inbound, MessageType, NetworkEvent, NetworkManager, RequestKey}; + +/// Deterministic loopback socket address for tests (`127.0.0.1:`). pub fn test_socket_address(id: u8) -> SocketAddr { - SocketAddr::from(([127, 0, 0, id], id as u16)) + use std::net::{IpAddr, Ipv4Addr}; + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 40000 + id as u16) } -/// Mock network manager for testing +struct Subscriber { + kinds: Vec, + tx: UnboundedSender, +} + +/// In-memory mock of the peer-to-peer network manager for unit tests. pub struct MockNetworkManager { - connected: bool, - connected_peer: SocketAddr, - headers_chain: Vec, - message_dispatcher: Mutex, - sent_messages: Vec, - /// Request sender for outgoing messages. - request_tx: UnboundedSender, - /// Receiver generated in the constructor. Can be taken out of the struct for testing. - request_rx: Option>, - /// Event bus for network events. - network_event_sender: broadcast::Sender, + sent: Mutex>, + sent_to: Mutex>, + broadcasts: Mutex>, + answered: Mutex>, + completed: Mutex>, + subscribers: Mutex>, + events_tx: broadcast::Sender, + tip: AtomicU32, + connected: AtomicU32, +} + +impl Default for MockNetworkManager { + fn default() -> Self { + Self::new() + } } impl MockNetworkManager { - /// Create a new mock network manager pub fn new() -> Self { - let (request_tx, request_rx) = unbounded_channel(); + let (events_tx, _) = broadcast::channel(1024); Self { - connected: true, - connected_peer: SocketAddr::new(std::net::Ipv4Addr::LOCALHOST.into(), 9999), - headers_chain: Vec::new(), - message_dispatcher: Mutex::new(MessageDispatcher::default()), - sent_messages: Vec::new(), - request_tx, - request_rx: Some(request_rx), - network_event_sender: broadcast::Sender::new(100000), + sent: Mutex::new(Vec::new()), + sent_to: Mutex::new(Vec::new()), + broadcasts: Mutex::new(Vec::new()), + answered: Mutex::new(Vec::new()), + completed: Mutex::new(Vec::new()), + subscribers: Mutex::new(Vec::new()), + events_tx, + tip: AtomicU32::new(0), + connected: AtomicU32::new(1), } } - pub fn take_receiver(&mut self) -> Option> { - self.request_rx.take() + /// Every message declared via [`NetworkManager::send`], in order. + pub fn sent_messages(&self) -> Vec { + self.sent.lock().expect("mock mutex poisoned").clone() } - /// Add a chain of headers for testing - pub fn add_headers_chain(&mut self, genesis_hash: BlockHash, count: usize) { - let mut headers = Vec::new(); - let mut prev_hash = genesis_hash; + /// Every `(peer, message)` sent via [`NetworkManager::send_to`], in order. + pub fn sent_to_messages(&self) -> Vec<(SocketAddr, NetworkMessage)> { + self.sent_to.lock().expect("mock mutex poisoned").clone() + } - // Skip genesis (height 0) as it's already in the storage - for i in 1..count { - let header = BlockHeader { - version: dashcore::block::Version::from_consensus(1), - prev_blockhash: prev_hash, - merkle_root: dashcore::hashes::sha256d::Hash::all_zeros().into(), - time: 1000000 + i as u32, - bits: dashcore::CompactTarget::from_consensus(0x207fffff), - nonce: i as u32, - }; + /// Every message broadcast via [`NetworkManager::broadcast`], in order. + pub fn broadcast_messages(&self) -> Vec { + self.broadcasts.lock().expect("mock mutex poisoned").clone() + } - prev_hash = header.block_hash(); - headers.push(header); - } + /// Every request key reported via [`NetworkManager::request_answered`]. + pub fn answered_keys(&self) -> Vec { + self.answered.lock().expect("mock mutex poisoned").clone() + } - self.headers_chain = headers; - } - - /// Process GetHeaders request and return appropriate headers - fn process_getheaders(&self, msg: &GetHeadersMessage) -> Vec { - // Find the starting point in our chain - let start_idx = if msg.locator_hashes.is_empty() { - 0 - } else { - // Find the first locator hash we recognize - let mut found_idx = None; - for locator in &msg.locator_hashes { - for (idx, header) in self.headers_chain.iter().enumerate() { - if header.block_hash() == *locator { - found_idx = Some(idx + 1); // Start from next header - break; - } - } - if found_idx.is_some() { - break; - } - } - found_idx.unwrap_or(0) - }; + /// Every `(peer, n)` reported via [`NetworkManager::request_completed`]. + pub fn completed_requests(&self) -> Vec<(SocketAddr, usize)> { + self.completed.lock().expect("mock mutex poisoned").clone() + } - // Return up to 2000 headers starting from start_idx - let end_idx = (start_idx + 2000).min(self.headers_chain.len()); + /// Clear all recorded sends (handy between phases of a test). + pub fn clear_sent(&self) { + self.sent.lock().expect("mock mutex poisoned").clear(); + self.sent_to.lock().expect("mock mutex poisoned").clear(); + self.broadcasts.lock().expect("mock mutex poisoned").clear(); + } - if start_idx < self.headers_chain.len() { - self.headers_chain[start_idx..end_idx].to_vec() - } else { - Vec::new() - } + /// Set the tip height reported by [`NetworkManager::tip`]. + pub fn set_tip(&self, tip: u32) { + self.tip.store(tip, Ordering::SeqCst); } - pub fn sent_messages(&self) -> &Vec { - &self.sent_messages + /// Set the count reported by [`NetworkManager::connected_count`]. + pub fn set_connected(&self, n: u32) { + self.connected.store(n, Ordering::SeqCst); } -} -impl Default for MockNetworkManager { - fn default() -> Self { - Self::new() + /// Deliver an inbound `(peer, message)` to every subscriber interested in + /// its message type, as the real pump would. + pub fn inject(&self, peer: SocketAddr, msg: NetworkMessage) { + let kind = MessageType::from_cmd(msg.cmd()); + let shared = std::sync::Arc::new(msg); + let subs = self.subscribers.lock().expect("mock mutex poisoned"); + for sub in subs.iter() { + if kind.map(|k| sub.kinds.contains(&k)).unwrap_or(false) { + let _ = sub.tx.send((peer, shared.clone())); + } + } + } + + /// Emit a peer-set lifecycle event to all [`NetworkManager::events`] subscribers. + pub fn emit_event(&self, event: NetworkEvent) { + let _ = self.events_tx.send(event); } } #[async_trait] impl NetworkManager for MockNetworkManager { - async fn message_receiver(&mut self, types: &[MessageType]) -> UnboundedReceiver { - self.message_dispatcher.lock().await.message_receiver(types) - } + fn start(&self) {} - fn request_sender(&self) -> RequestSender { - RequestSender::new(self.request_tx.clone()) - } + fn stop(&self) {} - async fn connect(&mut self) -> NetworkResult<()> { - self.connected = true; - Ok(()) + async fn send(&self, msg: NetworkMessage) { + self.sent.lock().expect("mock mutex poisoned").push(msg); } - async fn disconnect(&mut self) -> NetworkResult<()> { - self.connected = false; - Ok(()) + async fn send_to(&self, addr: SocketAddr, msg: NetworkMessage) -> bool { + self.sent_to.lock().expect("mock mutex poisoned").push((addr, msg)); + true } - async fn send_message(&mut self, message: NetworkMessage) -> NetworkResult<()> { - if !self.connected { - return Err(NetworkError::NotConnected); - } - - // Process GetHeaders requests - if let NetworkMessage::GetHeaders(ref getheaders) = message { - let headers = self.process_getheaders(getheaders); - if !headers.is_empty() { - let msg = Message::new(self.connected_peer, NetworkMessage::Headers(headers)); - self.message_dispatcher.lock().await.dispatch(&msg); - } - } - - self.sent_messages.push(message); + fn broadcast(&self, msg: NetworkMessage) { + self.broadcasts.lock().expect("mock mutex poisoned").push(msg); + } - Ok(()) + async fn dispatch_local(&self, msg: NetworkMessage) { + self.inject(test_socket_address(0), msg); } - fn peer_count(&self) -> usize { - if self.connected { - 1 - } else { - 0 - } + + async fn request_answered(&self, key: RequestKey) { + self.answered.lock().expect("mock mutex poisoned").push(key); } - async fn broadcast(&self, _message: NetworkMessage) -> NetworkResult<()> { - panic!("Broadcast not implemented for MockNetworkManager"); + async fn request_completed(&self, peer: SocketAddr, n: usize) { + self.completed.lock().expect("mock mutex poisoned").push((peer, n)); } - async fn dispatch_local(&self, message: NetworkMessage) { - let local_addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0)); - let msg = Message::new(local_addr, message); - self.message_dispatcher.lock().await.dispatch(&msg); + async fn subscribe(&self, kinds: &[MessageType]) -> UnboundedReceiver { + let (tx, rx) = unbounded_channel(); + self.subscribers.lock().expect("mock mutex poisoned").push(Subscriber { + kinds: kinds.to_vec(), + tx, + }); + rx } - async fn disconnect_peer(&self, _addr: &SocketAddr, _reason: &str) -> NetworkResult<()> { - panic!("Disconnect peer not implemented for MockNetworkManager"); + fn events(&self) -> broadcast::Receiver { + self.events_tx.subscribe() } - fn subscribe_network_events(&self) -> broadcast::Receiver { - self.network_event_sender.subscribe() + fn tip(&self) -> u32 { + self.tip.load(Ordering::SeqCst) } -} -impl Peer { - pub fn dummy(addr: SocketAddr) -> Self { - Peer::new(addr, Duration::from_secs(10), Network::Mainnet) + async fn connected_count(&self) -> u32 { + self.connected.load(Ordering::SeqCst) } } diff --git a/dash-spv/tests/dashd_masternode/setup.rs b/dash-spv/tests/dashd_masternode/setup.rs index 773dd5c07..67277838b 100644 --- a/dash-spv/tests/dashd_masternode/setup.rs +++ b/dash-spv/tests/dashd_masternode/setup.rs @@ -165,8 +165,7 @@ pub(super) async fn create_and_start_client( config: &ClientConfig, wallet: Arc>>, ) -> ClientHandle { - let network_manager = - PeerNetworkManager::new(config).await.expect("Failed to create network manager"); + let network_manager = PeerNetworkManager::new(config).await; let storage_manager = DiskStorageManager::new(config).await.expect("Failed to create storage manager"); diff --git a/dash-spv/tests/dashd_sync/helpers.rs b/dash-spv/tests/dashd_sync/helpers.rs index a3be238f3..d46586c01 100644 --- a/dash-spv/tests/dashd_sync/helpers.rs +++ b/dash-spv/tests/dashd_sync/helpers.rs @@ -220,6 +220,36 @@ pub(super) async fn wait_for_mempool_tx( } } +/// Wait for a wallet `TransactionDetected` event for a *specific* txid +pub(super) async fn wait_for_mempool_txid( + receiver: &mut broadcast::Receiver, + expected: Txid, +) -> bool { + let timeout = tokio::time::sleep(Duration::from_secs(30)); + tokio::pin!(timeout); + + loop { + tokio::select! { + _ = &mut timeout => return false, + result = receiver.recv() => { + match result { + Ok(WalletEvent::TransactionDetected { ref record, .. }) + if record.txid == expected + && matches!( + record.context, + TransactionContext::Mempool | TransactionContext::InstantSend(_) + ) => + { + return true; + } + Ok(_) => continue, + Err(_) => return false, + } + } + } + } +} + /// Wait for the mempool manager to reach `Synced` state via the progress watch channel. /// Returns `true` if the state is reached within the timeout, `false` otherwise. pub(super) async fn wait_for_mempool_synced( diff --git a/dash-spv/tests/dashd_sync/setup.rs b/dash-spv/tests/dashd_sync/setup.rs index 27c8b5415..2089063e4 100644 --- a/dash-spv/tests/dashd_sync/setup.rs +++ b/dash-spv/tests/dashd_sync/setup.rs @@ -1,6 +1,5 @@ use dash_spv::client::config::MempoolStrategy; use dash_spv::network::NetworkEvent; -use dash_spv::storage::{PeerStorage, PersistentPeerStorage, PersistentStorage}; use dash_spv::test_utils::{ create_test_wallet, init_test_logging, next_unused_receive_address, retain_test_dir, DashdTestContext, TestChain, TestEventHandler, @@ -12,8 +11,6 @@ use dash_spv::{ sync::{ProgressPercentage, SyncEvent, SyncProgress}, LoggingGuard, Network, }; -use dashcore::network::address::AddrV2Message; -use dashcore::network::constants::ServiceFlags; use dashcore::Txid; use key_wallet::wallet::managed_wallet_info::wallet_info_interface::WalletInfoInterface; use key_wallet::wallet::managed_wallet_info::ManagedWalletInfo; @@ -281,8 +278,7 @@ pub(super) async fn create_and_start_client( config: &ClientConfig, wallet: Arc>>, ) -> ClientHandle { - let network_manager = - PeerNetworkManager::new(config).await.expect("Failed to create network manager"); + let network_manager = PeerNetworkManager::new(config).await; let storage_manager = DiskStorageManager::new(config).await.expect("Failed to create storage manager"); @@ -329,12 +325,10 @@ pub(super) async fn create_non_exclusive_test_config( storage_path: PathBuf, peer_addr: std::net::SocketAddr, ) -> ClientConfig { - let config = ClientConfig::regtest().with_storage_path(storage_path).without_masternodes(); - // Seed the peer store so the client can discover our dashd node - let peer_store = PersistentPeerStorage::open(config.storage_path.clone()) - .await - .expect("Failed to open peer storage"); - let msg = AddrV2Message::new(peer_addr, ServiceFlags::NETWORK); - peer_store.save_peers(&[msg]).await.expect("Failed to seed peer store"); + let mut config = ClientConfig::regtest().with_storage_path(storage_path).without_masternodes(); + // Non-exclusive discovery: add the node as a configured peer while leaving + // `restrict_to_configured_peers` false. (Peer storage was removed, so this + // replaces seeding it on disk.) + config.add_peer(peer_addr); config } diff --git a/dash-spv/tests/dashd_sync/tests_mempool.rs b/dash-spv/tests/dashd_sync/tests_mempool.rs index becb4598f..34103b687 100644 --- a/dash-spv/tests/dashd_sync/tests_mempool.rs +++ b/dash-spv/tests/dashd_sync/tests_mempool.rs @@ -410,12 +410,12 @@ async fn test_mempool_peer_disconnect_reactivation() { let (fa_disc, bf_disc) = tokio::join!( wait_for_network_event( &mut fa_net_rx, - |e| matches!(e, NetworkEvent::PeerDisconnected { address } if *address == ctx.dashd.addr), + |e| matches!(e, NetworkEvent::PeerDisconnected(address) if *address == ctx.dashd.addr), Duration::from_secs(10), ), wait_for_network_event( &mut bf_net_rx, - |e| matches!(e, NetworkEvent::PeerDisconnected { address } if *address == ctx.dashd.addr), + |e| matches!(e, NetworkEvent::PeerDisconnected(address) if *address == ctx.dashd.addr), Duration::from_secs(10), ), ); @@ -449,10 +449,10 @@ async fn test_mempool_peer_disconnect_reactivation() { _ = &mut deadline => panic!("{}: timed out waiting for both peer disconnects", label), result = receiver.recv() => { match result { - Ok(NetworkEvent::PeerDisconnected { address }) if address == ctx.dashd.addr => { + Ok(NetworkEvent::PeerDisconnected(address)) if address == ctx.dashd.addr => { seen_dashd1 = true; } - Ok(NetworkEvent::PeerDisconnected { address }) if address == dashd2.addr => { + Ok(NetworkEvent::PeerDisconnected(address)) if address == dashd2.addr => { seen_dashd2 = true; } _ => {} diff --git a/dash-spv/tests/dashd_sync/tests_restart.rs b/dash-spv/tests/dashd_sync/tests_restart.rs index 89bce2421..5a939400a 100644 --- a/dash-spv/tests/dashd_sync/tests_restart.rs +++ b/dash-spv/tests/dashd_sync/tests_restart.rs @@ -113,20 +113,32 @@ async fn test_sync_restart_with_fresh_wallet() { /// Verify sync completes successfully despite repeated interruptions. /// /// Listens for key sync events (BlockHeadersStored, FilterHeadersStored, FiltersStored, -/// BlocksNeeded, BlockProcessed) and restarts the client on every 2nd occurrence until -/// sync completes. This exercises restart/resume from unpredictable points across the -/// full sync lifecycle. +/// BlocksNeeded, BlockProcessed) and restarts the client on every 2nd occurrence to +/// exercise restart/resume from unpredictable points across the full sync lifecycle. +/// After a bounded number of restarts it lets the final client run through to +/// completion, so the assertion confirms sync actually finishes rather than merely +/// surviving each interruption. #[tokio::test] async fn test_sync_with_multiple_restarts() { let Some(ctx) = TestContext::new(TestChain::Full).await else { return; }; + // Number of event-driven restarts to exercise before letting the client finish. + // Bounded so the test stays deterministic and fast: restarting on every 2nd + // progress event until natural completion would take hundreds of restart cycles + // on the full chain. + const MAX_RESTARTS: usize = 8; + let mut restart_count = 0; let final_progress = loop { tracing::info!("Starting sync (restart count: {})", restart_count); let mut client_handle = ctx.spawn_new_client().await; + // Once we've exercised enough restart points, let this client sync to + // completion instead of interrupting it again. + let final_run = restart_count >= MAX_RESTARTS; + // Wait for either sync completion or the 2nd matching event let mut events_seen = 0; let mut should_restart = false; @@ -146,7 +158,7 @@ async fn test_sync_with_multiple_restarts() { match result { Ok(ref event) if is_progress_event(event) => { events_seen += 1; - if events_seen % 2 == 0 { + if !final_run && events_seen % 2 == 0 { tracing::info!("Restarting on: {}", event); should_restart = true; break; diff --git a/dash-spv/tests/dashd_sync/tests_transaction.rs b/dash-spv/tests/dashd_sync/tests_transaction.rs index 0dfb5d7aa..95b3d02bd 100644 --- a/dash-spv/tests/dashd_sync/tests_transaction.rs +++ b/dash-spv/tests/dashd_sync/tests_transaction.rs @@ -6,8 +6,8 @@ use std::time::Duration; use tokio::sync::RwLock; use super::helpers::{ - count_wallet_transactions, get_spendable_balance, wait_for_mempool_tx, wait_for_sync, - wait_for_wallet_synced, EMPTY_MNEMONIC, SECONDARY_MNEMONIC, + count_wallet_transactions, get_spendable_balance, wait_for_mempool_tx, wait_for_mempool_txid, + wait_for_sync, wait_for_wallet_synced, EMPTY_MNEMONIC, SECONDARY_MNEMONIC, }; use super::setup::{create_and_start_client, TestContext}; use dash_spv::test_utils::{create_test_wallet, TestChain}; @@ -407,9 +407,10 @@ async fn test_spend_change_balance() { build_and_sign(&wallet, &wallet_id, &dest_a, 100_000_000).await.expect("build tx_a"); client_handle.client.broadcast_transaction(&tx_a).await.expect("broadcast tx_a"); - wait_for_mempool_tx(&mut client_handle.wallet_event_receiver, MEMPOOL_TIMEOUT) - .await - .expect("detect tx_a"); + assert!( + wait_for_mempool_txid(&mut client_handle.wallet_event_receiver, tx_a.txid()).await, + "tx_a not detected in mempool", + ); // The wallet's only UTXO now is the mempool change from tx_a, so a // successful build proves coin selection used it. @@ -423,9 +424,10 @@ async fn test_spend_change_balance() { ); client_handle.client.broadcast_transaction(&tx_b).await.expect("broadcast tx_b"); - wait_for_mempool_tx(&mut client_handle.wallet_event_receiver, MEMPOOL_TIMEOUT) - .await - .expect("detect tx_b"); + assert!( + wait_for_mempool_txid(&mut client_handle.wallet_event_receiver, tx_b.txid()).await, + "tx_b not detected in mempool", + ); client_handle.stop().await; } @@ -483,15 +485,15 @@ async fn test_concurrent_builds_do_not_double_spend() { // Both broadcasts succeed and both transactions reach the mempool: a // double-spend would have the second rejected by the network. client_handle.client.broadcast_transaction(&tx_a).await.expect("broadcast tx_a"); - let detected_a = wait_for_mempool_tx(&mut client_handle.wallet_event_receiver, MEMPOOL_TIMEOUT) - .await - .expect("detect tx_a"); - assert_eq!(detected_a, tx_a.txid()); + assert!( + wait_for_mempool_txid(&mut client_handle.wallet_event_receiver, tx_a.txid()).await, + "tx_a not detected in mempool", + ); client_handle.client.broadcast_transaction(&tx_b).await.expect("broadcast tx_b"); - let detected_b = wait_for_mempool_tx(&mut client_handle.wallet_event_receiver, MEMPOOL_TIMEOUT) - .await - .expect("detect tx_b"); - assert_eq!(detected_b, tx_b.txid()); + assert!( + wait_for_mempool_txid(&mut client_handle.wallet_event_receiver, tx_b.txid()).await, + "tx_b not detected in mempool", + ); client_handle.stop().await; } diff --git a/dash-spv/tests/peer_test.rs b/dash-spv/tests/peer_test.rs deleted file mode 100644 index 8634d4413..000000000 --- a/dash-spv/tests/peer_test.rs +++ /dev/null @@ -1,232 +0,0 @@ -//! Integration tests for peer networking - -use std::net::SocketAddr; -use std::sync::Arc; -use std::time::Duration; -use tempfile::TempDir; -use tokio::sync::RwLock; -use tokio::time; - -use dash_spv::client::{ClientConfig, DashSpvClient}; -use dash_spv::network::PeerNetworkManager; -use dash_spv::storage::DiskStorageManager; -use dash_spv::types::ValidationMode; -use dashcore::Network; -use key_wallet::wallet::managed_wallet_info::ManagedWalletInfo; -use key_wallet_manager::WalletManager; - -fn init_test_tracing() { - let _ = tracing_subscriber::fmt().with_test_writer().try_init(); -} - -/// Create a test configuration with the given network -fn create_test_config(network: Network) -> ClientConfig { - let mut config = ClientConfig::new(network); - - config.storage_path = TempDir::new().unwrap().path().to_path_buf(); - - config.validation_mode = ValidationMode::Basic; - config.enable_filters = false; - config.enable_masternodes = false; - config.max_peers = 3; - config.peers = vec![]; // Will be populated by DNS discovery - config -} - -#[tokio::test] -#[ignore] // Requires network access -async fn test_peer_connection() { - init_test_tracing(); - - let config = create_test_config(Network::Testnet); - - // Create network manager - let network_manager = PeerNetworkManager::new(&config).await.unwrap(); - - // Create storage manager - let storage_manager = DiskStorageManager::new(&config).await.unwrap(); - - // Create wallet manager - let wallet = Arc::new(RwLock::new(WalletManager::::new(config.network))); - - let client = - DashSpvClient::new(config, network_manager, storage_manager, wallet, vec![]).await.unwrap(); - - let run_client = client.clone(); - let handle = tokio::spawn(async move { run_client.run().await }); - - // Give it time to connect to peers - time::sleep(Duration::from_secs(5)).await; - - // Check that we have connected to at least one peer - let peer_count = client.peer_count().await; - assert!(peer_count > 0, "Should have connected to at least one peer"); - - client.stop().await.expect("Should stop"); - let _ = handle.await; -} - -#[tokio::test] -#[ignore] // Requires network access -async fn test_peer_persistence() { - init_test_tracing(); - - let config = create_test_config(Network::Testnet); - - // First run: connect and save peers - { - // Create network manager - let network_manager = PeerNetworkManager::new(&config).await.unwrap(); - - // Create storage manager - let storage_manager = DiskStorageManager::new(&config).await.unwrap(); - - // Create wallet manager - let wallet = Arc::new(RwLock::new(WalletManager::::new(config.network))); - - let client = - DashSpvClient::new(config.clone(), network_manager, storage_manager, wallet, vec![]) - .await - .unwrap(); - - let run_client = client.clone(); - let handle = tokio::spawn(async move { run_client.run().await }); - - time::sleep(Duration::from_secs(5)).await; - - let peer_count = client.peer_count().await; - assert!(peer_count > 0, "Should have connected to peers"); - - client.stop().await.expect("Should stop"); - let _ = handle.await; - } - - // Second run: should load saved peers - { - // Create network manager - let network_manager = PeerNetworkManager::new(&config).await.unwrap(); - - // Create storage manager - reuse same path - let storage_manager = DiskStorageManager::new(&config).await.unwrap(); - - // Create wallet manager - let wallet = Arc::new(RwLock::new(WalletManager::::new(config.network))); - - let client = DashSpvClient::new(config, network_manager, storage_manager, wallet, vec![]) - .await - .unwrap(); - - // Should connect faster due to saved peers - let run_client = client.clone(); - let start = tokio::time::Instant::now(); - let handle = tokio::spawn(async move { run_client.run().await }); - - // Wait for connection but with shorter timeout - time::sleep(Duration::from_secs(3)).await; - - let peer_count = client.peer_count().await; - assert!(peer_count > 0, "Should have connected using saved peers"); - - let elapsed = start.elapsed(); - println!("Connected to {} peers in {:?} (using saved peers)", peer_count, elapsed); - - client.stop().await.expect("Should stop"); - let _ = handle.await; - } -} - -#[tokio::test] -async fn test_peer_disconnection() { - init_test_tracing(); - - let mut config = create_test_config(Network::Regtest); - - // Add manual test peers (would need actual regtest nodes running) - config.peers = vec!["127.0.0.1:19899".parse().unwrap(), "127.0.0.1:19898".parse().unwrap()]; - - // Create network manager - let network_manager = PeerNetworkManager::new(&config).await.unwrap(); - - // Create storage manager - let storage_manager = DiskStorageManager::new(&config).await.unwrap(); - - // Create wallet manager - let wallet = Arc::new(RwLock::new(WalletManager::::new(config.network))); - - let client = - DashSpvClient::new(config, network_manager, storage_manager, wallet, vec![]).await.unwrap(); - - // Note: This test would require actual regtest nodes running - // For now, we just test that the API works - let test_addr: SocketAddr = "127.0.0.1:19899".parse().unwrap(); - - // Try to disconnect (will fail if not connected, but tests the API) - match client.disconnect_peer(&test_addr, "Test disconnection").await { - Ok(_) => println!("Disconnected peer {}", test_addr), - Err(e) => println!("Expected error disconnecting non-existent peer: {}", e), - } -} - -#[cfg(test)] -mod unit_tests { - use super::*; - use dash_spv::network::addrv2::AddrV2Handler; - use dash_spv::network::discovery::DnsDiscovery; - use dash_spv::network::pool::PeerPool; - use dashcore::network::constants::ServiceFlags; - - #[tokio::test] - async fn test_connection_pool_limits() { - let pool = PeerPool::new(8); - - // Should start empty - assert_eq!(pool.peer_count().await, 0); - assert!(pool.needs_more_peers().await); - assert!(pool.can_accept_peers().await); - - // Test marking as connecting - let addr1: SocketAddr = "127.0.0.1:9999".parse().unwrap(); - assert!(pool.mark_connecting(addr1).await); - assert!(!pool.mark_connecting(addr1).await); // Already marked - assert!(pool.is_connecting(&addr1).await); - } - - #[tokio::test] - async fn test_addrv2_handler() { - let handler = AddrV2Handler::new(); - - // Test tracking AddrV2 support - let peer: SocketAddr = "192.168.1.1:9999".parse().unwrap(); - handler.handle_sendaddrv2(peer).await; - assert!(handler.peer_supports_addrv2(&peer).await); - - // Test adding addresses - handler.add_known_address(peer, ServiceFlags::NETWORK).await; - let known = handler.get_known_addresses().await; - assert_eq!(known.len(), 1); - assert_eq!(known[0].socket_addr().unwrap(), peer); - - // Test getting addresses for sharing - let to_share = handler.get_addresses_for_peer(10).await; - assert_eq!(to_share.len(), 1); - } - - #[tokio::test] - #[ignore] // Requires network access - async fn test_dns_discovery() { - let discovery = DnsDiscovery::new(); - - // Test mainnet discovery - let peers = discovery.discover_peers(Network::Mainnet).await; - assert!(!peers.is_empty(), "Should discover mainnet peers"); - - // All peers should use correct port - for peer in &peers { - assert_eq!(peer.port(), 9999); - } - - // Test limited discovery - let limited = discovery.discover_peers_limited(Network::Mainnet, 5).await; - assert!(limited.len() <= 5); - } -} diff --git a/dash-spv/tests/test_handshake_logic.rs b/dash-spv/tests/test_handshake_logic.rs deleted file mode 100644 index d8ebdb6c9..000000000 --- a/dash-spv/tests/test_handshake_logic.rs +++ /dev/null @@ -1,16 +0,0 @@ -//! Unit tests for handshake logic - -use dash_spv::network::{HandshakeManager, HandshakeState}; -use dashcore::Network; - -#[test] -fn test_handshake_state_transitions() { - let mut handshake = HandshakeManager::new(Network::Mainnet, None); - - // Initial state should be Init - assert_eq!(*handshake.state(), HandshakeState::Init); - - // After reset, should be back to Init - handshake.reset(); - assert_eq!(*handshake.state(), HandshakeState::Init); -} diff --git a/dash-spv/tests/wallet_integration_test.rs b/dash-spv/tests/wallet_integration_test.rs index ef374ebfc..faba1be7c 100644 --- a/dash-spv/tests/wallet_integration_test.rs +++ b/dash-spv/tests/wallet_integration_test.rs @@ -24,7 +24,7 @@ async fn create_test_client( .with_restrict_to_configured_peers(true); // Create network manager - let network_manager = PeerNetworkManager::new(&config).await.unwrap(); + let network_manager = PeerNetworkManager::new(&config).await; // Create storage manager let storage_manager = DiskStorageManager::new(&config).await.expect("Failed to create storage"); diff --git a/dash/Cargo.toml b/dash/Cargo.toml index 22d697b9b..8362d666a 100644 --- a/dash/Cargo.toml +++ b/dash/Cargo.toml @@ -34,6 +34,7 @@ quorum_validation = ["bls"] message_verification = ["bls"] bincode = [ "dep:bincode", "dep:bincode_derive", "dashcore_hashes/bincode", "dash-network/bincode" ] test-utils = [] +tokio = ["dep:tokio-util", "dep:bytes"] [package.metadata.docs.rs] all-features = true @@ -57,6 +58,8 @@ bincode = { version = "2.0.1", optional = true } bincode_derive = { version = "2.0.1", optional = true } blsful = { git = "https://github.com/dashpay/agora-blsful", rev = "0c34a7a488a0bd1c9a9a2196e793b303ad35c900", optional = true } ed25519-dalek = { version = "2.1", features = ["rand_core"], optional = true } +tokio-util = { version = "0.7", default-features = false, features = ["codec"], optional = true } +bytes = { version = "1", optional = true } blake3 = "1.8.1" thiserror = "2" bitvec = "1.0" diff --git a/dash/src/network/message.rs b/dash/src/network/message.rs index 140d32d16..d64aaad53 100644 --- a/dash/src/network/message.rs +++ b/dash/src/network/message.rs @@ -666,6 +666,138 @@ impl Decodable for RawNetworkMessage { } } +/// A Tokio [`Decoder`](tokio_util::codec::Decoder)/[`Encoder`](tokio_util::codec::Encoder) +/// that frames one [`RawNetworkMessage`] per message, so a peer connection's read half can +/// be wrapped in `tokio_util::codec::FramedRead` (and the write half in `FramedWrite`) +/// instead of hand-rolling the length-prefixed framing. Enabled by the `tokio` feature. +/// The [`RawNetworkMessage`] decoder validates the payload checksum, so a corrupt +/// frame surfaces as an error rather than a bad message. +#[cfg(feature = "tokio")] +#[derive(Debug, Default, Clone, Copy)] +pub struct RawNetworkMessageCodec; + +#[cfg(feature = "tokio")] +impl tokio_util::codec::Decoder for RawNetworkMessageCodec { + type Item = RawNetworkMessage; + type Error = encode::Error; + + fn decode(&mut self, src: &mut bytes::BytesMut) -> Result, Self::Error> { + use bytes::Buf as _; + + /// Message header: `magic(4) + command(12) + payload_length(4) + checksum(4)`. + const HEADER_LEN: usize = 24; + /// Byte offset of the little-endian payload length within the header. + const LENGTH_OFFSET: usize = 16; + + // Need the whole header before we can read the declared payload length. + if src.len() < HEADER_LEN { + return Ok(None); + } + + let payload_len = u32::from_le_bytes([ + src[LENGTH_OFFSET], + src[LENGTH_OFFSET + 1], + src[LENGTH_OFFSET + 2], + src[LENGTH_OFFSET + 3], + ]) as usize; + + // Reject an absurd declared length *before* reserving buffer for it. + if payload_len > MAX_MSG_SIZE { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "declared payload length exceeds MAX_MSG_SIZE", + ) + .into()); + } + + let frame_len = HEADER_LEN + payload_len; + if src.len() < frame_len { + src.reserve(frame_len - src.len()); + return Ok(None); + } + + // Decode from the exact frame bytes (a finite, in-memory reader). + let mut cursor = io::Cursor::new(&src[..frame_len]); + let msg = RawNetworkMessage::consensus_decode_from_finite_reader(&mut cursor)?; + src.advance(frame_len); + + Ok(Some(msg)) + } +} + +#[cfg(feature = "tokio")] +impl tokio_util::codec::Encoder<&RawNetworkMessage> for RawNetworkMessageCodec { + type Error = encode::Error; + + fn encode( + &mut self, + item: &RawNetworkMessage, + dst: &mut bytes::BytesMut, + ) -> Result<(), Self::Error> { + dst.extend_from_slice(&serialize(item)); + Ok(()) + } +} + +#[cfg(feature = "tokio")] +impl tokio_util::codec::Encoder for RawNetworkMessageCodec { + type Error = encode::Error; + + fn encode( + &mut self, + item: RawNetworkMessage, + dst: &mut bytes::BytesMut, + ) -> Result<(), Self::Error> { + >::encode(self, &item, dst) + } +} + +#[cfg(all(test, feature = "tokio"))] +mod codec_tests { + use super::{NetworkMessage, RawNetworkMessage, RawNetworkMessageCodec}; + use bytes::BytesMut; + use tokio_util::codec::{Decoder, Encoder}; + + #[test] + fn roundtrip_frame() { + let msg = RawNetworkMessage { + magic: 0xBD6B_0CBF, + payload: NetworkMessage::Ping(0x0102_0304_0506_0708), + }; + let mut codec = RawNetworkMessageCodec; + let mut buf = BytesMut::new(); + codec.encode(&msg, &mut buf).unwrap(); + + let decoded = codec.decode(&mut buf).unwrap().expect("full frame decodes"); + assert_eq!(decoded.magic, msg.magic); + assert!(matches!(decoded.payload, NetworkMessage::Ping(0x0102_0304_0506_0708))); + assert!(buf.is_empty()); + } + + #[test] + fn partial_frame_yields_none_until_complete() { + let msg = RawNetworkMessage { + magic: 0xBD6B_0CBF, + payload: NetworkMessage::Ping(42), + }; + let mut codec = RawNetworkMessageCodec; + let mut full = BytesMut::new(); + codec.encode(&msg, &mut full).unwrap(); + + let mut partial = BytesMut::new(); + for (i, byte) in full.iter().enumerate() { + partial.extend_from_slice(&[*byte]); + let out = codec.decode(&mut partial).unwrap(); + if i + 1 < full.len() { + assert!(out.is_none()); + } else { + assert!(out.is_some()); + assert!(partial.is_empty()); + } + } + } +} + #[cfg(test)] mod test { use std::net::Ipv4Addr; diff --git a/masternode-seeds-fetcher/Cargo.toml b/masternode-seeds-fetcher/Cargo.toml index 012dcb74a..981f94431 100644 --- a/masternode-seeds-fetcher/Cargo.toml +++ b/masternode-seeds-fetcher/Cargo.toml @@ -13,11 +13,12 @@ name = "masternode-seeds-fetcher" path = "src/main.rs" [dependencies] -dashcore = { path = "../dash", features = ["core-block-hash-use-x11"] } +dashcore = { path = "../dash", features = ["core-block-hash-use-x11", "tokio"] } dash-spv = { path = "../dash-spv" } dash-network-seeds = { path = "../dash-network-seeds" } tokio = { version = "1.0", features = ["full"] } +tokio-util = "0.7" clap = { version = "4.0", features = ["derive", "env"] } anyhow = "1.0" diff --git a/masternode-seeds-fetcher/src/main.rs b/masternode-seeds-fetcher/src/main.rs index 6b1f68613..d7379dc68 100644 --- a/masternode-seeds-fetcher/src/main.rs +++ b/masternode-seeds-fetcher/src/main.rs @@ -33,12 +33,12 @@ use std::net::{IpAddr, SocketAddr}; use std::path::PathBuf; use std::time::Duration; +use crate::peer::Peer; use anyhow::{Context, Result, anyhow}; use clap::Parser; use dash_network_seeds::{ CoreStatus, MasternodeSeed, MasternodeType, PlatformStatus, Reachability, }; -use dash_spv::network::Peer; use dashcore::hashes::Hash; use dashcore::network::Address; use dashcore::network::constants::ServiceFlags; @@ -53,6 +53,7 @@ use std::sync::Arc; use tokio::sync::Semaphore; use tokio::time::Instant; +mod peer; mod probe; // ---------- CLI ---------- @@ -256,7 +257,7 @@ async fn fetch_from_peer(peer_addr: SocketAddr, network: Network) -> Result Some(d.clone()), + NetworkMessage::MnListDiff(d) => Some((*d).clone()), _ => None, }) .await diff --git a/masternode-seeds-fetcher/src/peer.rs b/masternode-seeds-fetcher/src/peer.rs new file mode 100644 index 000000000..5d6ee16f6 --- /dev/null +++ b/masternode-seeds-fetcher/src/peer.rs @@ -0,0 +1,74 @@ +//! Minimal Dash P2P connection for probing a single peer. +//! +//! The SPV client's network module drives its peers from background tasks and hands +//! messages to subscribers — the right shape for a sync, the wrong one for a probe that +//! wants to send a request and block on the reply. So this tool keeps its own socket +//! rather than depending on the client's internals. + +use std::net::SocketAddr; +use std::time::Duration; + +use anyhow::{Context, Result, anyhow}; +use dashcore::Network; +use dashcore::consensus::encode; +use dashcore::network::message::{NetworkMessage, RawNetworkMessage, RawNetworkMessageCodec}; +use futures::StreamExt; +use tokio::io::AsyncWriteExt; +use tokio::net::TcpStream; +use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf}; +use tokio_util::codec::FramedRead; + +/// A message received from a peer. +pub struct Message(NetworkMessage); + +impl Message { + pub fn inner(&self) -> &NetworkMessage { + &self.0 + } +} + +/// A connected peer: a framed reader and a raw writer over one TCP socket. +pub struct Peer { + reader: FramedRead, + writer: OwnedWriteHalf, + magic: u32, +} + +impl Peer { + /// Open a TCP connection to `addr`. No handshake — the caller drives it. + pub async fn connect(addr: SocketAddr, timeout_secs: u64, network: Network) -> Result { + let stream = + tokio::time::timeout(Duration::from_secs(timeout_secs), TcpStream::connect(addr)) + .await + .map_err(|_| anyhow!("connect to {addr} timed out after {timeout_secs}s"))? + .with_context(|| format!("connect to {addr}"))?; + + let (read_half, writer) = stream.into_split(); + + Ok(Self { + reader: FramedRead::new(read_half, RawNetworkMessageCodec), + writer, + magic: network.magic(), + }) + } + + pub async fn send_message(&mut self, message: NetworkMessage) -> Result<()> { + let raw = RawNetworkMessage { + magic: self.magic, + payload: message, + }; + + self.writer.write_all(&encode::serialize(&raw)).await.context("send message")?; + + Ok(()) + } + + /// Next message from the peer, or `None` once it closes the connection. + pub async fn receive_message(&mut self) -> Result> { + match self.reader.next().await { + Some(Ok(raw)) => Ok(Some(Message(raw.payload))), + Some(Err(e)) => Err(e).context("decode message"), + None => Ok(None), + } + } +} diff --git a/masternode-seeds-fetcher/src/probe.rs b/masternode-seeds-fetcher/src/probe.rs index 49a1de9cc..3020cbe45 100644 --- a/masternode-seeds-fetcher/src/probe.rs +++ b/masternode-seeds-fetcher/src/probe.rs @@ -27,7 +27,7 @@ use tokio::net::TcpStream; use x509_parser::prelude::FromDer; use x509_parser::x509::X509Version; -use dash_spv::network::Peer; +use crate::peer::Peer; /// How long to give a single TCP connect before giving up. const CONNECT_TIMEOUT: Duration = Duration::from_secs(5);