From b9b157720d21a0646ead171e2a1e7903201241c0 Mon Sep 17 00:00:00 2001 From: chuan Date: Wed, 5 Aug 2026 23:28:31 +0800 Subject: [PATCH] feat: implement BEP-51 InfoHash sampling and related metrics - Added new metrics for BEP-51 InfoHash sampling in `metrics.md`. - Enhanced logging in `main.rs` to include BEP-51 sampling statistics. - Introduced new fields in `crawl_engine.rs` to manage sampling state. - Implemented encoding for BEP-51 sampling queries in `krpc.rs`. - Created a new module `sample_infohashes.rs` to handle BEP-51 sampling logic. - Updated `peer_lookup.rs` to support requests for sampled InfoHashes. - Modified `scheduler.rs` to send `PeerLookupRequest` instead of raw info hashes. - Integrated BEP-51 sampling into the DHT server in `server.rs`. - Added configuration options for BEP-51 sampling in `types.rs`. - Updated runtime statistics to track BEP-51 sampling metrics in `runtime_stats.rs`. --- CHANGELOG.md | 15 ++ README.md | 10 +- docs/metrics.md | 11 ++ examples/main.rs | 11 +- src/crawl_engine.rs | 3 + src/krpc.rs | 33 ++++ src/lib.rs | 7 +- src/peer_lookup.rs | 137 +++++++++---- src/protocol.rs | 9 + src/runtime_stats.rs | 114 +++++++++++ src/sample_infohashes.rs | 410 +++++++++++++++++++++++++++++++++++++++ src/scheduler.rs | 16 +- src/server.rs | 24 ++- src/types.rs | 45 ++++- 14 files changed, 790 insertions(+), 55 deletions(-) create mode 100644 src/sample_infohashes.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 1b8b69d..fcd1299 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,21 @@ 本项目遵循语义化版本。0.2.1 是包含公开 API 变更的 breaking release。 +## Unreleased + +### Added + +- BEP-51 `sample_infohashes` 主动发现、按节点 interval/退避的采样 actor，以及有界 Hash + 去重。 +- 采样 Hash 优先查询来源节点，再通过迭代式 `get_peers` 补充 Peer。 + +### Changed + +- 主动 `get_peers` ingress 改为有界排队，避免突发采样在速率预算耗尽时直接丢弃。 +- 默认主动 Peer lookup 提升到每秒 128 个、最多 256 个并发 lookup。 +- 空节点池的 Bootstrap 默认改为 30 秒重试、每轮最多 16 个端点，降低坏 DNS + 地址导致冷启动停滞的概率。 + ## 0.2.1 - 2026-07-30 ### Breaking changes diff --git a/README.md b/README.md index 4397b56..0452c1d 100644 --- a/README.md +++ b/README.md @@ -4,13 +4,14 @@ [![Documentation](https://docs.rs/dht-crawler/badge.svg)](https://docs.rs/dht-crawler) [![License](https://img.shields.io/crates/l/dht-crawler.svg)](LICENSE) -基于 Rust 和 Tokio 的 BitTorrent DHT 爬虫库。它参与 BEP-5 DHT 网络，接收 -`announce_peer`，并通过 BEP-9 `ut_metadata` 获取、校验和解析 torrent 元数据。 +基于 Rust 和 Tokio 的 BitTorrent DHT 爬虫库。它参与 BEP-5 DHT 网络，通过 BEP-51 +`sample_infohashes` 主动发现 InfoHash，也接收 `announce_peer`，并通过 BEP-9 +`ut_metadata` 获取、校验和解析 torrent 元数据。 `dht-crawler` 提供： - IPv4、IPv6 和双栈 DHT； -- 主动节点发现与 `get_peers` 查询； +- 主动节点发现、BEP-51 InfoHash 采样与 `get_peers` 查询； - 有界、去重的 Metadata 下载队列； - InfoHash 过滤、异步准入、结果交付和完成通知； - 默认可用的运行时统计，以及可选的 `metrics` 集成。 @@ -177,6 +178,7 @@ let options = DHTOptions { | `DHTOptions` | 监听端口、网络模式和顶层队列 | | `MetadataOptions` | 下载超时、队列、并发和失败 Peer 缓存 | | `PeerLookupOptions` | 主动 `get_peers` 的速率与并发 | +| `SampleInfohashesOptions` | BEP-51 采样速率、并发、超时、退避和 Hash 去重容量 | | `RateLimitOptions` | `find_node`、在途请求和 UDP 回复预算 | | `PoolOptions` | 节点池、最近探测记录和响应节点缓存 | | `BootstrapOptions` | Bootstrap 节点与失败退避 | @@ -188,8 +190,10 @@ let options = DHTOptions { - `DHTOptions::default()` 使用 `Ipv4Only`； - `NetMode::DualStack` 会分别绑定 IPv4 和 IPv6 Socket； - `DHTServer::new()` 会立即在所有可用接口上绑定配置的 UDP 端口； +- 空节点池默认每 30 秒重新尝试 Bootstrap，每轮最多使用 16 个已解析端点； - `MetadataOptions::timeout_secs` 是单个 Peer 尝试的端到端期限； - `PeerLookupOptions::max_lookups_per_second = 0` 会关闭主动 `get_peers`； +- `SampleInfohashesOptions::max_queries_per_second = 0` 会关闭 BEP-51 主动采样； - Metadata 和爬取队列都是有界的，容量应与下游处理能力一起调整。 ## 数据与运行语义 diff --git a/docs/metrics.md b/docs/metrics.md index 675284a..fdcea15 100644 --- a/docs/metrics.md +++ b/docs/metrics.md @@ -45,6 +45,17 @@ recorder、不监听端口，也不依赖任何特定导出协议；Prometheus e actor 每秒把内部增量 flush 到 counter，因此 exporter 看到的 counter 可能最多延迟约一秒。 +## BEP-51 InfoHash 采样 + +| 指标 | 类型 | 标签 | 含义 | +|---|---|---|---| +| `dht_sample_infohashes_queries_total` | counter | — | 已发送的 BEP-51 查询 | +| `dht_sample_infohashes_responses_total` | counter | — | 匹配的 BEP-51 响应 | +| `dht_sample_infohashes_timeouts_total` | counter | — | 超时的 BEP-51 请求 | +| `dht_sample_infohashes_hashes_total` | counter | — | 已接受并送往 Peer lookup 的新 Hash | +| `dht_sample_infohashes_in_flight` | gauge | — | 当前在途采样请求 | +| `dht_sample_infohashes_dropped_total` | counter | `reason=response_queue_full|peer_lookup_queue_full` | 有界队列丢弃 | + ## announce 与 Metadata ingress | 指标 | 类型 | 标签 | 含义 | diff --git a/examples/main.rs b/examples/main.rs index 55eb06d..280e12a 100644 --- a/examples/main.rs +++ b/examples/main.rs @@ -111,10 +111,19 @@ async fn main() -> Result<()> { // ✅ 监控：爬虫运行状态 log::info!( - "📊 [监控] 时长: {}s | 成功抓取: ✨ {} | 节点: {} | Metadata: {}/{} | worker: {}", + "📊 [监控] 时长: {}s | 成功抓取: ✨ {} | 节点: {} | BEP51: hash={}, resp={}, timeout={} | Lookup: peer={} | Fetch: ok={}, fail={}, connect={}, timeout={}, noext={} | Metadata: {}/{} | worker: {}", uptime, success_fetch, runtime.node_pool_size, + runtime.sample_infohashes_hashes_discovered, + runtime.sample_infohashes_responses, + runtime.sample_infohashes_timeouts, + runtime.peer_lookup_peers_found, + runtime.metadata_peer_succeeded, + runtime.metadata_peer_failed, + runtime.metadata_connect_failed, + runtime.metadata_peer_timeouts, + runtime.metadata_no_extension, runtime.metadata_queue_depth, runtime.metadata_queue_max, runtime.metadata_in_flight, diff --git a/src/crawl_engine.rs b/src/crawl_engine.rs index 9f59e69..1c9238a 100644 --- a/src/crawl_engine.rs +++ b/src/crawl_engine.rs @@ -914,6 +914,9 @@ mod tests { nodes: None, nodes6: None, values: None, + samples: None, + num: None, + interval: 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 6e43f6e..805f7cb 100644 --- a/src/krpc.rs +++ b/src/krpc.rs @@ -118,6 +118,23 @@ pub(crate) fn encode_get_peers_query( buffer.extend_from_slice(b"1:y1:qe"); } +pub(crate) fn encode_sample_infohashes_query( + buffer: &mut BytesMut, + tid: &TransactionId, + target: &[u8; 20], + sender_id: &[u8; 20], +) { + buffer.clear(); + buffer.reserve(128); + buffer.extend_from_slice(b"d1:ad2:id20:"); + buffer.extend_from_slice(sender_id); + buffer.extend_from_slice(b"6:target20:"); + buffer.extend_from_slice(target); + buffer.extend_from_slice(b"e1:q17:sample_infohashes1:t8:"); + buffer.extend_from_slice(tid); + buffer.extend_from_slice(b"1:y1:qe"); +} + pub(crate) fn encode_response( buffer: &mut BytesMut, tid: &[u8], @@ -240,6 +257,19 @@ mod tests { assert_eq!(message.a.unwrap().info_hash.unwrap().as_ref(), &info_hash); } + #[test] + fn manual_sample_infohashes_encoding_round_trips() { + let mut buffer = BytesMut::new(); + let tid = [1; 8]; + let target = [2; 20]; + let sender = [3; 20]; + encode_sample_infohashes_query(&mut buffer, &tid, &target, &sender); + let message: DhtMessage = serde_bencode::from_bytes(&buffer).unwrap(); + assert_eq!(message.t.as_ref(), &tid); + assert_eq!(message.q.as_deref(), Some("sample_infohashes")); + assert_eq!(message.a.unwrap().target.unwrap().as_ref(), &target); + } + #[test] fn manual_response_encoding_round_trips() { let mut buffer = BytesMut::new(); @@ -264,6 +294,9 @@ mod tests { serde_bytes::ByteBuf::from(vec![10, 0, 0, 1, 0x1a, 0xe1]), serde_bytes::ByteBuf::from(vec![1, 2, 3]), ]), + samples: None, + num: None, + interval: None, }; let mut peers = Vec::new(); assert_eq!( diff --git a/src/lib.rs b/src/lib.rs index 7399e83..bc2e600 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -21,6 +21,7 @@ mod peer_lookup; pub mod protocol; mod routing_snapshot; mod runtime_stats; +mod sample_infohashes; /// Bounded, deduplicating Metadata scheduler. pub mod scheduler; mod server; @@ -38,7 +39,8 @@ pub use server::{DHTServer, HashDiscovered}; pub use types::{ BootstrapOptions, CrawlOptions, DHTOptions, FileInfo, MetadataFetchCompletion, MetadataFetchCompletionStatus, MetadataOptions, NetMode, NodeTuple, PeerLookupOptions, - PoolOptions, RateLimitOptions, SchedulerOptions, TargetOptions, TorrentInfo, + PoolOptions, RateLimitOptions, SampleInfohashesOptions, SchedulerOptions, TargetOptions, + TorrentInfo, }; /// Common server, configuration and callback payload imports. @@ -52,6 +54,7 @@ pub mod prelude { pub use crate::types::{ BootstrapOptions, CrawlOptions, DHTOptions, FileInfo, MetadataFetchCompletion, MetadataFetchCompletionStatus, MetadataOptions, NetMode, NodeTuple, PeerLookupOptions, - PoolOptions, RateLimitOptions, SchedulerOptions, TargetOptions, TorrentInfo, + PoolOptions, RateLimitOptions, SampleInfohashesOptions, SchedulerOptions, TargetOptions, + TorrentInfo, }; } diff --git a/src/peer_lookup.rs b/src/peer_lookup.rs index 53afa54..9ec8780 100644 --- a/src/peer_lookup.rs +++ b/src/peer_lookup.rs @@ -27,7 +27,7 @@ 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 REQUEST_CHANNEL_CAPACITY: usize = 16_384; const RESPONSE_CHANNEL_CAPACITY: usize = 4_096; #[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)] @@ -46,6 +46,7 @@ struct LookupState { info_hash: [u8; 20], info_hash_hex: String, frontier: Vec, + preferred: Option, seen_nodes: AHashSet, peers: AHashSet, queried: usize, @@ -55,6 +56,9 @@ struct LookupState { impl LookupState { fn pop_closest(&mut self) -> Option { + if let Some(preferred) = self.preferred.take() { + return Some(preferred); + } let index = self .frontier .iter() @@ -78,9 +82,24 @@ struct LookupResponse { response: DhtResponse, } +#[derive(Debug, Clone, Copy)] +pub(crate) struct PeerLookupRequest { + pub(crate) info_hash: [u8; 20], + pub(crate) preferred_node: Option, +} + +impl PeerLookupRequest { + pub(crate) fn new(info_hash: [u8; 20]) -> Self { + Self { + info_hash, + preferred_node: None, + } + } +} + #[derive(Clone)] pub(crate) struct PeerLookupHandle { - request_tx: mpsc::Sender<[u8; 20]>, + request_tx: mpsc::Sender, response_tx: mpsc::Sender, runtime_stats: DhtRuntimeStats, } @@ -92,7 +111,7 @@ pub(crate) struct PeerLookupRuntime { } impl PeerLookupHandle { - pub(crate) fn request_sender(&self) -> mpsc::Sender<[u8; 20]> { + pub(crate) fn request_sender(&self) -> mpsc::Sender { self.request_tx.clone() } @@ -160,6 +179,10 @@ pub(crate) fn spawn_peer_lookup( Instant::now(), ), max_active_lookups: options.max_active_lookups, + enabled: options.max_lookups_per_second > 0 && options.max_active_lookups > 0, + queued: VecDeque::new(), + queued_hashes: AHashSet::new(), + active_hashes: AHashSet::new(), active: AHashMap::new(), pending: AHashMap::new(), pending_expiry: VecDeque::new(), @@ -183,10 +206,14 @@ struct PeerLookupActor { socket_v6: Option>, snapshot: Arc>, hash_tx: mpsc::Sender, - request_rx: mpsc::Receiver<[u8; 20]>, + request_rx: mpsc::Receiver, response_rx: mpsc::Receiver, request_budget: RateBucket, max_active_lookups: usize, + enabled: bool, + queued: VecDeque, + queued_hashes: AHashSet<[u8; 20]>, + active_hashes: AHashSet<[u8; 20]>, active: AHashMap, pending: AHashMap, pending_expiry: VecDeque<(Instant, PendingKey)>, @@ -210,47 +237,83 @@ impl PeerLookupActor { } _ = 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; + let Some(request) = request else { break }; + self.queue_request(request); } } } } - async fn start_lookup(&mut self, info_hash: [u8; 20], now: Instant) { + fn queue_request(&mut self, request: PeerLookupRequest) { 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); + if !self.enabled { return; } + if self.queued_hashes.contains(&request.info_hash) + || self.active_hashes.contains(&request.info_hash) + { + return; + } + if self.queued.len() >= REQUEST_CHANNEL_CAPACITY { + self.runtime_stats.peer_lookup_rate_limited(); + #[cfg(feature = "metrics")] + counter!("dht_peer_lookup_dropped_total", "reason" => "request_queue_full") + .increment(1); + return; + } + self.queued_hashes.insert(request.info_hash); + self.queued.push_back(request); + } + + async fn start_queued(&mut self, now: Instant) { + while self.active.len() < self.max_active_lookups + && !self.queued.is_empty() + && self.request_budget.try_take_one(now) + { + let request = self.queued.pop_front().expect("queued request exists"); + self.queued_hashes.remove(&request.info_hash); + if !self.start_lookup(request, now).await { + self.request_budget.refund_one(); + } + } + } + + async fn start_lookup(&mut self, request: PeerLookupRequest, now: Instant) -> bool { + let info_hash = request.info_hash; 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 = + let mut frontier = self.snapshot .load() .closest_nodes(&info_hash, MAX_QUERIES_PER_LOOKUP, filter_ipv6); - if frontier.is_empty() { + let preferred = request.preferred_node; + if let Some(preferred) = preferred { + frontier.retain(|node| node.addr != preferred.addr); + } + if frontier.is_empty() && preferred.is_none() { self.runtime_stats.peer_lookup_empty(); - return; + return false; } 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(); + let mut seen_nodes: AHashSet<_> = frontier.iter().map(|node| node.addr).collect(); + if let Some(preferred) = preferred { + seen_nodes.insert(preferred.addr); + } self.active.insert( lookup_id, LookupState { info_hash, info_hash_hex: hex::encode(info_hash), frontier, + preferred, seen_nodes, peers: AHashSet::new(), queried: 0, @@ -258,29 +321,27 @@ impl PeerLookupActor { deadline: now + LOOKUP_TIMEOUT, }, ); + self.active_hashes.insert(info_hash); self.runtime_stats.peer_lookup_started(); #[cfg(feature = "metrics")] counter!("dht_peer_lookup_started_total").increment(1); self.dispatch_more(lookup_id, now).await; + true } 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; - }; - + while 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)) + }) { let tid = self.next_transaction_id(); let key = PendingKey { addr: node.addr, @@ -410,6 +471,7 @@ impl PeerLookupActor { for lookup_id in expired { self.finish_lookup(lookup_id); } + self.start_queued(now).await; } fn finish_if_complete(&mut self, lookup_id: u64, now: Instant) { @@ -423,7 +485,9 @@ impl PeerLookupActor { } fn finish_lookup(&mut self, lookup_id: u64) { - self.active.remove(&lookup_id); + if let Some(state) = self.active.remove(&lookup_id) { + self.active_hashes.remove(&state.info_hash); + } self.pending .retain(|_, pending| pending.lookup_id != lookup_id); } @@ -477,7 +541,11 @@ mod tests { shutdown: shutdown.clone(), }, ); - handle.request_tx.send([3; 20]).await.unwrap(); + handle + .request_tx + .send(PeerLookupRequest::new([3; 20])) + .await + .unwrap(); let mut buffer = [0u8; 512]; let (len, source) = @@ -499,6 +567,9 @@ mod tests { values: Some(vec![serde_bytes::ByteBuf::from(vec![ 8, 8, 4, 4, 0x1a, 0xe1, ])]), + samples: None, + num: None, + interval: None, }, ); diff --git a/src/protocol.rs b/src/protocol.rs index 6a5f3dc..ff291d5 100644 --- a/src/protocol.rs +++ b/src/protocol.rs @@ -51,4 +51,13 @@ pub struct DhtResponse { #[serde(default)] /// Compact Peer endpoints returned by `get_peers`. pub values: Option>, + #[serde(default)] + /// BEP-51 concatenated 20-byte sampled InfoHashes. + pub samples: Option, + #[serde(default)] + /// BEP-51 estimated number of InfoHashes held by the responder. + pub num: Option, + #[serde(default)] + /// BEP-51 requested delay before sampling this node again, in seconds. + pub interval: Option, } diff --git a/src/runtime_stats.rs b/src/runtime_stats.rs index b6a5b12..771df77 100644 --- a/src/runtime_stats.rs +++ b/src/runtime_stats.rs @@ -246,6 +246,22 @@ pub struct DhtRuntimeSnapshot { pub peer_lookup_response_dropped: u64, /// Discovered Peer endpoints dropped because Hash ingress was full. pub peer_lookup_output_dropped: u64, + /// BEP-51 requests sent. + pub sample_infohashes_queries: u64, + /// Matched BEP-51 responses. + pub sample_infohashes_responses: u64, + /// BEP-51 requests that timed out. + pub sample_infohashes_timeouts: u64, + /// BEP-51 UDP sends that failed immediately. + pub sample_infohashes_send_failures: u64, + /// BEP-51 responses dropped before reaching the sampler actor. + pub sample_infohashes_response_dropped: u64, + /// New sampled hashes accepted for Peer lookup. + pub sample_infohashes_hashes_discovered: u64, + /// Sampled hashes rejected by the bounded deduplicator. + pub sample_infohashes_hashes_duplicate: u64, + /// New sampled hashes dropped because Peer lookup ingress was full. + pub sample_infohashes_hashes_dropped: u64, /// Real Peer network attempts. pub metadata_peer_attempts: u64, /// Successful Peer downloads and parses. @@ -340,6 +356,14 @@ struct DhtRuntimeStatsInner { peer_lookup_peers_found: AtomicU64, peer_lookup_response_dropped: AtomicU64, peer_lookup_output_dropped: AtomicU64, + sample_infohashes_queries: AtomicU64, + sample_infohashes_responses: AtomicU64, + sample_infohashes_timeouts: AtomicU64, + sample_infohashes_send_failures: AtomicU64, + sample_infohashes_response_dropped: AtomicU64, + sample_infohashes_hashes_discovered: AtomicU64, + sample_infohashes_hashes_duplicate: AtomicU64, + sample_infohashes_hashes_dropped: AtomicU64, metadata_peer_attempts: AtomicU64, metadata_peer_succeeded: AtomicU64, metadata_peer_failed: AtomicU64, @@ -442,6 +466,14 @@ impl Default for DhtRuntimeStatsInner { peer_lookup_peers_found: AtomicU64::new(0), peer_lookup_response_dropped: AtomicU64::new(0), peer_lookup_output_dropped: AtomicU64::new(0), + sample_infohashes_queries: AtomicU64::new(0), + sample_infohashes_responses: AtomicU64::new(0), + sample_infohashes_timeouts: AtomicU64::new(0), + sample_infohashes_send_failures: AtomicU64::new(0), + sample_infohashes_response_dropped: AtomicU64::new(0), + sample_infohashes_hashes_discovered: AtomicU64::new(0), + sample_infohashes_hashes_duplicate: AtomicU64::new(0), + sample_infohashes_hashes_dropped: AtomicU64::new(0), metadata_peer_attempts: AtomicU64::new(0), metadata_peer_succeeded: AtomicU64::new(0), metadata_peer_failed: AtomicU64::new(0), @@ -569,6 +601,24 @@ impl DhtRuntimeStats { .peer_lookup_response_dropped .load(Ordering::Relaxed), peer_lookup_output_dropped: inner.peer_lookup_output_dropped.load(Ordering::Relaxed), + sample_infohashes_queries: inner.sample_infohashes_queries.load(Ordering::Relaxed), + sample_infohashes_responses: inner.sample_infohashes_responses.load(Ordering::Relaxed), + sample_infohashes_timeouts: inner.sample_infohashes_timeouts.load(Ordering::Relaxed), + sample_infohashes_send_failures: inner + .sample_infohashes_send_failures + .load(Ordering::Relaxed), + sample_infohashes_response_dropped: inner + .sample_infohashes_response_dropped + .load(Ordering::Relaxed), + sample_infohashes_hashes_discovered: inner + .sample_infohashes_hashes_discovered + .load(Ordering::Relaxed), + sample_infohashes_hashes_duplicate: inner + .sample_infohashes_hashes_duplicate + .load(Ordering::Relaxed), + sample_infohashes_hashes_dropped: inner + .sample_infohashes_hashes_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), @@ -807,6 +857,54 @@ impl DhtRuntimeStats { .fetch_add(1, Ordering::Relaxed); } + pub(crate) fn sample_query(&self) { + self.inner + .sample_infohashes_queries + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn sample_response(&self) { + self.inner + .sample_infohashes_responses + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn sample_timeout(&self) { + self.inner + .sample_infohashes_timeouts + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn sample_send_failed(&self) { + self.inner + .sample_infohashes_send_failures + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn sample_response_dropped(&self) { + self.inner + .sample_infohashes_response_dropped + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn sample_hash_discovered(&self) { + self.inner + .sample_infohashes_hashes_discovered + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn sample_hash_duplicate(&self) { + self.inner + .sample_infohashes_hashes_duplicate + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn sample_hash_dropped(&self) { + self.inner + .sample_infohashes_hashes_dropped + .fetch_add(1, Ordering::Relaxed); + } + pub(crate) fn metadata_peer_attempt(&self) { self.inner .metadata_peer_attempts @@ -1133,6 +1231,14 @@ mod tests { writer.peer_lookup_peer_found(); writer.peer_lookup_response_dropped(); writer.peer_lookup_output_dropped(); + writer.sample_query(); + writer.sample_response(); + writer.sample_timeout(); + writer.sample_send_failed(); + writer.sample_response_dropped(); + writer.sample_hash_discovered(); + writer.sample_hash_duplicate(); + writer.sample_hash_dropped(); writer.metadata_peer_attempt(); writer.metadata_peer_succeeded(); writer.metadata_peer_failed(); @@ -1192,6 +1298,14 @@ mod tests { peer_lookup_peers_found: 1, peer_lookup_response_dropped: 1, peer_lookup_output_dropped: 1, + sample_infohashes_queries: 1, + sample_infohashes_responses: 1, + sample_infohashes_timeouts: 1, + sample_infohashes_send_failures: 1, + sample_infohashes_response_dropped: 1, + sample_infohashes_hashes_discovered: 1, + sample_infohashes_hashes_duplicate: 1, + sample_infohashes_hashes_dropped: 1, metadata_peer_attempts: 1, metadata_peer_succeeded: 1, metadata_peer_failed: 1, diff --git a/src/sample_infohashes.rs b/src/sample_infohashes.rs new file mode 100644 index 0000000..ab5ece8 --- /dev/null +++ b/src/sample_infohashes.rs @@ -0,0 +1,410 @@ +use crate::budget::RateBucket; +use crate::crawl_engine::CrawlEngine; +use crate::krpc::{encode_sample_infohashes_query, for_each_response_node}; +use crate::node_id::{TransactionId, random_node_id}; +use crate::peer_lookup::{PeerLookupHandle, PeerLookupRequest}; +use crate::protocol::DhtResponse; +use crate::routing_snapshot::RoutingSnapshot; +use crate::runtime_stats::DhtRuntimeStats; +use crate::types::{NetMode, NodeTuple, SampleInfohashesOptions}; +use ahash::{AHashMap, AHashSet}; +use arc_swap::ArcSwap; +use bytes::BytesMut; +#[cfg(feature = "metrics")] +use metrics::{counter, gauge}; +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 SAMPLE_TID_TAG: u8 = 0x51; +const RESPONSE_CHANNEL_CAPACITY: usize = 4_096; +const MAINTENANCE_INTERVAL: Duration = Duration::from_millis(25); +const PRODUCTIVE_REVISIT: Duration = Duration::from_secs(60); +const MAX_PROTOCOL_INTERVAL: Duration = Duration::from_secs(6 * 60 * 60); + +#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)] +struct PendingKey { + addr: SocketAddr, + tid: TransactionId, +} + +#[derive(Debug, Clone, Copy)] +struct PendingRequest { + deadline: Instant, +} + +struct SampleResponse { + remote_addr: SocketAddr, + tid: TransactionId, + response: DhtResponse, +} + +#[derive(Clone)] +pub(crate) struct SampleInfohashesHandle { + response_tx: mpsc::Sender, + runtime_stats: DhtRuntimeStats, +} + +impl SampleInfohashesHandle { + pub(crate) fn route_response( + &self, + remote_addr: SocketAddr, + tid: TransactionId, + response: DhtResponse, + ) { + if self + .response_tx + .try_send(SampleResponse { + remote_addr, + tid, + response, + }) + .is_err() + { + self.runtime_stats.sample_response_dropped(); + #[cfg(feature = "metrics")] + counter!("dht_sample_infohashes_dropped_total", "reason" => "response_queue_full") + .increment(1); + } + } +} + +pub(crate) fn is_sample_infohashes_tid(tid: &TransactionId) -> bool { + tid[0] == SAMPLE_TID_TAG +} + +pub(crate) struct SampleInfohashesRuntime { + pub(crate) options: SampleInfohashesOptions, + pub(crate) stats: DhtRuntimeStats, + pub(crate) shutdown: CancellationToken, +} + +pub(crate) fn spawn_sample_infohashes( + netmode: NetMode, + local_id: [u8; 20], + sockets: &std::collections::HashMap>, + snapshot: Arc>, + crawl_engine: Arc, + peer_lookup: PeerLookupHandle, + runtime: SampleInfohashesRuntime, +) -> SampleInfohashesHandle { + let SampleInfohashesRuntime { + options, + stats, + shutdown, + } = runtime; + 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 now = Instant::now(); + let actor = SampleInfohashesActor { + netmode, + local_id, + socket_v4, + socket_v6, + snapshot, + crawl_engine, + peer_lookup, + response_rx, + query_budget: RateBucket::per_second( + options.max_queries_per_second, + options.burst, + true, + now, + ), + max_in_flight: options.max_in_flight, + request_timeout: Duration::from_millis(options.request_timeout_millis.max(100)), + unsupported_backoff: Duration::from_secs(options.unsupported_backoff_secs.max(1)), + dedup_capacity: options.dedup_capacity, + seen_hashes: AHashSet::new(), + seen_order: VecDeque::new(), + next_allowed: AHashMap::new(), + pending_addrs: AHashSet::new(), + pending: AHashMap::new(), + pending_expiry: VecDeque::new(), + next_tid: 1, + runtime_stats: stats.clone(), + shutdown, + }; + tokio::spawn(actor.run()); + SampleInfohashesHandle { + response_tx, + runtime_stats: stats, + } +} + +struct SampleInfohashesActor { + netmode: NetMode, + local_id: [u8; 20], + socket_v4: Option>, + socket_v6: Option>, + snapshot: Arc>, + crawl_engine: Arc, + peer_lookup: PeerLookupHandle, + response_rx: mpsc::Receiver, + query_budget: RateBucket, + max_in_flight: usize, + request_timeout: Duration, + unsupported_backoff: Duration, + dedup_capacity: usize, + seen_hashes: AHashSet<[u8; 20]>, + seen_order: VecDeque<[u8; 20]>, + next_allowed: AHashMap, + pending_addrs: AHashSet, + pending: AHashMap, + pending_expiry: VecDeque<(Instant, PendingKey)>, + next_tid: u64, + runtime_stats: DhtRuntimeStats, + shutdown: CancellationToken, +} + +impl SampleInfohashesActor { + 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()); + } + _ = maintenance.tick() => { + let now = Instant::now(); + self.expire(now); + self.dispatch(now).await; + } + } + } + } + + async fn dispatch(&mut self, now: Instant) { + if self.max_in_flight == 0 || self.pending.len() >= self.max_in_flight { + return; + } + let available = self.max_in_flight.saturating_sub(self.pending.len()); + let budget = self.query_budget.try_take(available, now); + if budget == 0 { + 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 candidates = self + .snapshot + .load() + .random_nodes((budget * 8).max(64), filter_ipv6); + let mut sent_count = 0usize; + for node in candidates { + if sent_count >= budget { + break; + } + if self.pending_addrs.contains(&node.addr) + || self + .next_allowed + .get(&node.addr) + .is_some_and(|deadline| *deadline > now) + { + continue; + } + if self.send_query(node, now).await { + sent_count += 1; + } + } + if sent_count < budget { + self.query_budget.refund(budget - sent_count); + } + } + + async fn send_query(&mut self, node: NodeTuple, now: Instant) -> bool { + let socket = if node.addr.is_ipv4() { + self.socket_v4.clone() + } else { + self.socket_v6.clone() + }; + let Some(socket) = socket else { + return false; + }; + let tid = self.next_transaction_id(); + let mut buffer = BytesMut::with_capacity(128); + encode_sample_infohashes_query(&mut buffer, &tid, &random_node_id(), &self.local_id); + if socket.send_to(&buffer, node.addr).await.is_err() { + self.runtime_stats.sample_send_failed(); + self.next_allowed + .insert(node.addr, now + self.unsupported_backoff); + return false; + } + let key = PendingKey { + addr: node.addr, + tid, + }; + let deadline = now + self.request_timeout; + self.pending.insert(key, PendingRequest { deadline }); + self.pending_expiry.push_back((deadline, key)); + self.pending_addrs.insert(node.addr); + self.runtime_stats.udp_sent(buffer.len()); + self.runtime_stats.sample_query(); + #[cfg(feature = "metrics")] + { + counter!("dht_sample_infohashes_queries_total").increment(1); + gauge!("dht_sample_infohashes_in_flight").set(self.pending.len() as f64); + } + true + } + + fn handle_response(&mut self, event: SampleResponse, now: Instant) { + let key = PendingKey { + addr: event.remote_addr, + tid: event.tid, + }; + if self.pending.remove(&key).is_none() { + return; + } + self.pending_addrs.remove(&event.remote_addr); + self.runtime_stats.sample_response(); + + for_each_response_node(&event.response, self.netmode, |node| { + self.crawl_engine.route_discovered(node) + }); + + let responder_id = event + .response + .id + .as_deref() + .and_then(|id| <[u8; 20]>::try_from(id.as_slice()).ok()) + .unwrap_or([0; 20]); + let preferred_node = NodeTuple { + id: responder_id, + addr: event.remote_addr, + }; + let mut discovered = 0usize; + if let Some(samples) = event.response.samples.as_deref() + && samples.len() % 20 == 0 + { + for chunk in samples.chunks_exact(20) { + let hash: [u8; 20] = chunk.try_into().expect("sample hash is 20 bytes"); + if self.seen_hashes.contains(&hash) { + self.runtime_stats.sample_hash_duplicate(); + continue; + } + let request = PeerLookupRequest { + info_hash: hash, + preferred_node: Some(preferred_node), + }; + if self.peer_lookup.request_sender().try_send(request).is_ok() { + self.remember_hash(hash); + discovered += 1; + self.runtime_stats.sample_hash_discovered(); + } else { + self.runtime_stats.sample_hash_dropped(); + #[cfg(feature = "metrics")] + counter!("dht_sample_infohashes_dropped_total", "reason" => "peer_lookup_queue_full") + .increment(1); + } + } + } + + let protocol_interval = Duration::from_secs(event.response.interval.unwrap_or(300)) + .clamp(Duration::from_secs(10), MAX_PROTOCOL_INTERVAL); + let delay = if discovered > 0 { + protocol_interval.min(PRODUCTIVE_REVISIT) + } else { + protocol_interval.saturating_add(self.unsupported_backoff) + }; + self.next_allowed.insert(event.remote_addr, now + delay); + #[cfg(feature = "metrics")] + { + counter!("dht_sample_infohashes_responses_total").increment(1); + counter!("dht_sample_infohashes_hashes_total").increment(discovered as u64); + gauge!("dht_sample_infohashes_in_flight").set(self.pending.len() as f64); + } + } + + fn remember_hash(&mut self, hash: [u8; 20]) { + if self.dedup_capacity == 0 { + return; + } + if self.seen_hashes.insert(hash) { + self.seen_order.push_back(hash); + } + while self.seen_order.len() > self.dedup_capacity { + if let Some(expired) = self.seen_order.pop_front() { + self.seen_hashes.remove(&expired); + } + } + } + + fn expire(&mut self, now: Instant) { + while let Some((deadline, key)) = self.pending_expiry.front().copied() { + if deadline > now { + break; + } + self.pending_expiry.pop_front(); + if !self + .pending + .get(&key) + .is_some_and(|pending| pending.deadline == deadline) + { + continue; + } + self.pending.remove(&key); + self.pending_addrs.remove(&key.addr); + self.next_allowed + .insert(key.addr, now + self.unsupported_backoff); + self.runtime_stats.sample_timeout(); + #[cfg(feature = "metrics")] + counter!("dht_sample_infohashes_timeouts_total").increment(1); + } + #[cfg(feature = "metrics")] + gauge!("dht_sample_infohashes_in_flight").set(self.pending.len() as f64); + } + + fn next_transaction_id(&mut self) -> TransactionId { + let mut tid = self.next_tid.to_be_bytes(); + tid[0] = SAMPLE_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 sample_transaction_ids_have_a_reserved_tag() { + assert!(is_sample_infohashes_tid(&[ + SAMPLE_TID_TAG, + 1, + 2, + 3, + 4, + 5, + 6, + 7 + ])); + assert!(!is_sample_infohashes_tid(&[0; 8])); + } + + #[test] + fn bep51_response_fields_decode() { + let bytes = b"d1:rd2:id20:aaaaaaaaaaaaaaaaaaaa8:intervali60e3:numi2e7:samples40:bbbbbbbbbbbbbbbbbbbbcccccccccccccccccccce1:t8:123456781:y1:re"; + let message: DhtMessage = serde_bencode::from_bytes(bytes).unwrap(); + let response = message.r.unwrap(); + assert_eq!(response.interval, Some(60)); + assert_eq!(response.num, Some(2)); + assert_eq!(response.samples.unwrap().len(), 40); + } +} diff --git a/src/scheduler.rs b/src/scheduler.rs index 128930b..4d10064 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -1,4 +1,5 @@ use crate::metadata::{FetchedMetadata, MetadataFetchOutcome, RbitFetcher}; +use crate::peer_lookup::PeerLookupRequest; #[cfg(test)] use crate::runtime_stats::DhtRuntimeLimits; use crate::runtime_stats::DhtRuntimeStats; @@ -60,7 +61,7 @@ pub struct MetadataSchedulerCallbacks { pub(crate) struct MetadataSchedulerRuntime { pub(crate) stats: DhtRuntimeStats, - pub(crate) peer_lookup_tx: Option>, + pub(crate) peer_lookup_tx: Option>, } #[derive(Debug, Clone, Copy)] @@ -149,7 +150,7 @@ impl PendingHashQueue { }); entry .peers - .sort_unstable_by(|left, right| right.discovered_at.cmp(&left.discovered_at)); + .sort_unstable_by_key(|peer| std::cmp::Reverse(peer.discovered_at)); 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); @@ -212,10 +213,7 @@ impl PendingHashQueue { fn expire(&mut self, now: Instant) -> usize { let mut expired = 0; - loop { - let Some((&oldest_key, oldest_hash)) = self.order.first_key_value() else { - break; - }; + while let Some((&oldest_key, oldest_hash)) = self.order.first_key_value() { if now.checked_duration_since(oldest_key.0).unwrap_or_default() <= self.ttl { break; } @@ -269,7 +267,7 @@ pub struct MetadataScheduler { total_completed: Arc, queue_len: Arc, runtime_stats: DhtRuntimeStats, - peer_lookup_tx: Option>, + peer_lookup_tx: Option>, shutdown: CancellationToken, } @@ -630,7 +628,7 @@ impl MetadataScheduler { callback: &Arc>, on_metadata_fetch: &Arc>, runtime_stats: &DhtRuntimeStats, - peer_lookup_tx: Option>, + peer_lookup_tx: Option>, ) -> JobResult { if let Some(gate) = on_metadata_fetch.load_full() && !gate(job.info_hash.clone()).await @@ -663,7 +661,7 @@ impl MetadataScheduler { 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 _ = peer_lookup_tx.try_send(PeerLookupRequest::new(info_hash_bytes)); }) }); let fetched = race_peer_fetches( diff --git a/src/server.rs b/src/server.rs index 1142371..fed4df5 100644 --- a/src/server.rs +++ b/src/server.rs @@ -11,6 +11,10 @@ use crate::peer_lookup::{ }; use crate::protocol::{DhtArgs, DhtMessage}; use crate::runtime_stats::{DhtRuntimeLimits, DhtRuntimeStats}; +use crate::sample_infohashes::{ + SampleInfohashesHandle, SampleInfohashesRuntime, is_sample_infohashes_tid, + spawn_sample_infohashes, +}; use crate::scheduler::{ MetadataCompletionCallback, MetadataFetchCallback, MetadataScheduler, MetadataSchedulerCallbacks, MetadataSchedulerLimits, MetadataSchedulerRuntime, @@ -257,6 +261,7 @@ pub struct DHTServer { error_callback: Arc>, crawl_engine: Arc, peer_lookup: PeerLookupHandle, + sample_infohashes: SampleInfohashesHandle, hash_events_tx: mpsc::Sender, metadata_queue_len: Arc, max_metadata_queue_size: usize, @@ -373,6 +378,19 @@ impl DHTServer { shutdown: shutdown.clone(), }, ); + let sample_infohashes = spawn_sample_infohashes( + options.netmode, + node_id, + &sockets_by_bind_addr, + crawl_engine.snapshot.clone(), + crawl_engine.clone(), + peer_lookup.clone(), + SampleInfohashesRuntime { + options: options.sample_infohashes.clone(), + stats: runtime_stats.clone(), + shutdown: shutdown.clone(), + }, + ); let scheduler = MetadataScheduler::new_with_runtime( hash_rx, @@ -409,6 +427,7 @@ impl DHTServer { error_callback: Arc::new(ArcSwapOption::empty()), crawl_engine, peer_lookup, + sample_infohashes, hash_events_tx, metadata_queue_len, max_metadata_queue_size, @@ -675,7 +694,10 @@ impl DHTServer { if let Some(response) = msg.r && let Some(tid) = transaction_id_from_bytes(&msg.t) { - if is_peer_lookup_tid(&tid) { + if is_sample_infohashes_tid(&tid) { + self.sample_infohashes + .route_response(remote_addr, tid, response); + } else 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 94fe5c4..18b81e9 100644 --- a/src/types.rs +++ b/src/types.rs @@ -119,6 +119,8 @@ pub struct DHTOptions { pub metadata: MetadataOptions, /// Active get_peers lookup rate and concurrency limits. pub peer_lookup: PeerLookupOptions, + /// Active BEP-51 InfoHash sampling limits. + pub sample_infohashes: SampleInfohashesOptions, /// Active crawl, node-pool and scheduler limits. pub crawl: CrawlOptions, } @@ -149,6 +151,23 @@ pub struct PeerLookupOptions { pub max_active_lookups: usize, } +#[derive(Debug, Clone)] +/// Active BEP-51 `sample_infohashes` discovery limits. +pub struct SampleInfohashesOptions { + /// Maximum BEP-51 queries started per second. Zero disables sampling. + pub max_queries_per_second: u32, + /// Maximum sampling budget consumed immediately after an idle period. + pub burst: u32, + /// Maximum outstanding BEP-51 requests. + pub max_in_flight: usize, + /// Per-request timeout in milliseconds. + pub request_timeout_millis: u64, + /// Retry delay for timeouts or nodes that do not return samples, in seconds. + pub unsupported_backoff_secs: u64, + /// Maximum sampled InfoHashes retained for bounded in-memory deduplication. + pub dedup_capacity: usize, +} + #[derive(Debug, Clone, Default)] /// Active crawl configuration grouped by responsibility. pub struct CrawlOptions { @@ -257,6 +276,7 @@ impl Default for DHTOptions { hash_queue_capacity: 10_000, metadata: MetadataOptions::default(), peer_lookup: PeerLookupOptions::default(), + sample_infohashes: SampleInfohashesOptions::default(), crawl: CrawlOptions::default(), } } @@ -277,9 +297,22 @@ impl Default for MetadataOptions { impl Default for PeerLookupOptions { fn default() -> Self { Self { - max_lookups_per_second: 32, - burst: 32, - max_active_lookups: 64, + max_lookups_per_second: 128, + burst: 128, + max_active_lookups: 256, + } + } +} + +impl Default for SampleInfohashesOptions { + fn default() -> Self { + Self { + max_queries_per_second: 16, + burst: 16, + max_in_flight: 64, + request_timeout_millis: 1_500, + unsupported_backoff_secs: 300, + dedup_capacity: 1_000_000, } } } @@ -323,9 +356,9 @@ impl Default for BootstrapOptions { "router.utorrent.com:6881".to_string(), "dht.aelitis.com:6881".to_string(), ], - interval_secs: 300, - max_nodes_per_round: 3, - source_backoff_base_secs: 300, + interval_secs: 30, + max_nodes_per_round: 16, + source_backoff_base_secs: 60, source_backoff_max_secs: 3_600, } }