tools/dht
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: improve metadata peer discovery

Browse Source
This commit is contained in:
桥下红药
2026-07-25 01:53:04 +08:00
parent 7c79bff593
commit fdf474f781
12 changed files with 1446 additions and 117 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+4 -1
View File
@@ -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
src/crawl_engine.rs
+1
View File
@@ -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());
src/krpc.rs
+85 -1
View File
@@ -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());
}
}
src/lib.rs
+5 -4
View File
@@ -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,
};
}
src/metadata.rs
+15
View File
@@ -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<F>(
&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")]
{
src/peer_lookup.rs
+518
View File
@@ -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<NodeTuple>,
seen_nodes: AHashSet<SocketAddr>,
peers: AHashSet<SocketAddr>,
queried: usize,
outstanding: usize,
deadline: Instant,
}
impl LookupState {
fn pop_closest(&mut self) -> Option<NodeTuple> {
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<LookupResponse>,
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<SocketAddr, Arc<UdpSocket>>,
snapshot: Arc<ArcSwap<RoutingSnapshot>>,
hash_tx: mpsc::Sender<HashDiscovered>,
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<Arc<UdpSocket>>,
socket_v6: Option<Arc<UdpSocket>>,
snapshot: Arc<ArcSwap<RoutingSnapshot>>,
hash_tx: mpsc::Sender<HashDiscovered>,
request_rx: mpsc::Receiver<[u8; 20]>,
response_rx: mpsc::Receiver<LookupResponse>,
request_budget: RateBucket,
max_active_lookups: usize,
active: AHashMap<u64, LookupState>,
pending: AHashMap<PendingKey, PendingQuery>,
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();
}
}
src/protocol.rs
+3
View File
@@ -48,4 +48,7 @@ pub struct DhtResponse {
#[serde(default)]
/// Compact IPv6 node tuples.
pub nodes6: Option<serde_bytes::ByteBuf>,
#[serde(default)]
/// Compact Peer endpoints returned by `get_peers`.
pub values: Option<Vec<serde_bytes::ByteBuf>>,
}
src/routing_snapshot.rs
+70
View File
@@ -34,4 +34,74 @@ impl RoutingSnapshot {
}
}
}
pub(crate) fn closest_nodes(
&self,
target: &[u8; 20],
count: usize,
filter_ipv6: Option<bool>,
) -> Vec<NodeTuple> {
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<_>>(),
vec![2, 3]
);
}
}
src/runtime_stats.rs
+195
View File
@@ -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,
src/scheduler.rs
+492 -106
View File
@@ -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<dyn Fn(TorrentInfo) -> bool + Send + Sync + 'static>;
@@ -52,6 +58,11 @@ pub struct MetadataSchedulerCallbacks {
pub completion: Arc<ArcSwapOption<MetadataCompletionCallback>>,
}
pub(crate) struct MetadataSchedulerRuntime {
pub(crate) stats: DhtRuntimeStats,
pub(crate) peer_lookup_tx: Option<mpsc::Sender<[u8; 20]>>,
}
#[derive(Debug, Clone, Copy)]
struct PeerCandidate {
addr: SocketAddr,
@@ -62,6 +73,8 @@ struct PeerCandidate {
struct QueuedHash {
info_hash: String,
peers: Vec<PeerCandidate>,
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<String, InFlightState>,
) -> Option<QueuedHash> {
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<PeerCandidate>,
already_admitted: bool,
peer_rx: mpsc::Receiver<PeerCandidate>,
}
#[derive(Debug)]
struct InFlightState {
attempts: usize,
peer_tx: mpsc::Sender<PeerCandidate>,
scheduled_peers: HashSet<SocketAddr>,
}
@@ -254,6 +269,7 @@ pub struct MetadataScheduler {
total_completed: Arc<AtomicU64>,
queue_len: Arc<AtomicUsize>,
runtime_stats: DhtRuntimeStats,
peer_lookup_tx: Option<mpsc::Sender<[u8; 20]>>,
shutdown: CancellationToken,
}
@@ -286,6 +302,29 @@ impl MetadataScheduler {
queue_len: Arc<AtomicUsize>,
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<HashDiscovered>,
fetcher: Arc<RbitFetcher>,
limits: MetadataSchedulerLimits,
callbacks: MetadataSchedulerCallbacks,
queue_len: Arc<AtomicUsize>,
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::<String, InFlightState>::new();
let mut tasks = JoinSet::<JobResult>::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<String, InFlightState>,
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<JobResult>,
) {
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<String, InFlightState>,
entry: QueuedHash,
peer_rx: mpsc::Receiver<PeerCandidate>,
tasks: &mut JoinSet<JobResult>,
) {
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<String, InFlightState>,
tasks: &mut JoinSet<JobResult>,
) {
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<RbitFetcher>,
callback: &Arc<ArcSwapOption<TorrentAckCallback>>,
on_metadata_fetch: &Arc<ArcSwapOption<MetadataFetchCallback>>,
runtime_stats: &DhtRuntimeStats,
peer_lookup_tx: Option<mpsc::Sender<[u8; 20]>>,
) -> 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<F, Fut>(
peers: Vec<PeerCandidate>,
peer_rx: &mut mpsc::Receiver<PeerCandidate>,
fetch: F,
runtime_stats: &DhtRuntimeStats,
) -> Option<(PeerCandidate, FetchedMetadata)>
where
F: Fn(PeerCandidate) -> Fut + Clone + Send + Sync + 'static,
Fut: Future<Output = (PeerCandidate, MetadataFetchOutcome)> + 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<AtomicBool>);
impl Drop for DropSignal {
fn drop(&mut self) {
self.0.store(true, Ordering::Relaxed);
}
}
fn empty_torrent_callback() -> Arc<ArcSwapOption<TorrentAckCallback>> {
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::<Vec<_>>());
}
#[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::<Vec<_>>());
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());
src/server.rs
+34 -5
View File
@@ -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<ArcSwapOption<MetadataCompletionCallback>>,
error_callback: Arc<ArcSwapOption<ErrorCallback>>,
crawl_engine: Arc<CrawlEngine>,
peer_lookup: PeerLookupHandle,
hash_events_tx: mpsc::Sender<HashDiscovered>,
metadata_queue_len: Arc<AtomicUsize>,
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);
}
}
}
_ => {}
src/types.rs
+24
View File
@@ -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 {