From fdf474f7816bb6643be00162193fdec9dae873f7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=A1=A5=E4=B8=8B=E7=BA=A2=E8=8D=AF?= Date: Sat, 25 Jul 2026 01:53:04 +0800 Subject: [PATCH] feat: improve metadata peer discovery --- README.md | 5 +- src/crawl_engine.rs | 1 + src/krpc.rs | 86 +++++- src/lib.rs | 9 +- src/metadata.rs | 15 + src/peer_lookup.rs | 518 ++++++++++++++++++++++++++++++++++ src/protocol.rs | 3 + src/routing_snapshot.rs | 70 +++++ src/runtime_stats.rs | 195 +++++++++++++ src/scheduler.rs | 598 +++++++++++++++++++++++++++++++++------- src/server.rs | 39 ++- src/types.rs | 24 ++ 12 files changed, 1446 insertions(+), 117 deletions(-) create mode 100644 src/peer_lookup.rs diff --git a/README.md b/README.md index 5fc5d45..5cdafd7 100644 --- a/README.md +++ b/README.md @@ -177,6 +177,9 @@ hash ingress → deduplicating Metadata queue → bounded workers → torrent ca | `metadata.max_worker_count` | `256` | 最大并发 Metadata job 数 | | `metadata.peer_failure_cache_capacity` | `200000` | 坏 Peer 缓存容量 | | `metadata.peer_failure_ttl_secs` | `60` | 坏 Peer 缓存 TTL | +| `peer_lookup.max_lookups_per_second` | `32` | 每秒启动的主动 `get_peers` 查询数；`0` 为关闭 | +| `peer_lookup.burst` | `32` | 空闲后可立即消费的查询预算 | +| `peer_lookup.max_active_lookups` | `64` | 同时活跃的 InfoHash 查询上限 | ### `crawl.rate_limit` @@ -328,7 +331,7 @@ PrometheusBuilder::new() ## 应用层集成 本 crate 只提供 DHT、BEP-9 Metadata、回调和观测能力，不包含 Redis、Manticore、 -HTTP 看板或 `P1`～`P16` 性能档位。同级的 `dht-crawler-node` 项目负责这些应用层策略， +HTTP 看板或 `P1`～`P17` 性能档位。同级的 `dht-crawler-node` 项目负责这些应用层策略， 并把档位转换成具体的 `DHTOptions`。开发两个项目时应保持下面的目录关系： ```text diff --git a/src/crawl_engine.rs b/src/crawl_engine.rs index 747f7dd..9f59e69 100644 --- a/src/crawl_engine.rs +++ b/src/crawl_engine.rs @@ -913,6 +913,7 @@ mod tests { id: None, nodes: None, nodes6: None, + values: None, }; engine.route_response("8.8.8.8:1".parse().unwrap(), [1; 8], response()); engine.route_response("1.1.1.1:2".parse().unwrap(), [2; 8], response()); diff --git a/src/krpc.rs b/src/krpc.rs index b51674c..6e43f6e 100644 --- a/src/krpc.rs +++ b/src/krpc.rs @@ -1,4 +1,4 @@ -use crate::addr::is_valid_node_addr; +use crate::addr::{addr_allowed_by_netmode, is_valid_node_addr}; use crate::node_id::TransactionId; use crate::protocol::DhtResponse; use crate::types::{NetMode, NodeTuple}; @@ -50,6 +50,40 @@ pub(crate) fn for_each_response_node( count } +pub(crate) fn for_each_response_peer( + response: &DhtResponse, + netmode: NetMode, + mut visit: impl FnMut(SocketAddr), +) -> usize { + let mut count = 0; + let Some(values) = response.values.as_ref() else { + return count; + }; + for value in values { + let bytes = value.as_ref(); + let addr = match bytes.len() { + 6 => { + let ip = Ipv4Addr::new(bytes[0], bytes[1], bytes[2], bytes[3]); + let port = u16::from_be_bytes([bytes[4], bytes[5]]); + SocketAddr::new(IpAddr::V4(ip), port) + } + 18 => { + let ip_bytes: [u8; 16] = bytes[..16] + .try_into() + .expect("compact IPv6 Peer address is 16 bytes"); + let port = u16::from_be_bytes([bytes[16], bytes[17]]); + SocketAddr::new(IpAddr::V6(Ipv6Addr::from(ip_bytes)), port) + } + _ => continue, + }; + if addr_allowed_by_netmode(&addr, netmode) && is_valid_node_addr(&addr) { + visit(addr); + count += 1; + } + } + count +} + pub(crate) fn encode_find_node_query( buffer: &mut BytesMut, tid: &TransactionId, @@ -67,6 +101,23 @@ pub(crate) fn encode_find_node_query( buffer.extend_from_slice(b"1:y1:qe"); } +pub(crate) fn encode_get_peers_query( + buffer: &mut BytesMut, + tid: &TransactionId, + info_hash: &[u8; 20], + sender_id: &[u8; 20], +) { + buffer.clear(); + buffer.reserve(111); + buffer.extend_from_slice(b"d1:ad2:id20:"); + buffer.extend_from_slice(sender_id); + buffer.extend_from_slice(b"9:info_hash20:"); + buffer.extend_from_slice(info_hash); + buffer.extend_from_slice(b"e1:q9:get_peers1:t8:"); + buffer.extend_from_slice(tid); + buffer.extend_from_slice(b"1:y1:qe"); +} + pub(crate) fn encode_response( buffer: &mut BytesMut, tid: &[u8], @@ -176,6 +227,19 @@ mod tests { assert_eq!(message.a.unwrap().target.unwrap().as_ref(), &target); } + #[test] + fn manual_get_peers_encoding_round_trips() { + let mut buffer = BytesMut::new(); + let tid = [1; 8]; + let info_hash = [2; 20]; + let sender = [3; 20]; + encode_get_peers_query(&mut buffer, &tid, &info_hash, &sender); + let message: DhtMessage = serde_bencode::from_bytes(&buffer).unwrap(); + assert_eq!(message.t.as_ref(), &tid); + assert_eq!(message.q.as_deref(), Some("get_peers")); + assert_eq!(message.a.unwrap().info_hash.unwrap().as_ref(), &info_hash); + } + #[test] fn manual_response_encoding_round_trips() { let mut buffer = BytesMut::new(); @@ -188,4 +252,24 @@ mod tests { let response = message.r.unwrap(); assert_eq!(response.nodes.unwrap().len(), 26); } + + #[test] + fn compact_get_peers_values_are_validated() { + let response = DhtResponse { + id: None, + nodes: None, + nodes6: None, + values: Some(vec![ + serde_bytes::ByteBuf::from(vec![8, 8, 8, 8, 0x1a, 0xe1]), + serde_bytes::ByteBuf::from(vec![10, 0, 0, 1, 0x1a, 0xe1]), + serde_bytes::ByteBuf::from(vec![1, 2, 3]), + ]), + }; + let mut peers = Vec::new(); + assert_eq!( + for_each_response_peer(&response, NetMode::Ipv4Only, |peer| peers.push(peer)), + 1 + ); + assert_eq!(peers[0], "8.8.8.8:6881".parse().unwrap()); + } } diff --git a/src/lib.rs b/src/lib.rs index d78c863..fb1fda5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -16,6 +16,7 @@ mod krpc; pub mod metadata; mod node_id; mod node_pool; +mod peer_lookup; /// Serializable BEP-5 KRPC wire types. pub mod protocol; mod routing_snapshot; @@ -36,8 +37,8 @@ pub use scheduler::{MetadataScheduler, MetadataSchedulerCallbacks, MetadataSched pub use server::{DHTServer, HashDiscovered}; pub use types::{ BootstrapOptions, CrawlOptions, DHTOptions, FileInfo, MetadataFetchCompletion, - MetadataFetchCompletionStatus, MetadataOptions, NetMode, NodeTuple, PoolOptions, - RateLimitOptions, SchedulerOptions, TargetOptions, TorrentInfo, + MetadataFetchCompletionStatus, MetadataOptions, NetMode, NodeTuple, PeerLookupOptions, + PoolOptions, RateLimitOptions, SchedulerOptions, TargetOptions, TorrentInfo, }; /// Common server, configuration and callback payload imports. @@ -50,8 +51,8 @@ pub mod prelude { pub use crate::server::DHTServer; pub use crate::types::{ BootstrapOptions, CrawlOptions, DHTOptions, FileInfo, MetadataFetchCompletion, - MetadataFetchCompletionStatus, MetadataOptions, NetMode, NodeTuple, PoolOptions, - RateLimitOptions, SchedulerOptions, TargetOptions, TorrentInfo, + MetadataFetchCompletionStatus, MetadataOptions, NetMode, NodeTuple, PeerLookupOptions, + PoolOptions, RateLimitOptions, SchedulerOptions, TargetOptions, TorrentInfo, }; } diff --git a/src/metadata.rs b/src/metadata.rs index 54ee385..517ddc6 100644 --- a/src/metadata.rs +++ b/src/metadata.rs @@ -194,11 +194,25 @@ impl RbitFetcher { /// The deadline covers TCP connect, both BitTorrent handshakes, all metadata /// piece I/O, hash validation and bencode parsing. Inner library timeouts can /// therefore never stack on top of the configured metadata timeout. + #[cfg(test)] pub(crate) async fn fetch( &self, info_hash: &[u8; 20], peer_addr: SocketAddr, ) -> MetadataFetchOutcome { + self.fetch_with_attempt_observer(info_hash, peer_addr, || {}) + .await + } + + pub(crate) async fn fetch_with_attempt_observer( + &self, + info_hash: &[u8; 20], + peer_addr: SocketAddr, + on_attempt: F, + ) -> MetadataFetchOutcome + where + F: FnOnce() + Send, + { let (cached_reason, cache_entries) = self.peer_failure_cache.get(peer_addr, Instant::now()); self.set_peer_failure_cache_entries(cache_entries); if let Some(reason) = cached_reason { @@ -215,6 +229,7 @@ impl RbitFetcher { return MetadataFetchOutcome::SkippedCached; } + on_attempt(); self.runtime_stats.metadata_peer_attempt(); #[cfg(feature = "metrics")] { diff --git a/src/peer_lookup.rs b/src/peer_lookup.rs new file mode 100644 index 0000000..53afa54 --- /dev/null +++ b/src/peer_lookup.rs @@ -0,0 +1,518 @@ +use crate::budget::RateBucket; +use crate::krpc::{encode_get_peers_query, for_each_response_node, for_each_response_peer}; +use crate::node_id::TransactionId; +use crate::protocol::DhtResponse; +use crate::routing_snapshot::{RoutingSnapshot, xor_distance_cmp}; +use crate::runtime_stats::DhtRuntimeStats; +use crate::server::HashDiscovered; +use crate::types::{NetMode, NodeTuple, PeerLookupOptions}; +use ahash::{AHashMap, AHashSet}; +use arc_swap::ArcSwap; +use bytes::BytesMut; +#[cfg(feature = "metrics")] +use metrics::counter; +use std::collections::VecDeque; +use std::net::SocketAddr; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use tokio::net::UdpSocket; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; + +const LOOKUP_TID_TAG: u8 = 0xa5; +const MAX_QUERIES_PER_LOOKUP: usize = 12; +const MAX_CONCURRENT_QUERIES_PER_LOOKUP: usize = 4; +const MAX_FRONTIER_NODES: usize = 64; +const MAX_PEERS_PER_LOOKUP: usize = 12; +const LOOKUP_TIMEOUT: Duration = Duration::from_secs(2); +const QUERY_TIMEOUT: Duration = Duration::from_millis(500); +const MAINTENANCE_INTERVAL: Duration = Duration::from_millis(25); +const REQUEST_CHANNEL_CAPACITY: usize = 1_024; +const RESPONSE_CHANNEL_CAPACITY: usize = 4_096; + +#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)] +struct PendingKey { + addr: SocketAddr, + tid: TransactionId, +} + +#[derive(Debug, Clone, Copy)] +struct PendingQuery { + lookup_id: u64, + deadline: Instant, +} + +struct LookupState { + info_hash: [u8; 20], + info_hash_hex: String, + frontier: Vec, + seen_nodes: AHashSet, + peers: AHashSet, + queried: usize, + outstanding: usize, + deadline: Instant, +} + +impl LookupState { + fn pop_closest(&mut self) -> Option { + let index = self + .frontier + .iter() + .enumerate() + .min_by(|(_, left), (_, right)| xor_distance_cmp(&left.id, &right.id, &self.info_hash)) + .map(|(index, _)| index)?; + Some(self.frontier.swap_remove(index)) + } + + fn is_complete(&self, now: Instant) -> bool { + self.deadline <= now + || self.peers.len() >= MAX_PEERS_PER_LOOKUP + || (self.outstanding == 0 + && (self.queried >= MAX_QUERIES_PER_LOOKUP || self.frontier.is_empty())) + } +} + +struct LookupResponse { + remote_addr: SocketAddr, + tid: TransactionId, + response: DhtResponse, +} + +#[derive(Clone)] +pub(crate) struct PeerLookupHandle { + request_tx: mpsc::Sender<[u8; 20]>, + response_tx: mpsc::Sender, + runtime_stats: DhtRuntimeStats, +} + +pub(crate) struct PeerLookupRuntime { + pub(crate) options: PeerLookupOptions, + pub(crate) stats: DhtRuntimeStats, + pub(crate) shutdown: CancellationToken, +} + +impl PeerLookupHandle { + pub(crate) fn request_sender(&self) -> mpsc::Sender<[u8; 20]> { + self.request_tx.clone() + } + + pub(crate) fn route_response( + &self, + remote_addr: SocketAddr, + tid: TransactionId, + response: DhtResponse, + ) { + if self + .response_tx + .try_send(LookupResponse { + remote_addr, + tid, + response, + }) + .is_err() + { + self.runtime_stats.peer_lookup_response_dropped(); + #[cfg(feature = "metrics")] + counter!("dht_peer_lookup_dropped_total", "reason" => "response_queue_full") + .increment(1); + } + } +} + +pub(crate) fn is_peer_lookup_tid(tid: &TransactionId) -> bool { + tid[0] == LOOKUP_TID_TAG +} + +pub(crate) fn spawn_peer_lookup( + netmode: NetMode, + local_id: [u8; 20], + sockets: &std::collections::HashMap>, + snapshot: Arc>, + hash_tx: mpsc::Sender, + runtime: PeerLookupRuntime, +) -> PeerLookupHandle { + let PeerLookupRuntime { + options, + stats, + shutdown, + } = runtime; + let (request_tx, request_rx) = mpsc::channel(REQUEST_CHANNEL_CAPACITY); + let (response_tx, response_rx) = mpsc::channel(RESPONSE_CHANNEL_CAPACITY); + let socket_v4 = sockets + .iter() + .find_map(|(addr, socket)| addr.is_ipv4().then(|| socket.clone())); + let socket_v6 = sockets + .iter() + .find_map(|(addr, socket)| addr.is_ipv6().then(|| socket.clone())); + let actor = PeerLookupActor { + netmode, + local_id, + socket_v4, + socket_v6, + snapshot, + hash_tx, + request_rx, + response_rx, + request_budget: RateBucket::per_second( + options.max_lookups_per_second, + options.burst, + true, + Instant::now(), + ), + max_active_lookups: options.max_active_lookups, + active: AHashMap::new(), + pending: AHashMap::new(), + pending_expiry: VecDeque::new(), + next_lookup_id: 1, + next_tid: 1, + runtime_stats: stats.clone(), + shutdown, + }; + tokio::spawn(actor.run()); + PeerLookupHandle { + request_tx, + response_tx, + runtime_stats: stats, + } +} + +struct PeerLookupActor { + netmode: NetMode, + local_id: [u8; 20], + socket_v4: Option>, + socket_v6: Option>, + snapshot: Arc>, + hash_tx: mpsc::Sender, + request_rx: mpsc::Receiver<[u8; 20]>, + response_rx: mpsc::Receiver, + request_budget: RateBucket, + max_active_lookups: usize, + active: AHashMap, + pending: AHashMap, + pending_expiry: VecDeque<(Instant, PendingKey)>, + next_lookup_id: u64, + next_tid: u64, + runtime_stats: DhtRuntimeStats, + shutdown: CancellationToken, +} + +impl PeerLookupActor { + async fn run(mut self) { + let mut maintenance = tokio::time::interval(MAINTENANCE_INTERVAL); + maintenance.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + tokio::select! { + biased; + _ = self.shutdown.cancelled() => break, + response = self.response_rx.recv() => { + let Some(response) = response else { break }; + self.handle_response(response, Instant::now()).await; + } + _ = maintenance.tick() => self.expire(Instant::now()).await, + request = self.request_rx.recv() => { + let Some(info_hash) = request else { break }; + self.start_lookup(info_hash, Instant::now()).await; + } + } + } + } + + async fn start_lookup(&mut self, info_hash: [u8; 20], now: Instant) { + self.runtime_stats.peer_lookup_requested(); + #[cfg(feature = "metrics")] + counter!("dht_peer_lookup_requests_total").increment(1); + if self.active.len() >= self.max_active_lookups || !self.request_budget.try_take_one(now) { + self.runtime_stats.peer_lookup_rate_limited(); + #[cfg(feature = "metrics")] + counter!("dht_peer_lookup_dropped_total", "reason" => "rate_limit").increment(1); + return; + } + + let filter_ipv6 = match (self.socket_v4.is_some(), self.socket_v6.is_some()) { + (true, false) => Some(false), + (false, true) => Some(true), + _ => None, + }; + let frontier = + self.snapshot + .load() + .closest_nodes(&info_hash, MAX_QUERIES_PER_LOOKUP, filter_ipv6); + if frontier.is_empty() { + self.runtime_stats.peer_lookup_empty(); + return; + } + + let lookup_id = self.next_lookup_id; + self.next_lookup_id = self.next_lookup_id.wrapping_add(1).max(1); + let seen_nodes = frontier.iter().map(|node| node.addr).collect(); + self.active.insert( + lookup_id, + LookupState { + info_hash, + info_hash_hex: hex::encode(info_hash), + frontier, + seen_nodes, + peers: AHashSet::new(), + queried: 0, + outstanding: 0, + deadline: now + LOOKUP_TIMEOUT, + }, + ); + self.runtime_stats.peer_lookup_started(); + #[cfg(feature = "metrics")] + counter!("dht_peer_lookup_started_total").increment(1); + self.dispatch_more(lookup_id, now).await; + } + + async fn dispatch_more(&mut self, lookup_id: u64, now: Instant) { + loop { + let Some((node, info_hash)) = self.active.get_mut(&lookup_id).and_then(|state| { + if state.deadline <= now + || state.peers.len() >= MAX_PEERS_PER_LOOKUP + || state.queried >= MAX_QUERIES_PER_LOOKUP + || state.outstanding >= MAX_CONCURRENT_QUERIES_PER_LOOKUP + { + return None; + } + let node = state.pop_closest()?; + state.queried += 1; + Some((node, state.info_hash)) + }) else { + break; + }; + + let tid = self.next_transaction_id(); + let key = PendingKey { + addr: node.addr, + tid, + }; + let mut buffer = BytesMut::with_capacity(128); + encode_get_peers_query(&mut buffer, &tid, &info_hash, &self.local_id); + let socket = if node.addr.is_ipv4() { + self.socket_v4.clone() + } else { + self.socket_v6.clone() + }; + let sent = match socket { + Some(socket) => socket.send_to(&buffer, node.addr).await.is_ok(), + None => false, + }; + if !sent { + self.runtime_stats.peer_lookup_send_failed(); + continue; + } + + let deadline = now + QUERY_TIMEOUT; + self.pending.insert( + key, + PendingQuery { + lookup_id, + deadline, + }, + ); + self.pending_expiry.push_back((deadline, key)); + if let Some(state) = self.active.get_mut(&lookup_id) { + state.outstanding += 1; + } + self.runtime_stats.udp_sent(buffer.len()); + self.runtime_stats.peer_lookup_query(); + #[cfg(feature = "metrics")] + counter!("dht_peer_lookup_queries_total").increment(1); + } + self.finish_if_complete(lookup_id, now); + } + + async fn handle_response(&mut self, event: LookupResponse, now: Instant) { + let key = PendingKey { + addr: event.remote_addr, + tid: event.tid, + }; + let Some(pending) = self.pending.remove(&key) else { + return; + }; + let Some(state) = self.active.get_mut(&pending.lookup_id) else { + return; + }; + state.outstanding = state.outstanding.saturating_sub(1); + self.runtime_stats.peer_lookup_response(); + + let mut discovered = Vec::new(); + for_each_response_peer(&event.response, self.netmode, |peer| { + if state.peers.len() < MAX_PEERS_PER_LOOKUP && state.peers.insert(peer) { + discovered.push(peer); + } + }); + let mut response_nodes = Vec::new(); + for_each_response_node(&event.response, self.netmode, |node| { + response_nodes.push(node) + }); + for node in response_nodes { + if state.frontier.len() >= MAX_FRONTIER_NODES { + break; + } + if state.seen_nodes.insert(node.addr) { + state.frontier.push(node); + } + } + let hash = state.info_hash_hex.clone(); + let lookup_id = pending.lookup_id; + let _ = state; + + for peer in discovered { + let event = HashDiscovered { + info_hash: hash.clone(), + peer_addr: peer, + discovered_at: now, + }; + if self.hash_tx.try_send(event).is_ok() { + self.runtime_stats.peer_lookup_peer_found(); + #[cfg(feature = "metrics")] + counter!("dht_peer_lookup_peers_found_total").increment(1); + } else { + self.runtime_stats.peer_lookup_output_dropped(); + #[cfg(feature = "metrics")] + counter!("dht_peer_lookup_dropped_total", "reason" => "hash_queue_full") + .increment(1); + } + } + self.dispatch_more(lookup_id, now).await; + } + + async fn expire(&mut self, now: Instant) { + let mut affected = AHashSet::new(); + while let Some((deadline, key)) = self.pending_expiry.front().copied() { + if deadline > now { + break; + } + self.pending_expiry.pop_front(); + let should_remove = self + .pending + .get(&key) + .is_some_and(|pending| pending.deadline == deadline); + if !should_remove { + continue; + } + let pending = self.pending.remove(&key).expect("pending lookup exists"); + if let Some(state) = self.active.get_mut(&pending.lookup_id) { + state.outstanding = state.outstanding.saturating_sub(1); + affected.insert(pending.lookup_id); + } + self.runtime_stats.peer_lookup_timeout(); + } + for lookup_id in affected { + self.dispatch_more(lookup_id, now).await; + } + let expired: Vec<_> = self + .active + .iter() + .filter_map(|(lookup_id, state)| (state.deadline <= now).then_some(*lookup_id)) + .collect(); + for lookup_id in expired { + self.finish_lookup(lookup_id); + } + } + + fn finish_if_complete(&mut self, lookup_id: u64, now: Instant) { + if self + .active + .get(&lookup_id) + .is_some_and(|state| state.is_complete(now)) + { + self.finish_lookup(lookup_id); + } + } + + fn finish_lookup(&mut self, lookup_id: u64) { + self.active.remove(&lookup_id); + self.pending + .retain(|_, pending| pending.lookup_id != lookup_id); + } + + fn next_transaction_id(&mut self) -> TransactionId { + let mut tid = self.next_tid.to_be_bytes(); + tid[0] = LOOKUP_TID_TAG; + self.next_tid = self.next_tid.wrapping_add(1).max(1); + tid + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::protocol::DhtMessage; + + #[test] + fn lookup_transaction_ids_have_a_reserved_tag() { + let tid = [LOOKUP_TID_TAG, 1, 2, 3, 4, 5, 6, 7]; + assert!(is_peer_lookup_tid(&tid)); + assert!(!is_peer_lookup_tid(&[0; 8])); + } + + #[tokio::test] + async fn lookup_response_feeds_discovered_peer_back_to_metadata_scheduler() { + let local_socket = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap()); + let remote_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let remote_addr = remote_socket.local_addr().unwrap(); + let local_addr = local_socket.local_addr().unwrap(); + let sockets = std::collections::HashMap::from([(local_addr, local_socket)]); + let snapshot = Arc::new(ArcSwap::from_pointee(RoutingSnapshot::from_nodes( + vec![NodeTuple { + id: [9; 20], + addr: remote_addr, + }], + 1, + ))); + let (hash_tx, mut hash_rx) = mpsc::channel(4); + let stats = DhtRuntimeStats::default(); + let shutdown = CancellationToken::new(); + let handle = spawn_peer_lookup( + NetMode::Ipv4Only, + [7; 20], + &sockets, + snapshot, + hash_tx, + PeerLookupRuntime { + options: PeerLookupOptions::default(), + stats: stats.clone(), + shutdown: shutdown.clone(), + }, + ); + handle.request_tx.send([3; 20]).await.unwrap(); + + let mut buffer = [0u8; 512]; + let (len, source) = + tokio::time::timeout(Duration::from_secs(1), remote_socket.recv_from(&mut buffer)) + .await + .unwrap() + .unwrap(); + assert_eq!(source, local_addr); + let query: DhtMessage = serde_bencode::from_bytes(&buffer[..len]).unwrap(); + assert_eq!(query.q.as_deref(), Some("get_peers")); + let tid: TransactionId = query.t.as_ref().try_into().unwrap(); + handle.route_response( + remote_addr, + tid, + DhtResponse { + id: Some(serde_bytes::ByteBuf::from(vec![9; 20])), + nodes: None, + nodes6: None, + values: Some(vec![serde_bytes::ByteBuf::from(vec![ + 8, 8, 4, 4, 0x1a, 0xe1, + ])]), + }, + ); + + let event = tokio::time::timeout(Duration::from_secs(1), hash_rx.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!(event.info_hash, hex::encode([3; 20])); + assert_eq!(event.peer_addr, "8.8.4.4:6881".parse().unwrap()); + let snapshot = stats.snapshot(); + assert_eq!(snapshot.peer_lookup_started, 1); + assert_eq!(snapshot.peer_lookup_queries, 1); + assert_eq!(snapshot.peer_lookup_responses, 1); + assert_eq!(snapshot.peer_lookup_peers_found, 1); + shutdown.cancel(); + } +} diff --git a/src/protocol.rs b/src/protocol.rs index e12f0ef..6a5f3dc 100644 --- a/src/protocol.rs +++ b/src/protocol.rs @@ -48,4 +48,7 @@ pub struct DhtResponse { #[serde(default)] /// Compact IPv6 node tuples. pub nodes6: Option, + #[serde(default)] + /// Compact Peer endpoints returned by `get_peers`. + pub values: Option>, } diff --git a/src/routing_snapshot.rs b/src/routing_snapshot.rs index b05c1e0..bdac0b0 100644 --- a/src/routing_snapshot.rs +++ b/src/routing_snapshot.rs @@ -34,4 +34,74 @@ impl RoutingSnapshot { } } } + + pub(crate) fn closest_nodes( + &self, + target: &[u8; 20], + count: usize, + filter_ipv6: Option, + ) -> Vec { + if count == 0 { + return Vec::new(); + } + let mut nodes = match filter_ipv6 { + Some(true) => self.v6.clone(), + Some(false) => self.v4.clone(), + None => { + let mut all = Vec::with_capacity(self.v4.len() + self.v6.len()); + all.extend_from_slice(&self.v4); + all.extend_from_slice(&self.v6); + all + } + }; + let compare = + |left: &NodeTuple, right: &NodeTuple| xor_distance_cmp(&left.id, &right.id, target); + if nodes.len() > count { + nodes.select_nth_unstable_by(count, compare); + nodes.truncate(count); + } + nodes.sort_unstable_by(compare); + nodes + } +} + +pub(crate) fn xor_distance_cmp( + left: &[u8; 20], + right: &[u8; 20], + target: &[u8; 20], +) -> std::cmp::Ordering { + for index in 0..20 { + let ordering = (left[index] ^ target[index]).cmp(&(right[index] ^ target[index])); + if ordering != std::cmp::Ordering::Equal { + return ordering; + } + } + std::cmp::Ordering::Equal +} + +#[cfg(test)] +mod tests { + use super::*; + use std::net::SocketAddr; + + fn node(id: u8, port: u16) -> NodeTuple { + NodeTuple { + id: [id; 20], + addr: SocketAddr::from(([8, 8, 8, 8], port)), + } + } + + #[test] + fn closest_nodes_orders_by_xor_distance_and_limits_results() { + let snapshot = + RoutingSnapshot::from_nodes(vec![node(0xf0, 1), node(0x01, 2), node(0x10, 3)], 3); + let closest = snapshot.closest_nodes(&[0; 20], 2, Some(false)); + assert_eq!( + closest + .iter() + .map(|node| node.addr.port()) + .collect::>(), + vec![2, 3] + ); + } } diff --git a/src/runtime_stats.rs b/src/runtime_stats.rs index 3b648aa..b6a5b12 100644 --- a/src/runtime_stats.rs +++ b/src/runtime_stats.rs @@ -216,6 +216,36 @@ pub struct DhtRuntimeSnapshot { pub metadata_queue_evicted: u64, /// Events rejected or expired because they were stale. pub metadata_queue_stale: u64, + /// Metadata Peer races started after admission. + pub metadata_races_started: u64, + /// Distinct Peer candidates launched into Metadata races. + pub metadata_peer_candidates: u64, + /// Peer candidates added while a Metadata race was already running. + pub metadata_live_peer_joins: u64, + /// Losing Peer attempts canceled after another Peer succeeded. + pub metadata_peer_canceled: u64, + /// Active `get_peers` lookup requests received from Metadata jobs. + pub peer_lookup_requested: u64, + /// Active `get_peers` lookups admitted by the bounded rate limiter. + pub peer_lookup_started: u64, + /// Active lookup requests rejected by concurrency or rate limits. + pub peer_lookup_rate_limited: u64, + /// Active lookup requests skipped because the routing snapshot was empty. + pub peer_lookup_empty: u64, + /// Active `get_peers` UDP queries sent. + pub peer_lookup_queries: u64, + /// Matched active `get_peers` responses. + pub peer_lookup_responses: u64, + /// Active `get_peers` queries that timed out. + pub peer_lookup_timeouts: u64, + /// Active `get_peers` UDP sends that failed immediately. + pub peer_lookup_send_failures: u64, + /// Unique Peer endpoints fed back into Metadata scheduling. + pub peer_lookup_peers_found: u64, + /// Tagged lookup responses dropped before reaching the lookup actor. + pub peer_lookup_response_dropped: u64, + /// Discovered Peer endpoints dropped because Hash ingress was full. + pub peer_lookup_output_dropped: u64, /// Real Peer network attempts. pub metadata_peer_attempts: u64, /// Successful Peer downloads and parses. @@ -295,6 +325,21 @@ struct DhtRuntimeStatsInner { metadata_queue_deduplicated: AtomicU64, metadata_queue_evicted: AtomicU64, metadata_queue_stale: AtomicU64, + metadata_races_started: AtomicU64, + metadata_peer_candidates: AtomicU64, + metadata_live_peer_joins: AtomicU64, + metadata_peer_canceled: AtomicU64, + peer_lookup_requested: AtomicU64, + peer_lookup_started: AtomicU64, + peer_lookup_rate_limited: AtomicU64, + peer_lookup_empty: AtomicU64, + peer_lookup_queries: AtomicU64, + peer_lookup_responses: AtomicU64, + peer_lookup_timeouts: AtomicU64, + peer_lookup_send_failures: AtomicU64, + peer_lookup_peers_found: AtomicU64, + peer_lookup_response_dropped: AtomicU64, + peer_lookup_output_dropped: AtomicU64, metadata_peer_attempts: AtomicU64, metadata_peer_succeeded: AtomicU64, metadata_peer_failed: AtomicU64, @@ -382,6 +427,21 @@ impl Default for DhtRuntimeStatsInner { metadata_queue_deduplicated: AtomicU64::new(0), metadata_queue_evicted: AtomicU64::new(0), metadata_queue_stale: AtomicU64::new(0), + metadata_races_started: AtomicU64::new(0), + metadata_peer_candidates: AtomicU64::new(0), + metadata_live_peer_joins: AtomicU64::new(0), + metadata_peer_canceled: AtomicU64::new(0), + peer_lookup_requested: AtomicU64::new(0), + peer_lookup_started: AtomicU64::new(0), + peer_lookup_rate_limited: AtomicU64::new(0), + peer_lookup_empty: AtomicU64::new(0), + peer_lookup_queries: AtomicU64::new(0), + peer_lookup_responses: AtomicU64::new(0), + peer_lookup_timeouts: AtomicU64::new(0), + peer_lookup_send_failures: AtomicU64::new(0), + peer_lookup_peers_found: AtomicU64::new(0), + peer_lookup_response_dropped: AtomicU64::new(0), + peer_lookup_output_dropped: AtomicU64::new(0), metadata_peer_attempts: AtomicU64::new(0), metadata_peer_succeeded: AtomicU64::new(0), metadata_peer_failed: AtomicU64::new(0), @@ -492,6 +552,23 @@ impl DhtRuntimeStats { metadata_queue_deduplicated: inner.metadata_queue_deduplicated.load(Ordering::Relaxed), metadata_queue_evicted: inner.metadata_queue_evicted.load(Ordering::Relaxed), metadata_queue_stale: inner.metadata_queue_stale.load(Ordering::Relaxed), + metadata_races_started: inner.metadata_races_started.load(Ordering::Relaxed), + metadata_peer_candidates: inner.metadata_peer_candidates.load(Ordering::Relaxed), + metadata_live_peer_joins: inner.metadata_live_peer_joins.load(Ordering::Relaxed), + metadata_peer_canceled: inner.metadata_peer_canceled.load(Ordering::Relaxed), + peer_lookup_requested: inner.peer_lookup_requested.load(Ordering::Relaxed), + peer_lookup_started: inner.peer_lookup_started.load(Ordering::Relaxed), + peer_lookup_rate_limited: inner.peer_lookup_rate_limited.load(Ordering::Relaxed), + peer_lookup_empty: inner.peer_lookup_empty.load(Ordering::Relaxed), + peer_lookup_queries: inner.peer_lookup_queries.load(Ordering::Relaxed), + peer_lookup_responses: inner.peer_lookup_responses.load(Ordering::Relaxed), + peer_lookup_timeouts: inner.peer_lookup_timeouts.load(Ordering::Relaxed), + peer_lookup_send_failures: inner.peer_lookup_send_failures.load(Ordering::Relaxed), + peer_lookup_peers_found: inner.peer_lookup_peers_found.load(Ordering::Relaxed), + peer_lookup_response_dropped: inner + .peer_lookup_response_dropped + .load(Ordering::Relaxed), + peer_lookup_output_dropped: inner.peer_lookup_output_dropped.load(Ordering::Relaxed), metadata_peer_attempts: inner.metadata_peer_attempts.load(Ordering::Relaxed), metadata_peer_succeeded: inner.metadata_peer_succeeded.load(Ordering::Relaxed), metadata_peer_failed: inner.metadata_peer_failed.load(Ordering::Relaxed), @@ -642,6 +719,94 @@ impl DhtRuntimeStats { .fetch_add(count as u64, Ordering::Relaxed); } + pub(crate) fn metadata_race_started(&self) { + self.inner + .metadata_races_started + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn metadata_peer_candidate(&self) { + self.inner + .metadata_peer_candidates + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn metadata_live_peer_join(&self) { + self.inner + .metadata_live_peer_joins + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn metadata_peer_canceled(&self, count: usize) { + self.inner + .metadata_peer_canceled + .fetch_add(count as u64, Ordering::Relaxed); + } + + pub(crate) fn peer_lookup_requested(&self) { + self.inner + .peer_lookup_requested + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn peer_lookup_started(&self) { + self.inner + .peer_lookup_started + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn peer_lookup_rate_limited(&self) { + self.inner + .peer_lookup_rate_limited + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn peer_lookup_empty(&self) { + self.inner.peer_lookup_empty.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn peer_lookup_query(&self) { + self.inner + .peer_lookup_queries + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn peer_lookup_response(&self) { + self.inner + .peer_lookup_responses + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn peer_lookup_timeout(&self) { + self.inner + .peer_lookup_timeouts + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn peer_lookup_send_failed(&self) { + self.inner + .peer_lookup_send_failures + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn peer_lookup_peer_found(&self) { + self.inner + .peer_lookup_peers_found + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn peer_lookup_response_dropped(&self) { + self.inner + .peer_lookup_response_dropped + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn peer_lookup_output_dropped(&self) { + self.inner + .peer_lookup_output_dropped + .fetch_add(1, Ordering::Relaxed); + } + pub(crate) fn metadata_peer_attempt(&self) { self.inner .metadata_peer_attempts @@ -953,6 +1118,21 @@ mod tests { writer.metadata_queue_deduplicated(); writer.metadata_queue_evicted(); writer.metadata_queue_stale(3); + writer.metadata_race_started(); + writer.metadata_peer_candidate(); + writer.metadata_live_peer_join(); + writer.metadata_peer_canceled(2); + writer.peer_lookup_requested(); + writer.peer_lookup_started(); + writer.peer_lookup_rate_limited(); + writer.peer_lookup_empty(); + writer.peer_lookup_query(); + writer.peer_lookup_response(); + writer.peer_lookup_timeout(); + writer.peer_lookup_send_failed(); + writer.peer_lookup_peer_found(); + writer.peer_lookup_response_dropped(); + writer.peer_lookup_output_dropped(); writer.metadata_peer_attempt(); writer.metadata_peer_succeeded(); writer.metadata_peer_failed(); @@ -997,6 +1177,21 @@ mod tests { metadata_queue_deduplicated: 1, metadata_queue_evicted: 1, metadata_queue_stale: 3, + metadata_races_started: 1, + metadata_peer_candidates: 1, + metadata_live_peer_joins: 1, + metadata_peer_canceled: 2, + peer_lookup_requested: 1, + peer_lookup_started: 1, + peer_lookup_rate_limited: 1, + peer_lookup_empty: 1, + peer_lookup_queries: 1, + peer_lookup_responses: 1, + peer_lookup_timeouts: 1, + peer_lookup_send_failures: 1, + peer_lookup_peers_found: 1, + peer_lookup_response_dropped: 1, + peer_lookup_output_dropped: 1, metadata_peer_attempts: 1, metadata_peer_succeeded: 1, metadata_peer_failed: 1, diff --git a/src/scheduler.rs b/src/scheduler.rs index 7def3f7..128930b 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -1,4 +1,4 @@ -use crate::metadata::{MetadataFetchOutcome, RbitFetcher}; +use crate::metadata::{FetchedMetadata, MetadataFetchOutcome, RbitFetcher}; #[cfg(test)] use crate::runtime_stats::DhtRuntimeLimits; use crate::runtime_stats::DhtRuntimeStats; @@ -8,6 +8,7 @@ use arc_swap::ArcSwapOption; #[cfg(feature = "metrics")] use metrics::{counter, gauge, histogram}; use std::collections::{BTreeMap, HashMap, HashSet}; +use std::future::Future; use std::net::SocketAddr; use std::panic::{AssertUnwindSafe, catch_unwind}; use std::sync::Arc; @@ -17,8 +18,13 @@ use tokio::sync::mpsc; use tokio::task::JoinSet; use tokio_util::sync::CancellationToken; -const MAX_PEERS_PER_HASH: usize = 3; +/// Maximum distinct Peer endpoints retained and raced for one InfoHash. +pub const MAX_METADATA_PEERS_PER_HASH: usize = 12; const HASH_QUEUE_TTL: Duration = Duration::from_secs(60); +const PEER_COALESCE_WINDOW: Duration = Duration::from_millis(25); +const LIVE_PEER_IDLE_GRACE: Duration = Duration::from_millis(300); +const ACTIVE_PEER_LOOKUP_DELAY: Duration = Duration::from_millis(100); +const DISPATCH_TICK: Duration = Duration::from_millis(25); /// Callback returning whether the application accepted a downloaded torrent. pub type TorrentAckCallback = Box bool + Send + Sync + 'static>; @@ -52,6 +58,11 @@ pub struct MetadataSchedulerCallbacks { pub completion: Arc>, } +pub(crate) struct MetadataSchedulerRuntime { + pub(crate) stats: DhtRuntimeStats, + pub(crate) peer_lookup_tx: Option>, +} + #[derive(Debug, Clone, Copy)] struct PeerCandidate { addr: SocketAddr, @@ -62,6 +73,8 @@ struct PeerCandidate { struct QueuedHash { info_hash: String, peers: Vec, + queued_at: Instant, + ready_at: Instant, latest_at: Instant, order_key: (Instant, u64), } @@ -137,7 +150,7 @@ impl PendingHashQueue { entry .peers .sort_unstable_by(|left, right| right.discovered_at.cmp(&left.discovered_at)); - entry.peers.truncate(MAX_PEERS_PER_HASH); + entry.peers.truncate(MAX_METADATA_PEERS_PER_HASH); entry.latest_at = entry.latest_at.max(event.discovered_at); entry.order_key = self.next_order_key(entry.latest_at); self.order.insert(entry.order_key, entry.info_hash.clone()); @@ -170,6 +183,8 @@ impl PendingHashQueue { addr: event.peer_addr, discovered_at: event.discovered_at, }], + queued_at: now, + ready_at: now + PEER_COALESCE_WINDOW, latest_at: event.discovered_at, order_key, }, @@ -183,15 +198,15 @@ impl PendingHashQueue { Some(entry) } - fn pop_newest_available( + fn pop_newest_ready( &mut self, + now: Instant, in_flight: &HashMap, ) -> Option { - let info_hash = self - .order - .iter() - .rev() - .find_map(|(_, hash)| (!in_flight.contains_key(hash)).then(|| hash.clone()))?; + let info_hash = self.order.iter().rev().find_map(|(_, hash)| { + let entry = self.entries.get(hash)?; + (!in_flight.contains_key(hash) && entry.ready_at <= now).then(|| hash.clone()) + })?; self.remove(&info_hash) } @@ -217,12 +232,12 @@ impl PendingHashQueue { struct MetadataJob { info_hash: String, peers: Vec, - already_admitted: bool, + peer_rx: mpsc::Receiver, } #[derive(Debug)] struct InFlightState { - attempts: usize, + peer_tx: mpsc::Sender, scheduled_peers: HashSet, } @@ -254,6 +269,7 @@ pub struct MetadataScheduler { total_completed: Arc, queue_len: Arc, runtime_stats: DhtRuntimeStats, + peer_lookup_tx: Option>, shutdown: CancellationToken, } @@ -286,6 +302,29 @@ impl MetadataScheduler { queue_len: Arc, shutdown: CancellationToken, runtime_stats: DhtRuntimeStats, + ) -> Self { + Self::new_with_runtime( + hash_rx, + fetcher, + limits, + callbacks, + queue_len, + shutdown, + MetadataSchedulerRuntime { + stats: runtime_stats, + peer_lookup_tx: None, + }, + ) + } + + pub(crate) fn new_with_runtime( + hash_rx: mpsc::Receiver, + fetcher: Arc, + limits: MetadataSchedulerLimits, + callbacks: MetadataSchedulerCallbacks, + queue_len: Arc, + shutdown: CancellationToken, + runtime: MetadataSchedulerRuntime, ) -> Self { Self { hash_rx, @@ -300,7 +339,8 @@ impl MetadataScheduler { total_dispatched: Arc::new(AtomicU64::new(0)), total_completed: Arc::new(AtomicU64::new(0)), queue_len, - runtime_stats, + runtime_stats: runtime.stats, + peer_lookup_tx: runtime.peer_lookup_tx, shutdown, } } @@ -326,7 +366,11 @@ impl MetadataScheduler { let mut in_flight = HashMap::::new(); let mut tasks = JoinSet::::new(); let mut maintenance = tokio::time::interval(Duration::from_secs(1)); + maintenance.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + let mut dispatch_tick = tokio::time::interval(DISPATCH_TICK); + dispatch_tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); let mut stats_interval = tokio::time::interval(Duration::from_secs(60)); + stats_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); let mut input_closed = false; loop { @@ -343,14 +387,14 @@ impl MetadataScheduler { self.runtime_stats .set_hash_ingress_queue_depth(self.hash_rx.len()); match result { - Some(hash) => self.enqueue(&mut queue, hash), + Some(hash) => self.enqueue(&mut queue, &mut in_flight, hash), None => input_closed = true, } } result = tasks.join_next(), if !tasks.is_empty() => { match result { Some(Ok(result)) => { - self.handle_job_result(result, &mut queue, &mut in_flight, &mut tasks); + self.handle_job_result(result, &mut queue, &mut in_flight); } Some(Err(error)) => { #[cfg(feature = "metrics")] @@ -360,6 +404,7 @@ impl MetadataScheduler { None => {} } } + _ = dispatch_tick.tick() => {} _ = maintenance.tick() => { let expired = queue.expire(Instant::now()); if expired > 0 { @@ -383,9 +428,50 @@ impl MetadataScheduler { gauge!("dht_metadata_queue_depth").set(0.0); } - fn enqueue(&self, queue: &mut PendingHashQueue, hash: HashDiscovered) { + fn enqueue( + &self, + queue: &mut PendingHashQueue, + in_flight: &mut HashMap, + hash: HashDiscovered, + ) { self.total_received.fetch_add(1, Ordering::Relaxed); - let kind = queue.push(hash, Instant::now()); + let now = Instant::now(); + if let Some(state) = in_flight.get_mut(&hash.info_hash) { + let kind = if now + .checked_duration_since(hash.discovered_at) + .unwrap_or_default() + > HASH_QUEUE_TTL + { + QueuePushKind::Stale + } else if state.scheduled_peers.contains(&hash.peer_addr) + || state.scheduled_peers.len() >= MAX_METADATA_PEERS_PER_HASH + { + QueuePushKind::Updated + } else { + let peer = PeerCandidate { + addr: hash.peer_addr, + discovered_at: hash.discovered_at, + }; + match state.peer_tx.try_send(peer) { + Ok(()) => { + state.scheduled_peers.insert(hash.peer_addr); + self.runtime_stats.metadata_live_peer_join(); + #[cfg(feature = "metrics")] + counter!("dht_metadata_live_peer_joins_total").increment(1); + QueuePushKind::Updated + } + Err(_) => QueuePushKind::Updated, + } + }; + self.record_queue_push(kind); + return; + } + + let kind = queue.push(hash, now); + self.record_queue_push(kind); + } + + fn record_queue_push(&self, kind: QueuePushKind) { let result = match kind { QueuePushKind::Inserted => { self.runtime_stats.metadata_queue_inserted(); @@ -420,68 +506,45 @@ impl MetadataScheduler { tasks: &mut JoinSet, ) { while tasks.len() < self.max_concurrent { - let Some(entry) = queue.pop_newest_available(in_flight) else { + let Some(mut entry) = queue.pop_newest_ready(Instant::now(), in_flight) else { break; }; let queue_wait = Instant::now() - .checked_duration_since(entry.latest_at) + .checked_duration_since(entry.queued_at) .unwrap_or_default(); self.runtime_stats.observe_metadata_queue_wait( queue_wait.as_millis().min(u128::from(u64::MAX)) as u64, ); #[cfg(feature = "metrics")] histogram!("dht_metadata_queue_wait_seconds").record(queue_wait.as_secs_f64()); + let now = Instant::now(); + entry.peers.retain(|peer| { + now.checked_duration_since(peer.discovered_at) + .unwrap_or_default() + <= HASH_QUEUE_TTL + }); + entry.peers.truncate(MAX_METADATA_PEERS_PER_HASH); + if entry.peers.is_empty() { + self.total_dropped.fetch_add(1, Ordering::Relaxed); + self.runtime_stats.metadata_queue_stale(1); + continue; + } + let (peer_tx, peer_rx) = mpsc::channel(MAX_METADATA_PEERS_PER_HASH); let state = InFlightState { - attempts: 0, - scheduled_peers: HashSet::new(), + peer_tx, + scheduled_peers: entry.peers.iter().map(|peer| peer.addr).collect(), }; in_flight.insert(entry.info_hash.clone(), state); - self.spawn_job(entry, false, in_flight, tasks); + self.spawn_job(entry, peer_rx, tasks); } } fn spawn_job( &self, - mut entry: QueuedHash, - already_admitted: bool, - in_flight: &mut HashMap, + entry: QueuedHash, + peer_rx: mpsc::Receiver, tasks: &mut JoinSet, ) { - let Some(state) = in_flight.get_mut(&entry.info_hash) else { - return; - }; - let remaining = MAX_PEERS_PER_HASH.saturating_sub(state.attempts); - let now = Instant::now(); - entry.peers.retain(|peer| { - !state.scheduled_peers.contains(&peer.addr) - && now - .checked_duration_since(peer.discovered_at) - .unwrap_or_default() - <= HASH_QUEUE_TTL - }); - entry.peers.truncate(remaining); - for peer in &entry.peers { - state.scheduled_peers.insert(peer.addr); - } - - if entry.peers.is_empty() { - let attempts = state.attempts; - if already_admitted { - self.finish( - entry.info_hash.clone(), - MetadataFetchCompletionStatus::FetchFailed, - attempts, - ); - } else { - self.total_dropped.fetch_add(1, Ordering::Relaxed); - self.runtime_stats.metadata_queue_stale(1); - #[cfg(feature = "metrics")] - counter!("dht_metadata_queue_events_total", "result" => "expired").increment(1); - } - in_flight.remove(&entry.info_hash); - return; - } - self.total_dispatched.fetch_add(1, Ordering::Relaxed); #[cfg(feature = "metrics")] counter!("dht_metadata_jobs_dispatched_total").increment(1); @@ -489,13 +552,23 @@ impl MetadataScheduler { let job = MetadataJob { info_hash: entry.info_hash, peers: entry.peers, - already_admitted, + peer_rx, }; let fetcher = self.fetcher.clone(); let callback = self.callback.clone(); let on_metadata_fetch = self.on_metadata_fetch.clone(); + let runtime_stats = self.runtime_stats.clone(); + let peer_lookup_tx = self.peer_lookup_tx.clone(); tasks.spawn(async move { - Self::process_hash(job, &fetcher, &callback, &on_metadata_fetch).await + Self::process_hash( + job, + &fetcher, + &callback, + &on_metadata_fetch, + &runtime_stats, + peer_lookup_tx, + ) + .await }); } @@ -504,12 +577,10 @@ impl MetadataScheduler { result: JobResult, queue: &mut PendingHashQueue, in_flight: &mut HashMap, - tasks: &mut JoinSet, ) { - let Some(state) = in_flight.get_mut(&result.info_hash) else { + if !in_flight.contains_key(&result.info_hash) { return; - }; - state.attempts = state.attempts.saturating_add(result.attempts); + } match result.outcome { JobOutcome::GateRejected => { @@ -519,26 +590,10 @@ impl MetadataScheduler { counter!("dht_metadata_jobs_completed_total", "result" => "gate_rejected") .increment(1); } - JobOutcome::Completed(MetadataFetchCompletionStatus::FetchFailed) - if state.attempts < MAX_PEERS_PER_HASH => - { - if let Some(entry) = queue.remove(&result.info_hash) { - self.spawn_job(entry, true, in_flight, tasks); - } else { - let attempts = state.attempts; - in_flight.remove(&result.info_hash); - self.finish( - result.info_hash, - MetadataFetchCompletionStatus::FetchFailed, - attempts, - ); - } - } JobOutcome::Completed(status) => { - let attempts = state.attempts; queue.remove(&result.info_hash); in_flight.remove(&result.info_hash); - self.finish(result.info_hash, status, attempts); + self.finish(result.info_hash, status, result.attempts); } } } @@ -570,13 +625,14 @@ impl MetadataScheduler { } async fn process_hash( - job: MetadataJob, + mut job: MetadataJob, fetcher: &Arc, callback: &Arc>, on_metadata_fetch: &Arc>, + runtime_stats: &DhtRuntimeStats, + peer_lookup_tx: Option>, ) -> JobResult { - if !job.already_admitted - && let Some(gate) = on_metadata_fetch.load_full() + if let Some(gate) = on_metadata_fetch.load_full() && !gate(job.info_hash.clone()).await { return JobResult { @@ -601,21 +657,39 @@ impl MetadataScheduler { } }; - let mut attempts = 0; - for peer in job.peers { - let (name, total_size, files, piece_length) = - match fetcher.fetch(&info_hash_bytes, peer.addr).await { - MetadataFetchOutcome::Fetched(metadata) => { - attempts += 1; - metadata - } - MetadataFetchOutcome::Failed => { - attempts += 1; - continue; - } - MetadataFetchOutcome::SkippedCached => continue, - }; + let attempts = Arc::new(AtomicUsize::new(0)); + let attempts_for_fetch = attempts.clone(); + let fetcher_for_race = fetcher.clone(); + let lookup_task = peer_lookup_tx.map(|peer_lookup_tx| { + tokio::spawn(async move { + tokio::time::sleep(ACTIVE_PEER_LOOKUP_DELAY).await; + let _ = peer_lookup_tx.try_send(info_hash_bytes); + }) + }); + let fetched = race_peer_fetches( + job.peers, + &mut job.peer_rx, + move |peer| { + let attempts = attempts_for_fetch.clone(); + let fetcher = fetcher_for_race.clone(); + async move { + let outcome = fetcher + .fetch_with_attempt_observer(&info_hash_bytes, peer.addr, move || { + attempts.fetch_add(1, Ordering::Relaxed); + }) + .await; + (peer, outcome) + } + }, + runtime_stats, + ) + .await; + if let Some(task) = lookup_task { + task.abort(); + } + let attempts = attempts.load(Ordering::Relaxed); + if let Some((peer, (name, total_size, files, piece_length))) = fetched { let metadata = TorrentInfo { info_hash: job.info_hash.clone(), name, @@ -686,9 +760,99 @@ impl MetadataScheduler { } } +async fn race_peer_fetches( + peers: Vec, + peer_rx: &mut mpsc::Receiver, + fetch: F, + runtime_stats: &DhtRuntimeStats, +) -> Option<(PeerCandidate, FetchedMetadata)> +where + F: Fn(PeerCandidate) -> Fut + Clone + Send + Sync + 'static, + Fut: Future + Send + 'static, +{ + runtime_stats.metadata_race_started(); + #[cfg(feature = "metrics")] + counter!("dht_metadata_peer_races_total").increment(1); + + let mut tasks = JoinSet::new(); + let mut candidates = 0usize; + for peer in peers { + let fetch = fetch.clone(); + tasks.spawn(fetch(peer)); + candidates += 1; + runtime_stats.metadata_peer_candidate(); + #[cfg(feature = "metrics")] + counter!("dht_metadata_peer_candidates_total", "source" => "initial").increment(1); + } + + let mut accepting_live_peers = candidates < MAX_METADATA_PEERS_PER_HASH; + loop { + if tasks.is_empty() { + if !accepting_live_peers { + return None; + } + match tokio::time::timeout(LIVE_PEER_IDLE_GRACE, peer_rx.recv()).await { + Ok(Some(peer)) => { + let fetch = fetch.clone(); + tasks.spawn(fetch(peer)); + candidates += 1; + runtime_stats.metadata_peer_candidate(); + #[cfg(feature = "metrics")] + counter!("dht_metadata_peer_candidates_total", "source" => "live").increment(1); + accepting_live_peers = candidates < MAX_METADATA_PEERS_PER_HASH; + } + Ok(None) | Err(_) => return None, + } + continue; + } + + tokio::select! { + result = tasks.join_next() => { + let Some(Ok((peer, outcome))) = result else { + continue; + }; + if let MetadataFetchOutcome::Fetched(metadata) = outcome { + let canceled = tasks.len(); + runtime_stats.metadata_peer_canceled(canceled); + #[cfg(feature = "metrics")] + counter!("dht_metadata_peer_canceled_total").increment(canceled as u64); + peer_rx.close(); + tasks.abort_all(); + while tasks.join_next().await.is_some() {} + return Some((peer, metadata)); + } + } + peer = peer_rx.recv(), if accepting_live_peers => { + match peer { + Some(peer) => { + let fetch = fetch.clone(); + tasks.spawn(fetch(peer)); + candidates += 1; + runtime_stats.metadata_peer_candidate(); + #[cfg(feature = "metrics")] + counter!("dht_metadata_peer_candidates_total", "source" => "live") + .increment(1); + accepting_live_peers = candidates < MAX_METADATA_PEERS_PER_HASH; + } + None => accepting_live_peers = false, + } + } + } + } +} + #[cfg(test)] mod tests { use super::*; + use std::sync::atomic::AtomicBool; + + struct DropSignal(Arc); + + impl Drop for DropSignal { + fn drop(&mut self) { + self.0.store(true, Ordering::Relaxed); + } + } fn empty_torrent_callback() -> Arc> { Arc::new(ArcSwapOption::empty()) @@ -749,22 +913,27 @@ mod tests { ); let start = Instant::now(); let mut queue = PendingHashQueue::new(2, HASH_QUEUE_TTL); + let mut in_flight = HashMap::new(); - scheduler.enqueue(&mut queue, event("old", 1000, start)); + scheduler.enqueue(&mut queue, &mut in_flight, event("old", 1000, start)); scheduler.enqueue( &mut queue, + &mut in_flight, event("old", 1001, start + Duration::from_secs(1)), ); scheduler.enqueue( &mut queue, + &mut in_flight, event("middle", 1002, start + Duration::from_secs(2)), ); scheduler.enqueue( &mut queue, + &mut in_flight, event("new", 1003, start + Duration::from_secs(3)), ); scheduler.enqueue( &mut queue, + &mut in_flight, event( "stale", 1004, @@ -786,10 +955,10 @@ mod tests { } #[test] - fn queue_deduplicates_hash_and_keeps_newest_unique_peers() { + fn queue_deduplicates_hash_and_keeps_twelve_newest_unique_peers() { let start = Instant::now(); let mut queue = PendingHashQueue::new(10, Duration::from_secs(60)); - for offset in 0..4 { + for offset in 0..13 { assert_ne!( queue.push( event( @@ -803,14 +972,224 @@ mod tests { ); } queue.push( - event("hash", 1003, start + Duration::from_secs(5)), - start + Duration::from_secs(5), + event("hash", 1012, start + Duration::from_secs(20)), + start + Duration::from_secs(20), ); assert_eq!(queue.len(), 1); let entry = queue.remove("hash").unwrap(); let ports: Vec<_> = entry.peers.iter().map(|peer| peer.addr.port()).collect(); - assert_eq!(ports, vec![1003, 1002, 1001]); + assert_eq!(ports, (1001..=1012).rev().collect::>()); + } + + #[test] + fn coalescing_window_is_fixed_from_first_enqueue() { + let start = Instant::now(); + let mut queue = PendingHashQueue::new(10, Duration::from_secs(60)); + let in_flight = HashMap::new(); + queue.push(event("hash", 1000, start), start); + queue.push( + event("hash", 1001, start + Duration::from_millis(15)), + start + Duration::from_millis(15), + ); + + assert!( + queue + .pop_newest_ready(start + Duration::from_millis(24), &in_flight) + .is_none() + ); + let entry = queue + .pop_newest_ready(start + Duration::from_millis(25), &in_flight) + .expect("duplicate arrivals must not extend the coalescing deadline"); + assert_eq!(entry.peers.len(), 2); + } + + #[test] + fn in_flight_hash_routes_at_most_twelve_unique_peers_to_live_race() { + let (_hash_tx, hash_rx) = mpsc::channel(4); + let stats = DhtRuntimeStats::default(); + let scheduler = MetadataScheduler::new_with_runtime_stats( + hash_rx, + Arc::new(RbitFetcher::new(1)), + MetadataSchedulerLimits { + queue_size: 4, + concurrency: 1, + }, + MetadataSchedulerCallbacks { + torrent: empty_torrent_callback(), + fetch_gate: Arc::new(ArcSwapOption::empty()), + completion: Arc::new(ArcSwapOption::empty()), + }, + Arc::new(AtomicUsize::new(0)), + CancellationToken::new(), + stats.clone(), + ); + let (peer_tx, mut peer_rx) = mpsc::channel(MAX_METADATA_PEERS_PER_HASH); + let mut in_flight = HashMap::from([( + "hash".to_string(), + InFlightState { + peer_tx, + scheduled_peers: HashSet::from([SocketAddr::from(([127, 0, 0, 1], 1000))]), + }, + )]); + let mut queue = PendingHashQueue::new(4, HASH_QUEUE_TTL); + let now = Instant::now(); + + for port in 1001..=1013 { + scheduler.enqueue(&mut queue, &mut in_flight, event("hash", port, now)); + } + scheduler.enqueue(&mut queue, &mut in_flight, event("hash", 1001, now)); + + let mut received = Vec::new(); + while let Ok(peer) = peer_rx.try_recv() { + received.push(peer.addr.port()); + } + assert_eq!(received, (1001..=1011).collect::>()); + assert_eq!( + in_flight["hash"].scheduled_peers.len(), + MAX_METADATA_PEERS_PER_HASH + ); + assert!(queue.is_empty()); + let snapshot = stats.snapshot(); + assert_eq!(snapshot.metadata_live_peer_joins, 11); + assert_eq!(snapshot.metadata_queue_deduplicated, 14); + } + + #[tokio::test] + async fn peer_race_returns_first_success_and_cancels_remaining_fetches() { + let now = Instant::now(); + let peers = vec![ + PeerCandidate { + addr: SocketAddr::from(([127, 0, 0, 1], 1000)), + discovered_at: now, + }, + PeerCandidate { + addr: SocketAddr::from(([127, 0, 0, 1], 1001)), + discovered_at: now, + }, + PeerCandidate { + addr: SocketAddr::from(([127, 0, 0, 1], 1002)), + discovered_at: now, + }, + ]; + let started = Arc::new(AtomicUsize::new(0)); + let cancelled = Arc::new(AtomicBool::new(false)); + let barrier = Arc::new(tokio::sync::Barrier::new(peers.len())); + let started_for_race = started.clone(); + let cancelled_for_race = cancelled.clone(); + let (_peer_tx, mut peer_rx) = mpsc::channel(MAX_METADATA_PEERS_PER_HASH); + let stats = DhtRuntimeStats::default(); + + let result = race_peer_fetches( + peers, + &mut peer_rx, + move |peer| { + let started = started_for_race.clone(); + let cancelled = cancelled_for_race.clone(); + let barrier = barrier.clone(); + async move { + started.fetch_add(1, Ordering::Relaxed); + barrier.wait().await; + match peer.addr.port() { + 1000 => { + tokio::time::sleep(Duration::from_millis(10)).await; + ( + peer, + MetadataFetchOutcome::Fetched(( + "winner".to_string(), + 1, + Vec::new(), + 0, + )), + ) + } + 1001 => { + tokio::time::sleep(Duration::from_secs(1)).await; + (peer, MetadataFetchOutcome::Failed) + } + _ => { + let _signal = DropSignal(cancelled); + std::future::pending::<()>().await; + unreachable!() + } + } + } + }, + &stats, + ) + .await + .expect("one peer should win the race"); + + assert_eq!(result.0.addr.port(), 1000); + assert_eq!(result.1.0, "winner"); + assert!(cancelled.load(Ordering::Relaxed)); + let snapshot = stats.snapshot(); + assert_eq!(snapshot.metadata_races_started, 1); + assert_eq!(snapshot.metadata_peer_candidates, 3); + assert_eq!(snapshot.metadata_peer_canceled, 2); + } + + #[tokio::test] + async fn live_peer_joins_running_race_and_cancels_slow_initial_peer() { + let now = Instant::now(); + let initial = vec![PeerCandidate { + addr: SocketAddr::from(([127, 0, 0, 1], 1000)), + discovered_at: now, + }]; + let live = PeerCandidate { + addr: SocketAddr::from(([127, 0, 0, 1], 1001)), + discovered_at: now, + }; + let (peer_tx, mut peer_rx) = mpsc::channel(MAX_METADATA_PEERS_PER_HASH); + let initial_started = Arc::new(tokio::sync::Notify::new()); + let initial_started_for_fetch = initial_started.clone(); + let initial_cancelled = Arc::new(AtomicBool::new(false)); + let initial_cancelled_for_fetch = initial_cancelled.clone(); + let stats = DhtRuntimeStats::default(); + let sender = tokio::spawn(async move { + initial_started.notified().await; + peer_tx.send(live).await.unwrap(); + }); + + let result = tokio::time::timeout( + Duration::from_secs(1), + race_peer_fetches( + initial, + &mut peer_rx, + move |peer| { + let initial_started = initial_started_for_fetch.clone(); + let initial_cancelled = initial_cancelled_for_fetch.clone(); + async move { + if peer.addr.port() == 1000 { + let _signal = DropSignal(initial_cancelled); + initial_started.notify_one(); + std::future::pending::<()>().await; + unreachable!() + } + ( + peer, + MetadataFetchOutcome::Fetched(( + "live-winner".to_string(), + 1, + Vec::new(), + 0, + )), + ) + } + }, + &stats, + ), + ) + .await + .unwrap() + .expect("live peer should win"); + sender.await.unwrap(); + + assert_eq!(result.0.addr.port(), 1001); + assert!(initial_cancelled.load(Ordering::Relaxed)); + let snapshot = stats.snapshot(); + assert_eq!(snapshot.metadata_peer_candidates, 2); + assert_eq!(snapshot.metadata_peer_canceled, 1); } #[test] @@ -844,7 +1223,14 @@ mod tests { ); let in_flight = HashMap::new(); - let newest = queue.pop_newest_available(&in_flight).unwrap(); + assert!( + queue + .pop_newest_ready(start + Duration::from_millis(24), &in_flight) + .is_none() + ); + let newest = queue + .pop_newest_ready(start + Duration::from_secs(2), &in_flight) + .unwrap(); assert_eq!(newest.info_hash, "newer"); assert_eq!(queue.expire(start + Duration::from_secs(11)), 1); assert!(queue.is_empty()); diff --git a/src/server.rs b/src/server.rs index d52f354..1142371 100644 --- a/src/server.rs +++ b/src/server.rs @@ -6,11 +6,15 @@ use crate::error::Result; use crate::krpc::encode_response; use crate::metadata::RbitFetcher; use crate::node_id::{neighbor_node_id, random_node_id, transaction_id_from_bytes}; +use crate::peer_lookup::{ + PeerLookupHandle, PeerLookupRuntime, is_peer_lookup_tid, spawn_peer_lookup, +}; use crate::protocol::{DhtArgs, DhtMessage}; use crate::runtime_stats::{DhtRuntimeLimits, DhtRuntimeStats}; use crate::scheduler::{ MetadataCompletionCallback, MetadataFetchCallback, MetadataScheduler, - MetadataSchedulerCallbacks, MetadataSchedulerLimits, TorrentAckCallback, + MetadataSchedulerCallbacks, MetadataSchedulerLimits, MetadataSchedulerRuntime, + TorrentAckCallback, }; use crate::types::{DHTOptions, MetadataFetchCompletion, NetMode, NodeTuple, TorrentInfo}; use crate::udp_buffer::UdpBufferPool; @@ -252,6 +256,7 @@ pub struct DHTServer { metadata_completion_callback: Arc>, error_callback: Arc>, crawl_engine: Arc, + peer_lookup: PeerLookupHandle, hash_events_tx: mpsc::Sender, metadata_queue_len: Arc, max_metadata_queue_size: usize, @@ -352,8 +357,24 @@ impl DHTServer { let metadata_completion_callback = Arc::new(ArcSwapOption::empty()); let metadata_queue_len = Arc::new(AtomicUsize::new(0)); let shutdown = CancellationToken::new(); + let crawl_engine = Arc::new(CrawlEngine::new( + crawl_config.clone(), + runtime_stats.clone(), + )); + let peer_lookup = spawn_peer_lookup( + options.netmode, + node_id, + &sockets_by_bind_addr, + crawl_engine.snapshot.clone(), + hash_events_tx.clone(), + PeerLookupRuntime { + options: options.peer_lookup.clone(), + stats: runtime_stats.clone(), + shutdown: shutdown.clone(), + }, + ); - let scheduler = MetadataScheduler::new_with_runtime_stats( + let scheduler = MetadataScheduler::new_with_runtime( hash_rx, fetcher, MetadataSchedulerLimits { @@ -367,7 +388,10 @@ impl DHTServer { }, metadata_queue_len.clone(), shutdown.clone(), - runtime_stats.clone(), + MetadataSchedulerRuntime { + stats: runtime_stats.clone(), + peer_lookup_tx: Some(peer_lookup.request_sender()), + }, ); tokio::spawn(scheduler.run()); @@ -383,7 +407,8 @@ impl DHTServer { on_metadata_fetch, metadata_completion_callback, error_callback: Arc::new(ArcSwapOption::empty()), - crawl_engine: Arc::new(CrawlEngine::new(crawl_config, runtime_stats.clone())), + crawl_engine, + peer_lookup, hash_events_tx, metadata_queue_len, max_metadata_queue_size, @@ -650,7 +675,11 @@ impl DHTServer { if let Some(response) = msg.r && let Some(tid) = transaction_id_from_bytes(&msg.t) { - self.crawl_engine.route_response(remote_addr, tid, response); + if is_peer_lookup_tid(&tid) { + self.peer_lookup.route_response(remote_addr, tid, response); + } else { + self.crawl_engine.route_response(remote_addr, tid, response); + } } } _ => {} diff --git a/src/types.rs b/src/types.rs index 4195d2a..94fe5c4 100644 --- a/src/types.rs +++ b/src/types.rs @@ -117,6 +117,8 @@ pub struct DHTOptions { pub hash_queue_capacity: usize, /// Metadata download and Peer-cache limits. pub metadata: MetadataOptions, + /// Active get_peers lookup rate and concurrency limits. + pub peer_lookup: PeerLookupOptions, /// Active crawl, node-pool and scheduler limits. pub crawl: CrawlOptions, } @@ -136,6 +138,17 @@ pub struct MetadataOptions { pub peer_failure_ttl_secs: u64, } +#[derive(Debug, Clone)] +/// Active get_peers lookup budgets used to discover additional Metadata Peers. +pub struct PeerLookupOptions { + /// Maximum new InfoHash lookups started per second. Zero disables active lookup. + pub max_lookups_per_second: u32, + /// Maximum lookup budget consumed immediately after an idle period. + pub burst: u32, + /// Maximum InfoHash lookups kept active at the same time. + pub max_active_lookups: usize, +} + #[derive(Debug, Clone, Default)] /// Active crawl configuration grouped by responsibility. pub struct CrawlOptions { @@ -243,6 +256,7 @@ impl Default for DHTOptions { netmode: NetMode::Ipv4Only, hash_queue_capacity: 10_000, metadata: MetadataOptions::default(), + peer_lookup: PeerLookupOptions::default(), crawl: CrawlOptions::default(), } } @@ -260,6 +274,16 @@ impl Default for MetadataOptions { } } +impl Default for PeerLookupOptions { + fn default() -> Self { + Self { + max_lookups_per_second: 32, + burst: 32, + max_active_lookups: 64, + } + } +} + impl Default for RateLimitOptions { fn default() -> Self { Self {