729 lines
24 KiB
Rust
729 lines
24 KiB
Rust
// 负责执行有界 BEP-51 采样并将准入后的 infohash 交给 Peer 查找
|
|
|
|
use crate::addr::is_valid_node_addr;
|
|
use crate::budget::{RateBucket, SharedRateBudget};
|
|
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::{DiscoverySource, NetMode, NodeTuple, SampleInfohashesOptions};
|
|
use ahash::{AHashMap, AHashSet};
|
|
use arc_swap::{ArcSwap, ArcSwapOption};
|
|
use bytes::BytesMut;
|
|
#[cfg(feature = "metrics")]
|
|
use metrics::{counter, gauge};
|
|
use std::collections::VecDeque;
|
|
use std::future::Future;
|
|
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 ADMISSION_BATCH_CHANNEL_CAPACITY: usize = 64;
|
|
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,
|
|
source: DiscoverySource,
|
|
}
|
|
|
|
struct SampleResponse {
|
|
remote_addr: SocketAddr,
|
|
tid: TransactionId,
|
|
response: DhtResponse,
|
|
}
|
|
|
|
struct SampleAdmissionBatch {
|
|
preferred_node: NodeTuple,
|
|
source: DiscoverySource,
|
|
hashes: Vec<[u8; 20]>,
|
|
}
|
|
|
|
pub(crate) type SampleHashAdmissionCallback = Box<
|
|
dyn Fn(Vec<[u8; 20]>) -> std::pin::Pin<Box<dyn Future<Output = Vec<[u8; 20]>> + Send>>
|
|
+ Send
|
|
+ Sync
|
|
+ 'static,
|
|
>;
|
|
|
|
#[derive(Clone)]
|
|
pub(crate) struct SampleInfohashesHandle {
|
|
response_tx: mpsc::Sender<SampleResponse>,
|
|
runtime_stats: DhtRuntimeStats,
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
pub(crate) struct SampleCandidateRouter {
|
|
sender: mpsc::Sender<NodeTuple>,
|
|
sample_percent: u8,
|
|
runtime_stats: DhtRuntimeStats,
|
|
}
|
|
|
|
impl SampleCandidateRouter {
|
|
pub(crate) fn route(&self, node: NodeTuple) -> bool {
|
|
if self.sample_percent == 0
|
|
|| !is_valid_node_addr(&node.addr)
|
|
|| stable_address_bucket(node.addr) >= self.sample_percent
|
|
{
|
|
return false;
|
|
}
|
|
match self.sender.try_send(node) {
|
|
Ok(()) => {
|
|
self.runtime_stats.sample_candidate_routed();
|
|
self.runtime_stats.set_sample_candidate_queue_depth(
|
|
self.sender
|
|
.max_capacity()
|
|
.saturating_sub(self.sender.capacity()),
|
|
);
|
|
true
|
|
}
|
|
Err(mpsc::error::TrySendError::Full(_)) => {
|
|
self.runtime_stats.sample_candidate_fallback();
|
|
self.runtime_stats
|
|
.set_sample_candidate_queue_depth(self.sender.max_capacity());
|
|
false
|
|
}
|
|
Err(mpsc::error::TrySendError::Closed(_)) => false,
|
|
}
|
|
}
|
|
}
|
|
|
|
pub(crate) fn sample_candidate_lane(
|
|
options: &SampleInfohashesOptions,
|
|
runtime_stats: DhtRuntimeStats,
|
|
) -> (SampleCandidateRouter, mpsc::Receiver<NodeTuple>) {
|
|
let capacity = options.candidate_queue_capacity.max(1);
|
|
let (sender, receiver) = mpsc::channel(capacity);
|
|
runtime_stats.configure_sample_candidate_queue(capacity);
|
|
(
|
|
SampleCandidateRouter {
|
|
sender,
|
|
sample_percent: options.new_node_sample_percent.min(100),
|
|
runtime_stats,
|
|
},
|
|
receiver,
|
|
)
|
|
}
|
|
|
|
fn stable_address_bucket(addr: SocketAddr) -> u8 {
|
|
let mut hash = 0xcbf2_9ce4_8422_2325u64;
|
|
let mut mix = |byte: u8| {
|
|
hash ^= u64::from(byte);
|
|
hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
|
|
};
|
|
match addr.ip() {
|
|
std::net::IpAddr::V4(ip) => ip.octets().into_iter().for_each(&mut mix),
|
|
std::net::IpAddr::V6(ip) => ip.octets().into_iter().for_each(&mut mix),
|
|
}
|
|
addr.port().to_be_bytes().into_iter().for_each(&mut mix);
|
|
(hash % 100) as u8
|
|
}
|
|
|
|
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) outbound_query_budget: SharedRateBudget,
|
|
pub(crate) hash_admission: Arc<ArcSwapOption<SampleHashAdmissionCallback>>,
|
|
pub(crate) shutdown: CancellationToken,
|
|
}
|
|
|
|
pub(crate) struct SampleInfohashesInputs<'a> {
|
|
pub(crate) sockets: &'a std::collections::HashMap<SocketAddr, Arc<UdpSocket>>,
|
|
pub(crate) snapshot: Arc<ArcSwap<RoutingSnapshot>>,
|
|
pub(crate) candidate_rx: mpsc::Receiver<NodeTuple>,
|
|
pub(crate) crawl_engine: Arc<CrawlEngine>,
|
|
pub(crate) peer_lookup: PeerLookupHandle,
|
|
}
|
|
|
|
pub(crate) fn spawn_sample_infohashes(
|
|
netmode: NetMode,
|
|
local_id: [u8; 20],
|
|
inputs: SampleInfohashesInputs<'_>,
|
|
runtime: SampleInfohashesRuntime,
|
|
) -> SampleInfohashesHandle {
|
|
let SampleInfohashesInputs {
|
|
sockets,
|
|
snapshot,
|
|
candidate_rx,
|
|
crawl_engine,
|
|
peer_lookup,
|
|
} = inputs;
|
|
let SampleInfohashesRuntime {
|
|
options,
|
|
stats,
|
|
outbound_query_budget,
|
|
hash_admission,
|
|
shutdown,
|
|
} = runtime;
|
|
let (response_tx, response_rx) = mpsc::channel(RESPONSE_CHANNEL_CAPACITY);
|
|
let (admission_tx, admission_rx) = mpsc::channel(ADMISSION_BATCH_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();
|
|
spawn_sample_admission(
|
|
admission_rx,
|
|
hash_admission,
|
|
peer_lookup.request_sender(),
|
|
options.fallback_to_iterative,
|
|
stats.clone(),
|
|
shutdown.clone(),
|
|
);
|
|
let actor = SampleInfohashesActor {
|
|
netmode,
|
|
local_id,
|
|
socket_v4,
|
|
socket_v6,
|
|
snapshot,
|
|
candidate_rx,
|
|
direct_only: options.new_node_sample_percent > 0,
|
|
crawl_engine,
|
|
admission_tx,
|
|
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(),
|
|
outbound_query_budget,
|
|
shutdown,
|
|
};
|
|
tokio::spawn(actor.run());
|
|
SampleInfohashesHandle {
|
|
response_tx,
|
|
runtime_stats: stats,
|
|
}
|
|
}
|
|
|
|
struct SampleInfohashesActor {
|
|
netmode: NetMode,
|
|
local_id: [u8; 20],
|
|
socket_v4: Option<Arc<UdpSocket>>,
|
|
socket_v6: Option<Arc<UdpSocket>>,
|
|
snapshot: Arc<ArcSwap<RoutingSnapshot>>,
|
|
candidate_rx: mpsc::Receiver<NodeTuple>,
|
|
direct_only: bool,
|
|
crawl_engine: Arc<CrawlEngine>,
|
|
admission_tx: mpsc::Sender<SampleAdmissionBatch>,
|
|
response_rx: mpsc::Receiver<SampleResponse>,
|
|
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<SocketAddr, Instant>,
|
|
pending_addrs: AHashSet<SocketAddr>,
|
|
pending: AHashMap<PendingKey, PendingRequest>,
|
|
pending_expiry: VecDeque<(Instant, PendingKey)>,
|
|
next_tid: u64,
|
|
runtime_stats: DhtRuntimeStats,
|
|
outbound_query_budget: SharedRateBudget,
|
|
shutdown: CancellationToken,
|
|
}
|
|
|
|
fn spawn_sample_admission(
|
|
mut receiver: mpsc::Receiver<SampleAdmissionBatch>,
|
|
hash_admission: Arc<ArcSwapOption<SampleHashAdmissionCallback>>,
|
|
peer_lookup_tx: mpsc::Sender<PeerLookupRequest>,
|
|
fallback_to_iterative: bool,
|
|
runtime_stats: DhtRuntimeStats,
|
|
shutdown: CancellationToken,
|
|
) {
|
|
tokio::spawn(async move {
|
|
loop {
|
|
let batch = tokio::select! {
|
|
_ = shutdown.cancelled() => break,
|
|
batch = receiver.recv() => {
|
|
let Some(batch) = batch else { break };
|
|
batch
|
|
}
|
|
};
|
|
let input_len = batch.hashes.len();
|
|
let admitted = match hash_admission.load_full() {
|
|
Some(callback) => callback(batch.hashes).await,
|
|
None => batch.hashes,
|
|
};
|
|
runtime_stats.sample_hash_filtered(input_len.saturating_sub(admitted.len()));
|
|
for info_hash in admitted {
|
|
let request = PeerLookupRequest {
|
|
info_hash,
|
|
preferred_node: Some(batch.preferred_node),
|
|
allow_iterative_fallback: fallback_to_iterative,
|
|
source: batch.source,
|
|
completion: None,
|
|
};
|
|
let sent = tokio::select! {
|
|
_ = shutdown.cancelled() => false,
|
|
result = peer_lookup_tx.send(request) => result.is_ok(),
|
|
};
|
|
if !sent {
|
|
return;
|
|
}
|
|
runtime_stats.sample_hash_discovered();
|
|
}
|
|
}
|
|
});
|
|
}
|
|
|
|
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;
|
|
}
|
|
}
|
|
}
|
|
self.runtime_stats.set_sample_candidate_queue_depth(0);
|
|
}
|
|
|
|
async fn dispatch(&mut self, now: Instant) {
|
|
if self.max_in_flight == 0
|
|
|| self.pending.len() >= self.max_in_flight
|
|
|| self.admission_tx.capacity() == 0
|
|
{
|
|
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;
|
|
while sent_count < budget {
|
|
let Some(node) = self.next_direct_candidate(now) else {
|
|
break;
|
|
};
|
|
if self
|
|
.send_query(node, DiscoverySource::SampleDirect, now)
|
|
.await
|
|
{
|
|
sent_count += 1;
|
|
}
|
|
}
|
|
if self.direct_only {
|
|
if sent_count < budget {
|
|
self.query_budget.refund(budget - sent_count);
|
|
}
|
|
return;
|
|
}
|
|
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, DiscoverySource::SampleSnapshot, now)
|
|
.await
|
|
{
|
|
sent_count += 1;
|
|
}
|
|
}
|
|
if sent_count < budget {
|
|
self.query_budget.refund(budget - sent_count);
|
|
}
|
|
}
|
|
|
|
fn next_direct_candidate(&mut self, now: Instant) -> Option<NodeTuple> {
|
|
for _ in 0..64 {
|
|
let node = self.candidate_rx.try_recv().ok()?;
|
|
self.runtime_stats
|
|
.set_sample_candidate_queue_depth(self.candidate_rx.len());
|
|
if self.pending_addrs.contains(&node.addr)
|
|
|| self
|
|
.next_allowed
|
|
.get(&node.addr)
|
|
.is_some_and(|deadline| *deadline > now)
|
|
{
|
|
continue;
|
|
}
|
|
return Some(node);
|
|
}
|
|
None
|
|
}
|
|
|
|
async fn send_query(&mut self, node: NodeTuple, source: DiscoverySource, 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;
|
|
};
|
|
if !self.outbound_query_budget.try_take_one(now) {
|
|
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.outbound_query_budget.refund_one();
|
|
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, source });
|
|
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(source == DiscoverySource::SampleDirect);
|
|
#[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,
|
|
};
|
|
let Some(pending) = self.pending.remove(&key) else {
|
|
return;
|
|
};
|
|
self.pending_addrs.remove(&event.remote_addr);
|
|
self.runtime_stats
|
|
.sample_response(pending.source == DiscoverySource::SampleDirect);
|
|
|
|
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 hashes = Vec::new();
|
|
let mut batch_hashes = AHashSet::new();
|
|
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;
|
|
}
|
|
if batch_hashes.insert(hash) {
|
|
hashes.push(hash);
|
|
}
|
|
}
|
|
}
|
|
|
|
let candidate_count = hashes.len();
|
|
let admitted_to_triage = if hashes.is_empty() {
|
|
false
|
|
} else {
|
|
let hashes_to_remember = hashes.clone();
|
|
if self
|
|
.admission_tx
|
|
.try_send(SampleAdmissionBatch {
|
|
preferred_node,
|
|
source: pending.source,
|
|
hashes,
|
|
})
|
|
.is_ok()
|
|
{
|
|
for hash in hashes_to_remember {
|
|
self.remember_hash(hash);
|
|
}
|
|
true
|
|
} else {
|
|
for _ in 0..candidate_count {
|
|
self.runtime_stats.sample_hash_dropped();
|
|
}
|
|
#[cfg(feature = "metrics")]
|
|
counter!("dht_sample_infohashes_dropped_total", "reason" => "admission_queue_full")
|
|
.increment(candidate_count as u64);
|
|
false
|
|
}
|
|
};
|
|
let discovered = usize::from(admitted_to_triage) * candidate_count;
|
|
|
|
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;
|
|
|
|
fn node_for_lane(want_direct: bool, percent: u8) -> NodeTuple {
|
|
(1..=u16::MAX)
|
|
.map(|port| NodeTuple {
|
|
id: [port as u8; 20],
|
|
addr: SocketAddr::from(([8, 8, 8, 8], port)),
|
|
})
|
|
.find(|node| (stable_address_bucket(node.addr) < percent) == want_direct)
|
|
.expect("both routing lanes have an address")
|
|
}
|
|
|
|
#[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);
|
|
}
|
|
|
|
#[test]
|
|
fn new_node_routing_is_stable_and_bounded() {
|
|
let options = SampleInfohashesOptions {
|
|
new_node_sample_percent: 50,
|
|
candidate_queue_capacity: 1,
|
|
..SampleInfohashesOptions::default()
|
|
};
|
|
let stats = DhtRuntimeStats::default();
|
|
let (router, mut receiver) = sample_candidate_lane(&options, stats.clone());
|
|
let direct = node_for_lane(true, 50);
|
|
let crawl = node_for_lane(false, 50);
|
|
|
|
assert!(router.route(direct));
|
|
assert!(!router.route(crawl));
|
|
assert!(!router.route(direct));
|
|
assert_eq!(receiver.try_recv().unwrap(), direct);
|
|
|
|
let snapshot = stats.snapshot();
|
|
assert_eq!(snapshot.sample_candidate_queue_capacity, 1);
|
|
assert_eq!(snapshot.sample_candidates_routed, 1);
|
|
assert_eq!(snapshot.sample_candidates_fallback, 1);
|
|
}
|
|
|
|
#[test]
|
|
fn zero_percent_preserves_snapshot_only_sampling() {
|
|
let options = SampleInfohashesOptions {
|
|
new_node_sample_percent: 0,
|
|
..SampleInfohashesOptions::default()
|
|
};
|
|
let (router, mut receiver) = sample_candidate_lane(&options, DhtRuntimeStats::default());
|
|
assert!(!router.route(node_for_lane(false, 0)));
|
|
assert!(receiver.try_recv().is_err());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn batch_admission_only_forwards_application_approved_hashes() {
|
|
let (batch_tx, batch_rx) = mpsc::channel(1);
|
|
let (lookup_tx, mut lookup_rx) = mpsc::channel(2);
|
|
let admission = Arc::new(ArcSwapOption::empty());
|
|
let callback: Arc<SampleHashAdmissionCallback> = Arc::new(Box::new(|hashes| {
|
|
Box::pin(async move { hashes.into_iter().filter(|hash| *hash == [2; 20]).collect() })
|
|
}));
|
|
admission.store(Some(callback));
|
|
let stats = DhtRuntimeStats::default();
|
|
let shutdown = CancellationToken::new();
|
|
spawn_sample_admission(
|
|
batch_rx,
|
|
admission,
|
|
lookup_tx,
|
|
false,
|
|
stats.clone(),
|
|
shutdown.clone(),
|
|
);
|
|
batch_tx
|
|
.send(SampleAdmissionBatch {
|
|
preferred_node: NodeTuple {
|
|
id: [9; 20],
|
|
addr: "127.0.0.1:6881".parse().unwrap(),
|
|
},
|
|
source: DiscoverySource::SampleDirect,
|
|
hashes: vec![[1; 20], [2; 20]],
|
|
})
|
|
.await
|
|
.unwrap();
|
|
|
|
let request = tokio::time::timeout(Duration::from_secs(1), lookup_rx.recv())
|
|
.await
|
|
.unwrap()
|
|
.unwrap();
|
|
assert_eq!(request.info_hash, [2; 20]);
|
|
assert!(!request.allow_iterative_fallback);
|
|
assert_eq!(request.source, DiscoverySource::SampleDirect);
|
|
let snapshot = stats.snapshot();
|
|
assert_eq!(snapshot.sample_infohashes_hashes_filtered, 1);
|
|
assert_eq!(snapshot.sample_infohashes_hashes_discovered, 1);
|
|
shutdown.cancel();
|
|
}
|
|
}
|