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

refactor: 拆分核心运行和存储模块

Browse Source
This commit is contained in:
chuan
2026-08-10 15:48:24 +08:00
parent 9f0911c596
commit c70d7e5c7b
13 changed files with 2459 additions and 2394 deletions
Download Patch File Download Diff File Expand all files Collapse all files
TODOS.md
+1
View File
@@ -42,6 +42,7 @@
- [x] 将最终应用与可复用 DHT 基础库分离
- [x] 删除被内容组索引替代的旧状态和无效兼容代码
- [x] 收紧仅供测试或存储内部使用的接口和依赖
- [x] 按运行统计调度队列 Peer 竞速和 RocksDB 生命周期职责拆分核心大模块
- [x] 实现 BEP-51 `sample_infohashes` 主动发现
- [x] 实现主动 Peer 查找和 Metadata 获取
- [x] 验证远程公网设备能够持续获取 Metadata
src/crawler/src/runtime_stats.rs
+8 -848
View File
@@ -1,107 +1,16 @@
// 负责收集 DHT 抓取调度网络和 Metadata 运行指标
use crate::types::DiscoverySource;
use std::sync::Arc;
use std::sync::atomic::{AtomicU32, AtomicU64, AtomicUsize, Ordering};
pub const METADATA_QUEUE_WAIT_BUCKETS_MS: [u64; 8] = [10, 50, 100, 250, 500, 1_000, 2_000, 5_000];
pub const METADATA_FETCH_BUCKETS_MS: [u64; 7] = [250, 500, 1_000, 2_000, 4_000, 6_000, 10_000];
pub const METADATA_SIZE_BUCKETS_BYTES: [u64; 8] = [
16 * 1024,
32 * 1024,
64 * 1024,
128 * 1024,
256 * 1024,
512 * 1024,
1024 * 1024,
10 * 1024 * 1024,
];
mod histogram;
mod recorders;
#[derive(Debug, Clone, PartialEq, Eq)]
/// Snapshot of a non-cumulative fixed-bucket histogram.
pub struct FixedHistogramSnapshot {
/// Inclusive upper bound for each bucket.
pub bounds: Vec<u64>,
/// Per-bucket counts; values are not cumulative.
pub counts: Vec<u64>,
/// Values larger than the final bound.
pub overflow: u64,
/// Total observations including overflow.
pub count: u64,
/// Wrapping sum of all observed values.
pub sum: u64,
}
impl FixedHistogramSnapshot {
/// Returns the inclusive bucket bound containing the requested approximate percentile.
///
/// `percentile` is clamped to `0.0..=1.0`. Overflow observations return the final finite
/// bound because the histogram intentionally stores no dynamic maximum.
pub fn percentile(&self, percentile: f64) -> Option<u64> {
if self.count == 0 {
return None;
}
let rank = ((self.count as f64 * percentile.clamp(0.0, 1.0)).ceil() as u64).max(1);
let mut cumulative = 0u64;
for (bound, count) in self.bounds.iter().zip(&self.counts) {
cumulative = cumulative.saturating_add(*count);
if cumulative >= rank {
return Some(*bound);
}
}
self.bounds.last().copied()
}
}
struct AtomicFixedHistogram<const N: usize> {
bounds: [u64; N],
counts: [AtomicU64; N],
overflow: AtomicU64,
count: AtomicU64,
sum: AtomicU64,
}
impl<const N: usize> AtomicFixedHistogram<N> {
fn new(bounds: [u64; N]) -> Self {
Self {
bounds,
counts: std::array::from_fn(|_| AtomicU64::new(0)),
overflow: AtomicU64::new(0),
count: AtomicU64::new(0),
sum: AtomicU64::new(0),
}
}
fn record(&self, value: u64) {
if let Some(index) = self.bounds.iter().position(|bound| value <= *bound) {
self.counts[index].fetch_add(1, Ordering::Relaxed);
} else {
self.overflow.fetch_add(1, Ordering::Relaxed);
}
self.count.fetch_add(1, Ordering::Relaxed);
self.sum.fetch_add(value, Ordering::Relaxed);
}
fn snapshot(&self) -> FixedHistogramSnapshot {
FixedHistogramSnapshot {
bounds: self.bounds.to_vec(),
counts: self
.counts
.iter()
.map(|count| count.load(Ordering::Relaxed))
.collect(),
overflow: self.overflow.load(Ordering::Relaxed),
count: self.count.load(Ordering::Relaxed),
sum: self.sum.load(Ordering::Relaxed),
}
}
}
impl<const N: usize> Default for AtomicFixedHistogram<N> {
fn default() -> Self {
Self::new([0; N])
}
}
use histogram::AtomicFixedHistogram;
pub use histogram::{
FixedHistogramSnapshot, METADATA_FETCH_BUCKETS_MS, METADATA_QUEUE_WAIT_BUCKETS_MS,
METADATA_SIZE_BUCKETS_BYTES,
};
/// Transport-neutral counters and fixed-bucket histograms used by dashboards.
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -809,756 +718,7 @@ impl DhtRuntimeStats {
metadata_size_bytes: inner.metadata_size_bytes.snapshot(),
}
}
pub(crate) fn hash_received(&self) {
self.inner.hashes_received.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn hash_ingress_dropped(&self) {
self.inner
.hash_ingress_dropped
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn set_hash_ingress_queue_depth(&self, depth: usize) {
self.inner
.hash_ingress_queue_depth
.store(depth, Ordering::Relaxed);
}
pub(crate) fn set_crawl_priority_queue_depth(&self, depth: usize) {
self.inner
.crawl_priority_queue_depth
.store(depth, Ordering::Relaxed);
}
pub(crate) fn set_crawl_discovery_queue_depth(&self, depth: usize) {
self.inner
.crawl_discovery_queue_depth
.store(depth, Ordering::Relaxed);
}
pub(crate) fn set_metadata_queue(&self, depth: usize, in_flight: usize) {
self.inner
.metadata_queue_depth
.store(depth, Ordering::Relaxed);
self.inner
.metadata_in_flight
.store(in_flight, Ordering::Relaxed);
}
pub(crate) fn metadata_queue_deduplicated(&self) {
self.inner
.metadata_queue_deduplicated
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_queue_inserted(&self) {
self.inner
.metadata_queue_inserted
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_queue_evicted(&self) {
self.inner
.metadata_queue_evicted
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_queue_stale(&self, count: usize) {
self.inner
.metadata_queue_stale
.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 peer_lookup_preferred_succeeded(&self) {
self.inner
.peer_lookup_preferred_succeeded
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn peer_lookup_fallback(&self) {
self.inner
.peer_lookup_fallbacks
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn configure_sample_candidate_queue(&self, capacity: usize) {
self.inner
.sample_candidate_queue_capacity
.store(capacity, Ordering::Relaxed);
}
pub(crate) fn set_sample_candidate_queue_depth(&self, depth: usize) {
self.inner
.sample_candidate_queue_depth
.store(depth, Ordering::Relaxed);
}
pub(crate) fn sample_candidate_routed(&self) {
self.inner
.sample_candidates_routed
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn sample_candidate_fallback(&self) {
self.inner
.sample_candidates_fallback
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn sample_query(&self, direct: bool) {
self.inner
.sample_infohashes_queries
.fetch_add(1, Ordering::Relaxed);
let source = if direct {
&self.inner.sample_infohashes_direct_queries
} else {
&self.inner.sample_infohashes_snapshot_queries
};
source.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn sample_response(&self, direct: bool) {
self.inner
.sample_infohashes_responses
.fetch_add(1, Ordering::Relaxed);
let source = if direct {
&self.inner.sample_infohashes_direct_responses
} else {
&self.inner.sample_infohashes_snapshot_responses
};
source.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_filtered(&self, count: usize) {
self.inner
.sample_infohashes_hashes_filtered
.fetch_add(count as u64, 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
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_peer_succeeded(&self) {
self.inner
.metadata_peer_succeeded
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_success_from(&self, source: DiscoverySource) {
let counter = match source {
DiscoverySource::AnnouncePeer => &self.inner.metadata_success_from_announce,
DiscoverySource::SampleDirect => &self.inner.metadata_success_from_sample_direct,
DiscoverySource::SampleSnapshot => &self.inner.metadata_success_from_sample_snapshot,
DiscoverySource::ActiveLookup => &self.inner.metadata_success_from_active_lookup,
};
counter.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_peer_failed(&self) {
self.inner
.metadata_peer_failed
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_peer_timeout(&self) {
self.inner
.metadata_peer_timeouts
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_connect_failed(&self) {
self.inner
.metadata_connect_failed
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_no_extension(&self) {
self.inner
.metadata_no_extension
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_peer_failure_cache_hit(&self) {
self.inner
.metadata_peer_failure_cache_hits
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn set_metadata_peer_failure_cache_entries(&self, count: usize) {
self.inner
.metadata_peer_failure_cache_entries
.store(count, Ordering::Relaxed);
}
pub(crate) fn set_node_pool_size(&self, size: usize) {
self.inner.node_pool_size.store(size, Ordering::Relaxed);
}
pub(crate) fn set_find_node_in_flight(&self, count: usize) {
self.inner
.find_node_in_flight
.store(count, Ordering::Relaxed);
}
pub(crate) fn set_find_node_effective_rate(&self, rate: u32) {
self.inner
.find_node_effective_rate_per_sec
.store(rate, Ordering::Relaxed);
}
pub(crate) fn query_new(&self) {
self.inner.queries_new.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn query_revisit(&self) {
self.inner.queries_revisit.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn query_bootstrap(&self) {
self.inner.queries_bootstrap.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn response(&self) {
self.inner.responses.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn unmatched_response(&self) {
self.inner
.unmatched_responses
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn timeout(&self) {
self.inner.timeouts.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn send_failure(&self) {
self.inner.send_failures.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn crawl_event_dropped_response(&self) {
self.inner
.crawl_events_dropped_response
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn crawl_event_dropped_discovered(&self) {
self.inner
.crawl_events_dropped_discovered
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn udp_received(&self) {
self.inner.udp_received.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn udp_received_bytes(&self, bytes: usize) {
self.inner
.udp_rx_bytes
.fetch_add(bytes as u64, Ordering::Relaxed);
}
pub(crate) fn udp_sent(&self, bytes: usize) {
self.inner.udp_tx_packets.fetch_add(1, Ordering::Relaxed);
self.inner
.udp_tx_bytes
.fetch_add(bytes as u64, Ordering::Relaxed);
}
pub(crate) fn inbound_query(&self, query: &str) {
let counter = match query {
"ping" => &self.inner.inbound_ping,
"find_node" => &self.inner.inbound_find_node,
"get_peers" => &self.inner.inbound_get_peers,
"announce_peer" => &self.inner.inbound_announce_peer,
_ => &self.inner.inbound_other,
};
counter.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn response_normal(&self) {
self.inner.response_normal.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn response_send_failed(&self) {
self.inner
.response_send_failed
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn announce_accepted(&self) {
self.inner.announce_accepted.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn announce_invalid_token(&self) {
self.inner
.announce_invalid_token
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn announce_filtered(&self) {
self.inner.announce_filtered.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn node_admitted(&self) {
self.inner.node_admitted.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn node_replaced(&self) {
self.inner.node_replaced.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn node_dropped_duplicate(&self) {
self.inner
.node_dropped_duplicate
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn node_dropped_rate_limited(&self) {
self.inner
.node_dropped_rate_limited
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn node_dropped_invalid(&self) {
self.inner
.node_dropped_invalid
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_bytes_downloaded(&self, bytes: usize) {
self.inner
.metadata_bytes_downloaded
.fetch_add(bytes as u64, Ordering::Relaxed);
}
pub(crate) fn metadata_failure_timeout(&self) {
self.inner
.metadata_failure_timeout
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_failure_connect(&self) {
self.inner
.metadata_failure_connect
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_failure_no_extension(&self) {
self.inner
.metadata_failure_no_extension
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_failure_send(&self) {
self.inner
.metadata_failure_send
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_failure_size_limit(&self) {
self.inner
.metadata_failure_size_limit
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_failure_sha1(&self) {
self.inner
.metadata_failure_sha1
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_failure_parse(&self) {
self.inner
.metadata_failure_parse
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_failure_other(&self) {
self.inner
.metadata_failure_other
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn peer_cache_hit_timeout(&self) {
self.inner
.peer_cache_timeout_hits
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn peer_cache_hit_connect(&self) {
self.inner
.peer_cache_connect_hits
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn observe_metadata_queue_wait(&self, millis: u64) {
self.inner.queue_wait_ms.record(millis);
}
pub(crate) fn observe_metadata_fetch_duration(&self, millis: u64) {
self.inner.fetch_duration_ms.record(millis);
}
pub(crate) fn observe_metadata_size(&self, bytes: usize) {
self.inner.metadata_size_bytes.record(bytes as u64);
}
pub(crate) fn udp_queue_full(&self) {
self.inner.udp_queue_full.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn udp_invalid(&self) {
self.inner.udp_invalid.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn udp_response_rate_limited(&self) {
self.inner
.udp_responses_rate_limited
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn udp_response_priority_reserved(&self) {
self.inner
.udp_responses_priority_reserved
.fetch_add(1, Ordering::Relaxed);
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn cloned_handle_updates_one_snapshot() {
let stats = DhtRuntimeStats::with_limits(DhtRuntimeLimits {
metadata_queue: 100,
node_pool: 1_000,
node_pool_low_watermark: 10,
find_node_in_flight: 512,
initial_find_node_rate: 200,
hash_ingress_queue: 20,
crawl_priority_queue: 30,
crawl_discovery_queue: 40,
});
let writer = stats.clone();
writer.hash_received();
writer.hash_ingress_dropped();
writer.set_hash_ingress_queue_depth(11);
writer.set_crawl_priority_queue_depth(12);
writer.set_crawl_discovery_queue_depth(13);
writer.set_metadata_queue(42, 7);
writer.metadata_queue_inserted();
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.peer_lookup_preferred_succeeded();
writer.peer_lookup_fallback();
writer.configure_sample_candidate_queue(20);
writer.set_sample_candidate_queue_depth(14);
writer.sample_candidate_routed();
writer.sample_candidate_fallback();
writer.sample_query(true);
writer.sample_response(true);
writer.sample_timeout();
writer.sample_send_failed();
writer.sample_response_dropped();
writer.sample_hash_discovered();
writer.sample_hash_filtered(1);
writer.sample_hash_duplicate();
writer.sample_hash_dropped();
writer.metadata_peer_attempt();
writer.metadata_peer_succeeded();
writer.metadata_success_from(DiscoverySource::AnnouncePeer);
writer.metadata_success_from(DiscoverySource::SampleDirect);
writer.metadata_success_from(DiscoverySource::SampleSnapshot);
writer.metadata_success_from(DiscoverySource::ActiveLookup);
writer.metadata_peer_failed();
writer.metadata_peer_timeout();
writer.metadata_connect_failed();
writer.metadata_no_extension();
writer.metadata_peer_failure_cache_hit();
writer.set_metadata_peer_failure_cache_entries(9);
writer.set_node_pool_size(321);
writer.set_find_node_in_flight(12);
writer.set_find_node_effective_rate(150);
writer.query_new();
writer.query_revisit();
writer.query_bootstrap();
writer.response();
writer.unmatched_response();
writer.timeout();
writer.send_failure();
writer.crawl_event_dropped_response();
writer.crawl_event_dropped_discovered();
writer.udp_received();
writer.udp_queue_full();
writer.udp_invalid();
writer.udp_response_rate_limited();
writer.udp_response_priority_reserved();
assert_eq!(
stats.snapshot(),
DhtRuntimeSnapshot {
hashes_received: 1,
hash_ingress_dropped: 1,
hash_ingress_queue_depth: 11,
hash_ingress_queue_capacity: 20,
crawl_priority_queue_depth: 12,
crawl_priority_queue_capacity: 30,
crawl_discovery_queue_depth: 13,
crawl_discovery_queue_capacity: 40,
metadata_queue_depth: 42,
metadata_queue_max: 100,
metadata_in_flight: 7,
metadata_queue_inserted: 1,
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,
peer_lookup_preferred_succeeded: 1,
peer_lookup_fallbacks: 1,
sample_candidate_queue_depth: 14,
sample_candidate_queue_capacity: 20,
sample_candidates_routed: 1,
sample_candidates_fallback: 1,
sample_infohashes_queries: 1,
sample_infohashes_direct_queries: 1,
sample_infohashes_snapshot_queries: 0,
sample_infohashes_responses: 1,
sample_infohashes_direct_responses: 1,
sample_infohashes_snapshot_responses: 0,
sample_infohashes_timeouts: 1,
sample_infohashes_send_failures: 1,
sample_infohashes_response_dropped: 1,
sample_infohashes_hashes_discovered: 1,
sample_infohashes_hashes_filtered: 1,
sample_infohashes_hashes_duplicate: 1,
sample_infohashes_hashes_dropped: 1,
metadata_peer_attempts: 1,
metadata_peer_succeeded: 1,
metadata_success_from_announce: 1,
metadata_success_from_sample_direct: 1,
metadata_success_from_sample_snapshot: 1,
metadata_success_from_active_lookup: 1,
metadata_peer_failed: 1,
metadata_peer_timeouts: 1,
metadata_connect_failed: 1,
metadata_no_extension: 1,
metadata_peer_failure_cache_hits: 1,
metadata_peer_failure_cache_entries: 9,
node_pool_size: 321,
node_pool_capacity: 1_000,
node_pool_low_watermark: 10,
find_node_in_flight: 12,
find_node_in_flight_max: 512,
find_node_effective_rate_per_sec: 150,
queries_new: 1,
queries_revisit: 1,
queries_bootstrap: 1,
responses: 1,
unmatched_responses: 1,
timeouts: 1,
send_failures: 1,
crawl_events_dropped_response: 1,
crawl_events_dropped_discovered: 1,
udp_received: 1,
udp_queue_full: 1,
udp_invalid: 1,
udp_responses_rate_limited: 1,
udp_responses_priority_reserved: 1,
}
);
}
#[test]
fn observability_snapshot_tracks_fixed_histograms_and_categories() {
let stats = DhtRuntimeStats::default();
stats.udp_received();
stats.udp_received_bytes(128);
stats.udp_sent(64);
stats.inbound_query("ping");
stats.inbound_query("unknown");
stats.node_admitted();
stats.metadata_bytes_downloaded(1_024);
stats.metadata_failure_parse();
stats.observe_metadata_queue_wait(75);
stats.observe_metadata_queue_wait(600);
stats.observe_metadata_fetch_duration(1_500);
stats.observe_metadata_size(70_000);
let snapshot = stats.observability_snapshot();
assert_eq!(snapshot.udp_rx_packets, 1);
assert_eq!(snapshot.udp_rx_bytes, 128);
assert_eq!(snapshot.udp_tx_packets, 1);
assert_eq!(snapshot.udp_tx_bytes, 64);
assert_eq!(snapshot.inbound_ping, 1);
assert_eq!(snapshot.inbound_other, 1);
assert_eq!(snapshot.node_admitted, 1);
assert_eq!(snapshot.metadata_bytes_downloaded, 1_024);
assert_eq!(snapshot.metadata_failure_parse, 1);
assert_eq!(snapshot.queue_wait_ms.percentile(0.50), Some(100));
assert_eq!(snapshot.queue_wait_ms.percentile(0.95), Some(1_000));
assert_eq!(snapshot.fetch_duration_ms.percentile(0.95), Some(2_000));
assert_eq!(snapshot.metadata_size_bytes.percentile(0.50), Some(131_072));
}
}
mod tests;
src/crawler/src/runtime_stats/histogram.rs
+102
View File
@@ -0,0 +1,102 @@
// 负责固定分桶直方图的记录快照和近似分位数计算
use std::sync::atomic::{AtomicU64, Ordering};
pub const METADATA_QUEUE_WAIT_BUCKETS_MS: [u64; 8] = [10, 50, 100, 250, 500, 1_000, 2_000, 5_000];
pub const METADATA_FETCH_BUCKETS_MS: [u64; 7] = [250, 500, 1_000, 2_000, 4_000, 6_000, 10_000];
pub const METADATA_SIZE_BUCKETS_BYTES: [u64; 8] = [
16 * 1024,
32 * 1024,
64 * 1024,
128 * 1024,
256 * 1024,
512 * 1024,
1024 * 1024,
10 * 1024 * 1024,
];
#[derive(Debug, Clone, PartialEq, Eq)]
/// Snapshot of a non-cumulative fixed-bucket histogram.
pub struct FixedHistogramSnapshot {
/// Inclusive upper bound for each bucket.
pub bounds: Vec<u64>,
/// Per-bucket counts; values are not cumulative.
pub counts: Vec<u64>,
/// Values larger than the final bound.
pub overflow: u64,
/// Total observations including overflow.
pub count: u64,
/// Wrapping sum of all observed values.
pub sum: u64,
}
impl FixedHistogramSnapshot {
/// Returns the inclusive bucket bound containing the requested approximate percentile.
///
/// `percentile` is clamped to `0.0..=1.0`. Overflow observations return the final finite
/// bound because the histogram intentionally stores no dynamic maximum.
pub fn percentile(&self, percentile: f64) -> Option<u64> {
if self.count == 0 {
return None;
}
let rank = ((self.count as f64 * percentile.clamp(0.0, 1.0)).ceil() as u64).max(1);
let mut cumulative = 0u64;
for (bound, count) in self.bounds.iter().zip(&self.counts) {
cumulative = cumulative.saturating_add(*count);
if cumulative >= rank {
return Some(*bound);
}
}
self.bounds.last().copied()
}
}
pub(super) struct AtomicFixedHistogram<const N: usize> {
bounds: [u64; N],
counts: [AtomicU64; N],
overflow: AtomicU64,
count: AtomicU64,
sum: AtomicU64,
}
impl<const N: usize> AtomicFixedHistogram<N> {
pub(super) fn new(bounds: [u64; N]) -> Self {
Self {
bounds,
counts: std::array::from_fn(|_| AtomicU64::new(0)),
overflow: AtomicU64::new(0),
count: AtomicU64::new(0),
sum: AtomicU64::new(0),
}
}
pub(super) fn record(&self, value: u64) {
if let Some(index) = self.bounds.iter().position(|bound| value <= *bound) {
self.counts[index].fetch_add(1, Ordering::Relaxed);
} else {
self.overflow.fetch_add(1, Ordering::Relaxed);
}
self.count.fetch_add(1, Ordering::Relaxed);
self.sum.fetch_add(value, Ordering::Relaxed);
}
pub(super) fn snapshot(&self) -> FixedHistogramSnapshot {
FixedHistogramSnapshot {
bounds: self.bounds.to_vec(),
counts: self
.counts
.iter()
.map(|count| count.load(Ordering::Relaxed))
.collect(),
overflow: self.overflow.load(Ordering::Relaxed),
count: self.count.load(Ordering::Relaxed),
sum: self.sum.load(Ordering::Relaxed),
}
}
}
impl<const N: usize> Default for AtomicFixedHistogram<N> {
fn default() -> Self {
Self::new([0; N])
}
}
src/crawler/src/runtime_stats/recorders.rs
+549
View File
@@ -0,0 +1,549 @@
// 负责写入 DHT 运行计数器状态和观测直方图
use super::{DhtRuntimeStats, Ordering};
use crate::types::DiscoverySource;
impl DhtRuntimeStats {
pub(crate) fn hash_received(&self) {
self.inner.hashes_received.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn hash_ingress_dropped(&self) {
self.inner
.hash_ingress_dropped
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn set_hash_ingress_queue_depth(&self, depth: usize) {
self.inner
.hash_ingress_queue_depth
.store(depth, Ordering::Relaxed);
}
pub(crate) fn set_crawl_priority_queue_depth(&self, depth: usize) {
self.inner
.crawl_priority_queue_depth
.store(depth, Ordering::Relaxed);
}
pub(crate) fn set_crawl_discovery_queue_depth(&self, depth: usize) {
self.inner
.crawl_discovery_queue_depth
.store(depth, Ordering::Relaxed);
}
pub(crate) fn set_metadata_queue(&self, depth: usize, in_flight: usize) {
self.inner
.metadata_queue_depth
.store(depth, Ordering::Relaxed);
self.inner
.metadata_in_flight
.store(in_flight, Ordering::Relaxed);
}
pub(crate) fn metadata_queue_deduplicated(&self) {
self.inner
.metadata_queue_deduplicated
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_queue_inserted(&self) {
self.inner
.metadata_queue_inserted
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_queue_evicted(&self) {
self.inner
.metadata_queue_evicted
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_queue_stale(&self, count: usize) {
self.inner
.metadata_queue_stale
.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 peer_lookup_preferred_succeeded(&self) {
self.inner
.peer_lookup_preferred_succeeded
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn peer_lookup_fallback(&self) {
self.inner
.peer_lookup_fallbacks
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn configure_sample_candidate_queue(&self, capacity: usize) {
self.inner
.sample_candidate_queue_capacity
.store(capacity, Ordering::Relaxed);
}
pub(crate) fn set_sample_candidate_queue_depth(&self, depth: usize) {
self.inner
.sample_candidate_queue_depth
.store(depth, Ordering::Relaxed);
}
pub(crate) fn sample_candidate_routed(&self) {
self.inner
.sample_candidates_routed
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn sample_candidate_fallback(&self) {
self.inner
.sample_candidates_fallback
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn sample_query(&self, direct: bool) {
self.inner
.sample_infohashes_queries
.fetch_add(1, Ordering::Relaxed);
let source = if direct {
&self.inner.sample_infohashes_direct_queries
} else {
&self.inner.sample_infohashes_snapshot_queries
};
source.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn sample_response(&self, direct: bool) {
self.inner
.sample_infohashes_responses
.fetch_add(1, Ordering::Relaxed);
let source = if direct {
&self.inner.sample_infohashes_direct_responses
} else {
&self.inner.sample_infohashes_snapshot_responses
};
source.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_filtered(&self, count: usize) {
self.inner
.sample_infohashes_hashes_filtered
.fetch_add(count as u64, 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
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_peer_succeeded(&self) {
self.inner
.metadata_peer_succeeded
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_success_from(&self, source: DiscoverySource) {
let counter = match source {
DiscoverySource::AnnouncePeer => &self.inner.metadata_success_from_announce,
DiscoverySource::SampleDirect => &self.inner.metadata_success_from_sample_direct,
DiscoverySource::SampleSnapshot => &self.inner.metadata_success_from_sample_snapshot,
DiscoverySource::ActiveLookup => &self.inner.metadata_success_from_active_lookup,
};
counter.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_peer_failed(&self) {
self.inner
.metadata_peer_failed
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_peer_timeout(&self) {
self.inner
.metadata_peer_timeouts
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_connect_failed(&self) {
self.inner
.metadata_connect_failed
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_no_extension(&self) {
self.inner
.metadata_no_extension
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_peer_failure_cache_hit(&self) {
self.inner
.metadata_peer_failure_cache_hits
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn set_metadata_peer_failure_cache_entries(&self, count: usize) {
self.inner
.metadata_peer_failure_cache_entries
.store(count, Ordering::Relaxed);
}
pub(crate) fn set_node_pool_size(&self, size: usize) {
self.inner.node_pool_size.store(size, Ordering::Relaxed);
}
pub(crate) fn set_find_node_in_flight(&self, count: usize) {
self.inner
.find_node_in_flight
.store(count, Ordering::Relaxed);
}
pub(crate) fn set_find_node_effective_rate(&self, rate: u32) {
self.inner
.find_node_effective_rate_per_sec
.store(rate, Ordering::Relaxed);
}
pub(crate) fn query_new(&self) {
self.inner.queries_new.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn query_revisit(&self) {
self.inner.queries_revisit.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn query_bootstrap(&self) {
self.inner.queries_bootstrap.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn response(&self) {
self.inner.responses.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn unmatched_response(&self) {
self.inner
.unmatched_responses
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn timeout(&self) {
self.inner.timeouts.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn send_failure(&self) {
self.inner.send_failures.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn crawl_event_dropped_response(&self) {
self.inner
.crawl_events_dropped_response
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn crawl_event_dropped_discovered(&self) {
self.inner
.crawl_events_dropped_discovered
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn udp_received(&self) {
self.inner.udp_received.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn udp_received_bytes(&self, bytes: usize) {
self.inner
.udp_rx_bytes
.fetch_add(bytes as u64, Ordering::Relaxed);
}
pub(crate) fn udp_sent(&self, bytes: usize) {
self.inner.udp_tx_packets.fetch_add(1, Ordering::Relaxed);
self.inner
.udp_tx_bytes
.fetch_add(bytes as u64, Ordering::Relaxed);
}
pub(crate) fn inbound_query(&self, query: &str) {
let counter = match query {
"ping" => &self.inner.inbound_ping,
"find_node" => &self.inner.inbound_find_node,
"get_peers" => &self.inner.inbound_get_peers,
"announce_peer" => &self.inner.inbound_announce_peer,
_ => &self.inner.inbound_other,
};
counter.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn response_normal(&self) {
self.inner.response_normal.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn response_send_failed(&self) {
self.inner
.response_send_failed
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn announce_accepted(&self) {
self.inner.announce_accepted.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn announce_invalid_token(&self) {
self.inner
.announce_invalid_token
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn announce_filtered(&self) {
self.inner.announce_filtered.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn node_admitted(&self) {
self.inner.node_admitted.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn node_replaced(&self) {
self.inner.node_replaced.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn node_dropped_duplicate(&self) {
self.inner
.node_dropped_duplicate
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn node_dropped_rate_limited(&self) {
self.inner
.node_dropped_rate_limited
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn node_dropped_invalid(&self) {
self.inner
.node_dropped_invalid
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_bytes_downloaded(&self, bytes: usize) {
self.inner
.metadata_bytes_downloaded
.fetch_add(bytes as u64, Ordering::Relaxed);
}
pub(crate) fn metadata_failure_timeout(&self) {
self.inner
.metadata_failure_timeout
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_failure_connect(&self) {
self.inner
.metadata_failure_connect
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_failure_no_extension(&self) {
self.inner
.metadata_failure_no_extension
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_failure_send(&self) {
self.inner
.metadata_failure_send
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_failure_size_limit(&self) {
self.inner
.metadata_failure_size_limit
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_failure_sha1(&self) {
self.inner
.metadata_failure_sha1
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_failure_parse(&self) {
self.inner
.metadata_failure_parse
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn metadata_failure_other(&self) {
self.inner
.metadata_failure_other
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn peer_cache_hit_timeout(&self) {
self.inner
.peer_cache_timeout_hits
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn peer_cache_hit_connect(&self) {
self.inner
.peer_cache_connect_hits
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn observe_metadata_queue_wait(&self, millis: u64) {
self.inner.queue_wait_ms.record(millis);
}
pub(crate) fn observe_metadata_fetch_duration(&self, millis: u64) {
self.inner.fetch_duration_ms.record(millis);
}
pub(crate) fn observe_metadata_size(&self, bytes: usize) {
self.inner.metadata_size_bytes.record(bytes as u64);
}
pub(crate) fn udp_queue_full(&self) {
self.inner.udp_queue_full.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn udp_invalid(&self) {
self.inner.udp_invalid.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn udp_response_rate_limited(&self) {
self.inner
.udp_responses_rate_limited
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn udp_response_priority_reserved(&self) {
self.inner
.udp_responses_priority_reserved
.fetch_add(1, Ordering::Relaxed);
}
}
src/crawler/src/runtime_stats/tests.rs
+208
View File
@@ -0,0 +1,208 @@
// 负责验证 DHT 运行指标快照和直方图统计行为
use super::*;
use crate::types::DiscoverySource;
#[test]
fn cloned_handle_updates_one_snapshot() {
let stats = DhtRuntimeStats::with_limits(DhtRuntimeLimits {
metadata_queue: 100,
node_pool: 1_000,
node_pool_low_watermark: 10,
find_node_in_flight: 512,
initial_find_node_rate: 200,
hash_ingress_queue: 20,
crawl_priority_queue: 30,
crawl_discovery_queue: 40,
});
let writer = stats.clone();
writer.hash_received();
writer.hash_ingress_dropped();
writer.set_hash_ingress_queue_depth(11);
writer.set_crawl_priority_queue_depth(12);
writer.set_crawl_discovery_queue_depth(13);
writer.set_metadata_queue(42, 7);
writer.metadata_queue_inserted();
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.peer_lookup_preferred_succeeded();
writer.peer_lookup_fallback();
writer.configure_sample_candidate_queue(20);
writer.set_sample_candidate_queue_depth(14);
writer.sample_candidate_routed();
writer.sample_candidate_fallback();
writer.sample_query(true);
writer.sample_response(true);
writer.sample_timeout();
writer.sample_send_failed();
writer.sample_response_dropped();
writer.sample_hash_discovered();
writer.sample_hash_filtered(1);
writer.sample_hash_duplicate();
writer.sample_hash_dropped();
writer.metadata_peer_attempt();
writer.metadata_peer_succeeded();
writer.metadata_success_from(DiscoverySource::AnnouncePeer);
writer.metadata_success_from(DiscoverySource::SampleDirect);
writer.metadata_success_from(DiscoverySource::SampleSnapshot);
writer.metadata_success_from(DiscoverySource::ActiveLookup);
writer.metadata_peer_failed();
writer.metadata_peer_timeout();
writer.metadata_connect_failed();
writer.metadata_no_extension();
writer.metadata_peer_failure_cache_hit();
writer.set_metadata_peer_failure_cache_entries(9);
writer.set_node_pool_size(321);
writer.set_find_node_in_flight(12);
writer.set_find_node_effective_rate(150);
writer.query_new();
writer.query_revisit();
writer.query_bootstrap();
writer.response();
writer.unmatched_response();
writer.timeout();
writer.send_failure();
writer.crawl_event_dropped_response();
writer.crawl_event_dropped_discovered();
writer.udp_received();
writer.udp_queue_full();
writer.udp_invalid();
writer.udp_response_rate_limited();
writer.udp_response_priority_reserved();
assert_eq!(
stats.snapshot(),
DhtRuntimeSnapshot {
hashes_received: 1,
hash_ingress_dropped: 1,
hash_ingress_queue_depth: 11,
hash_ingress_queue_capacity: 20,
crawl_priority_queue_depth: 12,
crawl_priority_queue_capacity: 30,
crawl_discovery_queue_depth: 13,
crawl_discovery_queue_capacity: 40,
metadata_queue_depth: 42,
metadata_queue_max: 100,
metadata_in_flight: 7,
metadata_queue_inserted: 1,
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,
peer_lookup_preferred_succeeded: 1,
peer_lookup_fallbacks: 1,
sample_candidate_queue_depth: 14,
sample_candidate_queue_capacity: 20,
sample_candidates_routed: 1,
sample_candidates_fallback: 1,
sample_infohashes_queries: 1,
sample_infohashes_direct_queries: 1,
sample_infohashes_snapshot_queries: 0,
sample_infohashes_responses: 1,
sample_infohashes_direct_responses: 1,
sample_infohashes_snapshot_responses: 0,
sample_infohashes_timeouts: 1,
sample_infohashes_send_failures: 1,
sample_infohashes_response_dropped: 1,
sample_infohashes_hashes_discovered: 1,
sample_infohashes_hashes_filtered: 1,
sample_infohashes_hashes_duplicate: 1,
sample_infohashes_hashes_dropped: 1,
metadata_peer_attempts: 1,
metadata_peer_succeeded: 1,
metadata_success_from_announce: 1,
metadata_success_from_sample_direct: 1,
metadata_success_from_sample_snapshot: 1,
metadata_success_from_active_lookup: 1,
metadata_peer_failed: 1,
metadata_peer_timeouts: 1,
metadata_connect_failed: 1,
metadata_no_extension: 1,
metadata_peer_failure_cache_hits: 1,
metadata_peer_failure_cache_entries: 9,
node_pool_size: 321,
node_pool_capacity: 1_000,
node_pool_low_watermark: 10,
find_node_in_flight: 12,
find_node_in_flight_max: 512,
find_node_effective_rate_per_sec: 150,
queries_new: 1,
queries_revisit: 1,
queries_bootstrap: 1,
responses: 1,
unmatched_responses: 1,
timeouts: 1,
send_failures: 1,
crawl_events_dropped_response: 1,
crawl_events_dropped_discovered: 1,
udp_received: 1,
udp_queue_full: 1,
udp_invalid: 1,
udp_responses_rate_limited: 1,
udp_responses_priority_reserved: 1,
}
);
}
#[test]
fn observability_snapshot_tracks_fixed_histograms_and_categories() {
let stats = DhtRuntimeStats::default();
stats.udp_received();
stats.udp_received_bytes(128);
stats.udp_sent(64);
stats.inbound_query("ping");
stats.inbound_query("unknown");
stats.node_admitted();
stats.metadata_bytes_downloaded(1_024);
stats.metadata_failure_parse();
stats.observe_metadata_queue_wait(75);
stats.observe_metadata_queue_wait(600);
stats.observe_metadata_fetch_duration(1_500);
stats.observe_metadata_size(70_000);
let snapshot = stats.observability_snapshot();
assert_eq!(snapshot.udp_rx_packets, 1);
assert_eq!(snapshot.udp_rx_bytes, 128);
assert_eq!(snapshot.udp_tx_packets, 1);
assert_eq!(snapshot.udp_tx_bytes, 64);
assert_eq!(snapshot.inbound_ping, 1);
assert_eq!(snapshot.inbound_other, 1);
assert_eq!(snapshot.node_admitted, 1);
assert_eq!(snapshot.metadata_bytes_downloaded, 1_024);
assert_eq!(snapshot.metadata_failure_parse, 1);
assert_eq!(snapshot.queue_wait_ms.percentile(0.50), Some(100));
assert_eq!(snapshot.queue_wait_ms.percentile(0.95), Some(1_000));
assert_eq!(snapshot.fetch_duration_ms.percentile(0.95), Some(2_000));
assert_eq!(snapshot.metadata_size_bytes.percentile(0.50), Some(131_072));
}
src/crawler/src/scheduler.rs
+10 -764
View File
@@ -1,19 +1,17 @@
// 负责管理有界 Metadata 任务队列 Peer 竞速和完成回调
use crate::metadata::{FetchedMetadata, MetadataFetchOutcome, RbitFetcher};
use crate::metadata::RbitFetcher;
use crate::peer_lookup::PeerLookupRequest;
#[cfg(test)]
use crate::runtime_stats::DhtRuntimeLimits;
use crate::runtime_stats::DhtRuntimeStats;
use crate::types::{
DiscoverySource, HashDiscovered, MetadataFetchCompletion, MetadataFetchCompletionStatus,
TorrentInfo,
HashDiscovered, MetadataFetchCompletion, MetadataFetchCompletionStatus, TorrentInfo,
};
use arc_swap::ArcSwapOption;
#[cfg(feature = "metrics")]
use metrics::{counter, gauge, histogram};
use std::collections::{BTreeMap, HashMap, HashSet};
use std::future::Future;
use std::collections::{HashMap, HashSet};
use std::net::SocketAddr;
use std::panic::{AssertUnwindSafe, catch_unwind};
use std::sync::Arc;
@@ -31,6 +29,12 @@ 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);
mod peer_race;
mod pending_queue;
use peer_race::race_peer_fetches;
use pending_queue::{PeerCandidate, PendingHashQueue, QueuePushKind, QueuedHash};
/// Callback returning whether the application accepted a downloaded torrent.
pub type TorrentAckCallback = Box<dyn Fn(TorrentInfo) -> bool + Send + Sync + 'static>;
/// Callback invoked once when an admitted Metadata job reaches a terminal state.
@@ -68,171 +72,6 @@ pub(crate) struct MetadataSchedulerRuntime {
pub(crate) peer_lookup_tx: Option<mpsc::Sender<PeerLookupRequest>>,
}
#[derive(Debug, Clone, Copy)]
struct PeerCandidate {
addr: SocketAddr,
source: DiscoverySource,
discovered_at: Instant,
}
#[derive(Debug)]
struct QueuedHash {
info_hash: String,
peers: Vec<PeerCandidate>,
queued_at: Instant,
ready_at: Instant,
latest_at: Instant,
order_key: (Instant, u64),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum QueuePushKind {
Inserted,
Updated,
EvictedOldest,
Stale,
}
#[derive(Debug)]
struct PendingHashQueue {
capacity: usize,
ttl: Duration,
entries: HashMap<String, QueuedHash>,
order: BTreeMap<(Instant, u64), String>,
sequence: u64,
}
impl PendingHashQueue {
fn new(capacity: usize, ttl: Duration) -> Self {
Self {
capacity: capacity.max(1),
ttl,
entries: HashMap::new(),
order: BTreeMap::new(),
sequence: 0,
}
}
fn len(&self) -> usize {
self.entries.len()
}
fn is_empty(&self) -> bool {
self.entries.is_empty()
}
#[cfg(test)]
fn contains(&self, info_hash: &str) -> bool {
self.entries.contains_key(info_hash)
}
fn next_order_key(&mut self, at: Instant) -> (Instant, u64) {
self.sequence = self.sequence.wrapping_add(1);
(at, self.sequence)
}
fn push(&mut self, event: HashDiscovered, now: Instant) -> QueuePushKind {
if now
.checked_duration_since(event.discovered_at)
.unwrap_or_default()
> self.ttl
{
return QueuePushKind::Stale;
}
if let Some(mut entry) = self.entries.remove(&event.info_hash) {
self.order.remove(&entry.order_key);
let discovered_at = entry
.peers
.iter()
.find(|peer| peer.addr == event.peer_addr)
.map(|peer| peer.discovered_at.max(event.discovered_at))
.unwrap_or(event.discovered_at);
entry.peers.retain(|peer| peer.addr != event.peer_addr);
entry.peers.push(PeerCandidate {
addr: event.peer_addr,
source: event.source,
discovered_at,
});
entry
.peers
.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);
self.order.insert(entry.order_key, entry.info_hash.clone());
self.entries.insert(entry.info_hash.clone(), entry);
return QueuePushKind::Updated;
}
let mut result = QueuePushKind::Inserted;
if self.entries.len() >= self.capacity {
let Some((&oldest_key, oldest_hash)) = self.order.first_key_value() else {
return QueuePushKind::Stale;
};
if event.discovered_at <= oldest_key.0 {
return QueuePushKind::Stale;
}
let oldest_hash = oldest_hash.clone();
self.order.remove(&oldest_key);
self.entries.remove(&oldest_hash);
result = QueuePushKind::EvictedOldest;
}
let order_key = self.next_order_key(event.discovered_at);
let info_hash = event.info_hash;
self.order.insert(order_key, info_hash.clone());
self.entries.insert(
info_hash.clone(),
QueuedHash {
info_hash,
peers: vec![PeerCandidate {
addr: event.peer_addr,
source: event.source,
discovered_at: event.discovered_at,
}],
queued_at: now,
ready_at: now + PEER_COALESCE_WINDOW,
latest_at: event.discovered_at,
order_key,
},
);
result
}
fn remove(&mut self, info_hash: &str) -> Option<QueuedHash> {
let entry = self.entries.remove(info_hash)?;
self.order.remove(&entry.order_key);
Some(entry)
}
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)| {
let entry = self.entries.get(hash)?;
(!in_flight.contains_key(hash) && entry.ready_at <= now).then(|| hash.clone())
})?;
self.remove(&info_hash)
}
fn expire(&mut self, now: Instant) -> usize {
let mut expired = 0;
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;
}
let oldest_hash = oldest_hash.clone();
self.order.remove(&oldest_key);
self.entries.remove(&oldest_hash);
expired += 1;
}
expired
}
}
#[derive(Debug)]
struct MetadataJob {
info_hash: String,
@@ -767,598 +606,5 @@ 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())
}
fn fetch_gate<F, Fut>(callback: F) -> Arc<ArcSwapOption<MetadataFetchCallback>>
where
F: Fn(String) -> Fut + Send + Sync + 'static,
Fut: std::future::Future<Output = bool> + Send + 'static,
{
let holder = Arc::new(ArcSwapOption::empty());
let callback: Arc<MetadataFetchCallback> =
Arc::new(Box::new(move |hash| Box::pin(callback(hash))));
holder.store(Some(callback));
holder
}
fn completion_callback<F>(callback: F) -> Arc<ArcSwapOption<MetadataCompletionCallback>>
where
F: Fn(MetadataFetchCompletion) + Send + Sync + 'static,
{
let holder = Arc::new(ArcSwapOption::empty());
let callback: Arc<MetadataCompletionCallback> = Arc::new(Box::new(callback));
holder.store(Some(callback));
holder
}
fn event(hash: &str, port: u16, discovered_at: Instant) -> HashDiscovered {
HashDiscovered {
info_hash: hash.to_string(),
peer_addr: SocketAddr::from(([127, 0, 0, 1], port)),
source: DiscoverySource::AnnouncePeer,
discovered_at,
}
}
#[test]
fn runtime_stats_track_queue_outcomes_and_depth() {
let (_hash_tx, hash_rx) = mpsc::channel(4);
let stats = DhtRuntimeStats::with_limits(DhtRuntimeLimits {
metadata_queue: 2,
..DhtRuntimeLimits::default()
});
let scheduler = MetadataScheduler::new_with_runtime_stats(
hash_rx,
Arc::new(RbitFetcher::new(1)),
MetadataSchedulerLimits {
queue_size: 2,
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 start = Instant::now();
let mut queue = PendingHashQueue::new(2, HASH_QUEUE_TTL);
let mut in_flight = HashMap::new();
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,
start
.checked_sub(HASH_QUEUE_TTL + Duration::from_secs(1))
.unwrap(),
),
);
scheduler.sync_queue_len(queue.len(), 1);
let snapshot = stats.snapshot();
assert_eq!(snapshot.metadata_queue_inserted, 3);
assert_eq!(snapshot.metadata_queue_deduplicated, 1);
assert_eq!(snapshot.metadata_queue_evicted, 1);
assert_eq!(snapshot.metadata_queue_stale, 1);
assert_eq!(snapshot.metadata_queue_depth, 3);
assert_eq!(snapshot.metadata_queue_max, 2);
assert_eq!(snapshot.metadata_in_flight, 1);
}
#[test]
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..13 {
assert_ne!(
queue.push(
event(
"hash",
1000 + offset,
start + Duration::from_secs(offset as u64)
),
start + Duration::from_secs(offset as u64),
),
QueuePushKind::Stale
);
}
queue.push(
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, (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)),
source: DiscoverySource::AnnouncePeer,
discovered_at: now,
},
PeerCandidate {
addr: SocketAddr::from(([127, 0, 0, 1], 1001)),
source: DiscoverySource::SampleDirect,
discovered_at: now,
},
PeerCandidate {
addr: SocketAddr::from(([127, 0, 0, 1], 1002)),
source: DiscoverySource::ActiveLookup,
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)),
source: DiscoverySource::AnnouncePeer,
discovered_at: now,
}];
let live = PeerCandidate {
addr: SocketAddr::from(([127, 0, 0, 1], 1001)),
source: DiscoverySource::ActiveLookup,
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]
fn full_queue_evicts_oldest_for_newer_hash() {
let start = Instant::now();
let mut queue = PendingHashQueue::new(2, Duration::from_secs(60));
queue.push(event("old", 1000, start), start);
queue.push(
event("middle", 1001, start + Duration::from_secs(1)),
start + Duration::from_secs(1),
);
let result = queue.push(
event("new", 1002, start + Duration::from_secs(2)),
start + Duration::from_secs(2),
);
assert_eq!(result, QueuePushKind::EvictedOldest);
assert!(!queue.contains("old"));
assert!(queue.contains("middle"));
assert!(queue.contains("new"));
}
#[test]
fn queue_preserves_peer_discovery_source() {
let start = Instant::now();
let mut queue = PendingHashQueue::new(2, Duration::from_secs(60));
let mut direct = event("direct", 1000, start);
direct.source = DiscoverySource::SampleDirect;
queue.push(direct, start);
let entry = queue
.pop_newest_ready(start + Duration::from_secs(1), &HashMap::new())
.unwrap();
assert_eq!(entry.peers[0].source, DiscoverySource::SampleDirect);
}
#[test]
fn queue_expires_stale_hashes_and_pops_newest_first() {
let start = Instant::now();
let mut queue = PendingHashQueue::new(4, Duration::from_secs(10));
queue.push(event("older", 1000, start), start);
queue.push(
event("newer", 1001, start + Duration::from_secs(1)),
start + Duration::from_secs(1),
);
let in_flight = HashMap::new();
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());
}
#[tokio::test]
async fn gate_rejection_does_not_emit_completion() {
let (hash_tx, hash_rx) = mpsc::channel(4);
let gate_calls = Arc::new(AtomicUsize::new(0));
let gate_calls_for_callback = gate_calls.clone();
let gate = fetch_gate(move |_| {
gate_calls_for_callback.fetch_add(1, Ordering::Relaxed);
async { false }
});
let (completion_tx, mut completion_rx) = mpsc::unbounded_channel();
let completion = completion_callback(move |result| {
let _ = completion_tx.send(result);
});
let shutdown = CancellationToken::new();
let scheduler = MetadataScheduler::new(
hash_rx,
Arc::new(RbitFetcher::new(1)),
MetadataSchedulerLimits {
queue_size: 4,
concurrency: 1,
},
MetadataSchedulerCallbacks {
torrent: empty_torrent_callback(),
fetch_gate: gate,
completion,
},
Arc::new(AtomicUsize::new(0)),
shutdown.clone(),
);
let scheduler_task = tokio::spawn(scheduler.run());
hash_tx
.send(event("bad", 1000, Instant::now()))
.await
.unwrap();
tokio::time::timeout(Duration::from_secs(1), async {
while gate_calls.load(Ordering::Relaxed) == 0 {
tokio::task::yield_now().await;
}
})
.await
.unwrap();
assert!(
tokio::time::timeout(Duration::from_millis(50), completion_rx.recv())
.await
.is_err()
);
shutdown.cancel();
scheduler_task.await.unwrap();
}
#[tokio::test]
async fn admitted_failure_emits_one_completion() {
let (hash_tx, hash_rx) = mpsc::channel(4);
let gate = fetch_gate(|_| async { true });
let (completion_tx, mut completion_rx) = mpsc::unbounded_channel();
let completion = completion_callback(move |result| {
let _ = completion_tx.send(result);
});
let shutdown = CancellationToken::new();
let scheduler = MetadataScheduler::new(
hash_rx,
Arc::new(RbitFetcher::new(1)),
MetadataSchedulerLimits {
queue_size: 4,
concurrency: 1,
},
MetadataSchedulerCallbacks {
torrent: empty_torrent_callback(),
fetch_gate: gate,
completion,
},
Arc::new(AtomicUsize::new(0)),
shutdown.clone(),
);
let scheduler_task = tokio::spawn(scheduler.run());
hash_tx
.send(event("bad", 1000, Instant::now()))
.await
.unwrap();
let result = tokio::time::timeout(Duration::from_secs(1), completion_rx.recv())
.await
.unwrap()
.unwrap();
assert_eq!(result.info_hash, "bad");
assert_eq!(result.status, MetadataFetchCompletionStatus::FetchFailed);
assert_eq!(result.attempts, 0);
assert!(
tokio::time::timeout(Duration::from_millis(50), completion_rx.recv())
.await
.is_err()
);
shutdown.cancel();
scheduler_task.await.unwrap();
}
}
mod tests;
src/crawler/src/scheduler/peer_race.rs
+91
View File
@@ -0,0 +1,91 @@
// 负责并发竞速 Metadata Peer 并在首个成功结果后取消剩余任务
use super::{LIVE_PEER_IDLE_GRACE, MAX_METADATA_PEERS_PER_HASH, PeerCandidate};
use crate::metadata::{FetchedMetadata, MetadataFetchOutcome};
use crate::runtime_stats::DhtRuntimeStats;
#[cfg(feature = "metrics")]
use metrics::counter;
use std::future::Future;
use tokio::sync::mpsc;
use tokio::task::JoinSet;
pub(super) 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,
}
}
}
}
}
src/crawler/src/scheduler/pending_queue.rs
+172
View File
@@ -0,0 +1,172 @@
// 负责维护按新鲜度排序并按 infohash 去重的有界待处理队列
use super::{InFlightState, MAX_METADATA_PEERS_PER_HASH, PEER_COALESCE_WINDOW};
use crate::types::{DiscoverySource, HashDiscovered};
use std::collections::{BTreeMap, HashMap};
use std::net::SocketAddr;
use std::time::{Duration, Instant};
#[derive(Debug, Clone, Copy)]
pub(super) struct PeerCandidate {
pub(super) addr: SocketAddr,
pub(super) source: DiscoverySource,
pub(super) discovered_at: Instant,
}
#[derive(Debug)]
pub(super) struct QueuedHash {
pub(super) info_hash: String,
pub(super) peers: Vec<PeerCandidate>,
pub(super) queued_at: Instant,
pub(super) ready_at: Instant,
pub(super) latest_at: Instant,
pub(super) order_key: (Instant, u64),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum QueuePushKind {
Inserted,
Updated,
EvictedOldest,
Stale,
}
#[derive(Debug)]
pub(super) struct PendingHashQueue {
capacity: usize,
ttl: Duration,
entries: HashMap<String, QueuedHash>,
order: BTreeMap<(Instant, u64), String>,
sequence: u64,
}
impl PendingHashQueue {
pub(super) fn new(capacity: usize, ttl: Duration) -> Self {
Self {
capacity: capacity.max(1),
ttl,
entries: HashMap::new(),
order: BTreeMap::new(),
sequence: 0,
}
}
pub(super) fn len(&self) -> usize {
self.entries.len()
}
pub(super) fn is_empty(&self) -> bool {
self.entries.is_empty()
}
#[cfg(test)]
pub(super) fn contains(&self, info_hash: &str) -> bool {
self.entries.contains_key(info_hash)
}
pub(super) fn next_order_key(&mut self, at: Instant) -> (Instant, u64) {
self.sequence = self.sequence.wrapping_add(1);
(at, self.sequence)
}
pub(super) fn push(&mut self, event: HashDiscovered, now: Instant) -> QueuePushKind {
if now
.checked_duration_since(event.discovered_at)
.unwrap_or_default()
> self.ttl
{
return QueuePushKind::Stale;
}
if let Some(mut entry) = self.entries.remove(&event.info_hash) {
self.order.remove(&entry.order_key);
let discovered_at = entry
.peers
.iter()
.find(|peer| peer.addr == event.peer_addr)
.map(|peer| peer.discovered_at.max(event.discovered_at))
.unwrap_or(event.discovered_at);
entry.peers.retain(|peer| peer.addr != event.peer_addr);
entry.peers.push(PeerCandidate {
addr: event.peer_addr,
source: event.source,
discovered_at,
});
entry
.peers
.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);
self.order.insert(entry.order_key, entry.info_hash.clone());
self.entries.insert(entry.info_hash.clone(), entry);
return QueuePushKind::Updated;
}
let mut result = QueuePushKind::Inserted;
if self.entries.len() >= self.capacity {
let Some((&oldest_key, oldest_hash)) = self.order.first_key_value() else {
return QueuePushKind::Stale;
};
if event.discovered_at <= oldest_key.0 {
return QueuePushKind::Stale;
}
let oldest_hash = oldest_hash.clone();
self.order.remove(&oldest_key);
self.entries.remove(&oldest_hash);
result = QueuePushKind::EvictedOldest;
}
let order_key = self.next_order_key(event.discovered_at);
let info_hash = event.info_hash;
self.order.insert(order_key, info_hash.clone());
self.entries.insert(
info_hash.clone(),
QueuedHash {
info_hash,
peers: vec![PeerCandidate {
addr: event.peer_addr,
source: event.source,
discovered_at: event.discovered_at,
}],
queued_at: now,
ready_at: now + PEER_COALESCE_WINDOW,
latest_at: event.discovered_at,
order_key,
},
);
result
}
pub(super) fn remove(&mut self, info_hash: &str) -> Option<QueuedHash> {
let entry = self.entries.remove(info_hash)?;
self.order.remove(&entry.order_key);
Some(entry)
}
pub(super) 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)| {
let entry = self.entries.get(hash)?;
(!in_flight.contains_key(hash) && entry.ready_at <= now).then(|| hash.clone())
})?;
self.remove(&info_hash)
}
pub(super) fn expire(&mut self, now: Instant) -> usize {
let mut expired = 0;
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;
}
let oldest_hash = oldest_hash.clone();
self.order.remove(&oldest_key);
self.entries.remove(&oldest_hash);
expired += 1;
}
expired
}
}
src/crawler/src/scheduler/tests.rs
+510
View File
@@ -0,0 +1,510 @@
// 负责验证 Metadata 队列调度 Peer 竞速和完成回调行为
use super::*;
use crate::metadata::MetadataFetchOutcome;
use crate::types::DiscoverySource;
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())
}
fn fetch_gate<F, Fut>(callback: F) -> Arc<ArcSwapOption<MetadataFetchCallback>>
where
F: Fn(String) -> Fut + Send + Sync + 'static,
Fut: std::future::Future<Output = bool> + Send + 'static,
{
let holder = Arc::new(ArcSwapOption::empty());
let callback: Arc<MetadataFetchCallback> =
Arc::new(Box::new(move |hash| Box::pin(callback(hash))));
holder.store(Some(callback));
holder
}
fn completion_callback<F>(callback: F) -> Arc<ArcSwapOption<MetadataCompletionCallback>>
where
F: Fn(MetadataFetchCompletion) + Send + Sync + 'static,
{
let holder = Arc::new(ArcSwapOption::empty());
let callback: Arc<MetadataCompletionCallback> = Arc::new(Box::new(callback));
holder.store(Some(callback));
holder
}
fn event(hash: &str, port: u16, discovered_at: Instant) -> HashDiscovered {
HashDiscovered {
info_hash: hash.to_string(),
peer_addr: SocketAddr::from(([127, 0, 0, 1], port)),
source: DiscoverySource::AnnouncePeer,
discovered_at,
}
}
#[test]
fn runtime_stats_track_queue_outcomes_and_depth() {
let (_hash_tx, hash_rx) = mpsc::channel(4);
let stats = DhtRuntimeStats::with_limits(DhtRuntimeLimits {
metadata_queue: 2,
..DhtRuntimeLimits::default()
});
let scheduler = MetadataScheduler::new_with_runtime_stats(
hash_rx,
Arc::new(RbitFetcher::new(1)),
MetadataSchedulerLimits {
queue_size: 2,
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 start = Instant::now();
let mut queue = PendingHashQueue::new(2, HASH_QUEUE_TTL);
let mut in_flight = HashMap::new();
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,
start
.checked_sub(HASH_QUEUE_TTL + Duration::from_secs(1))
.unwrap(),
),
);
scheduler.sync_queue_len(queue.len(), 1);
let snapshot = stats.snapshot();
assert_eq!(snapshot.metadata_queue_inserted, 3);
assert_eq!(snapshot.metadata_queue_deduplicated, 1);
assert_eq!(snapshot.metadata_queue_evicted, 1);
assert_eq!(snapshot.metadata_queue_stale, 1);
assert_eq!(snapshot.metadata_queue_depth, 3);
assert_eq!(snapshot.metadata_queue_max, 2);
assert_eq!(snapshot.metadata_in_flight, 1);
}
#[test]
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..13 {
assert_ne!(
queue.push(
event(
"hash",
1000 + offset,
start + Duration::from_secs(offset as u64)
),
start + Duration::from_secs(offset as u64),
),
QueuePushKind::Stale
);
}
queue.push(
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, (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)),
source: DiscoverySource::AnnouncePeer,
discovered_at: now,
},
PeerCandidate {
addr: SocketAddr::from(([127, 0, 0, 1], 1001)),
source: DiscoverySource::SampleDirect,
discovered_at: now,
},
PeerCandidate {
addr: SocketAddr::from(([127, 0, 0, 1], 1002)),
source: DiscoverySource::ActiveLookup,
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)),
source: DiscoverySource::AnnouncePeer,
discovered_at: now,
}];
let live = PeerCandidate {
addr: SocketAddr::from(([127, 0, 0, 1], 1001)),
source: DiscoverySource::ActiveLookup,
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]
fn full_queue_evicts_oldest_for_newer_hash() {
let start = Instant::now();
let mut queue = PendingHashQueue::new(2, Duration::from_secs(60));
queue.push(event("old", 1000, start), start);
queue.push(
event("middle", 1001, start + Duration::from_secs(1)),
start + Duration::from_secs(1),
);
let result = queue.push(
event("new", 1002, start + Duration::from_secs(2)),
start + Duration::from_secs(2),
);
assert_eq!(result, QueuePushKind::EvictedOldest);
assert!(!queue.contains("old"));
assert!(queue.contains("middle"));
assert!(queue.contains("new"));
}
#[test]
fn queue_preserves_peer_discovery_source() {
let start = Instant::now();
let mut queue = PendingHashQueue::new(2, Duration::from_secs(60));
let mut direct = event("direct", 1000, start);
direct.source = DiscoverySource::SampleDirect;
queue.push(direct, start);
let entry = queue
.pop_newest_ready(start + Duration::from_secs(1), &HashMap::new())
.unwrap();
assert_eq!(entry.peers[0].source, DiscoverySource::SampleDirect);
}
#[test]
fn queue_expires_stale_hashes_and_pops_newest_first() {
let start = Instant::now();
let mut queue = PendingHashQueue::new(4, Duration::from_secs(10));
queue.push(event("older", 1000, start), start);
queue.push(
event("newer", 1001, start + Duration::from_secs(1)),
start + Duration::from_secs(1),
);
let in_flight = HashMap::new();
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());
}
#[tokio::test]
async fn gate_rejection_does_not_emit_completion() {
let (hash_tx, hash_rx) = mpsc::channel(4);
let gate_calls = Arc::new(AtomicUsize::new(0));
let gate_calls_for_callback = gate_calls.clone();
let gate = fetch_gate(move |_| {
gate_calls_for_callback.fetch_add(1, Ordering::Relaxed);
async { false }
});
let (completion_tx, mut completion_rx) = mpsc::unbounded_channel();
let completion = completion_callback(move |result| {
let _ = completion_tx.send(result);
});
let shutdown = CancellationToken::new();
let scheduler = MetadataScheduler::new(
hash_rx,
Arc::new(RbitFetcher::new(1)),
MetadataSchedulerLimits {
queue_size: 4,
concurrency: 1,
},
MetadataSchedulerCallbacks {
torrent: empty_torrent_callback(),
fetch_gate: gate,
completion,
},
Arc::new(AtomicUsize::new(0)),
shutdown.clone(),
);
let scheduler_task = tokio::spawn(scheduler.run());
hash_tx
.send(event("bad", 1000, Instant::now()))
.await
.unwrap();
tokio::time::timeout(Duration::from_secs(1), async {
while gate_calls.load(Ordering::Relaxed) == 0 {
tokio::task::yield_now().await;
}
})
.await
.unwrap();
assert!(
tokio::time::timeout(Duration::from_millis(50), completion_rx.recv())
.await
.is_err()
);
shutdown.cancel();
scheduler_task.await.unwrap();
}
#[tokio::test]
async fn admitted_failure_emits_one_completion() {
let (hash_tx, hash_rx) = mpsc::channel(4);
let gate = fetch_gate(|_| async { true });
let (completion_tx, mut completion_rx) = mpsc::unbounded_channel();
let completion = completion_callback(move |result| {
let _ = completion_tx.send(result);
});
let shutdown = CancellationToken::new();
let scheduler = MetadataScheduler::new(
hash_rx,
Arc::new(RbitFetcher::new(1)),
MetadataSchedulerLimits {
queue_size: 4,
concurrency: 1,
},
MetadataSchedulerCallbacks {
torrent: empty_torrent_callback(),
fetch_gate: gate,
completion,
},
Arc::new(AtomicUsize::new(0)),
shutdown.clone(),
);
let scheduler_task = tokio::spawn(scheduler.run());
hash_tx
.send(event("bad", 1000, Instant::now()))
.await
.unwrap();
let result = tokio::time::timeout(Duration::from_secs(1), completion_rx.recv())
.await
.unwrap()
.unwrap();
assert_eq!(result.info_hash, "bad");
assert_eq!(result.status, MetadataFetchCompletionStatus::FetchFailed);
assert_eq!(result.attempts, 0);
assert!(
tokio::time::timeout(Duration::from_millis(50), completion_rx.recv())
.await
.is_err()
);
shutdown.cancel();
scheduler_task.await.unwrap();
}
src/search/src/storage/rocks.rs
+15 -782
View File
@@ -1,38 +1,32 @@
// 负责实现 RocksDB 打开配置批量写入精确查询和关闭流程
use std::{
collections::BTreeMap,
path::Path,
sync::{Arc, Mutex},
};
use std::sync::{Arc, Mutex};
use rocksdb::{
BlockBasedOptions, Cache, DB, DBCompressionType, Direction, IteratorMode, Options,
SliceTransform, WriteBatch, checkpoint::Checkpoint,
};
use rocksdb::{DB, Direction, IteratorMode, WriteBatch};
use crate::domain::{
ContentFilter, ContentGroupBuilder, InfoHash, MetadataLimits, RejectedMetadata, TorrentRecord,
ContentFilter, ContentGroupBuilder, InfoHash, RejectedMetadata, TorrentRecord,
VerificationResult,
};
use super::{
keys::{
CONTENT_FILTER_FINGERPRINT_KEY, CONTENT_FILTER_MIGRATION_KEY, DATABASE_FORMAT_KEY,
DATABASE_FORMAT_VALUE, VERIFICATION_QUEUE_COUNT_KEY, content_group_key,
content_group_prefix, content_member_key, content_member_prefix, content_members_prefix,
decode_content_member_info_hash, decode_pending_content_key, decode_verification_lease,
decode_verification_task, pending_index_key, pending_index_prefix, rejected_metadata_key,
torrent_key, torrent_prefix, verification_lease_key, verification_lease_prefix,
verification_locator_key, verification_task_key, verification_task_prefix,
VERIFICATION_QUEUE_COUNT_KEY, content_group_key, content_group_prefix, content_member_key,
content_members_prefix, decode_content_member_info_hash, decode_pending_content_key,
decode_verification_lease, decode_verification_task, pending_index_key,
pending_index_prefix, rejected_metadata_key, torrent_key, verification_lease_key,
verification_lease_prefix, verification_locator_key, verification_task_key,
verification_task_prefix,
},
repository::{
CheckpointSummary, ContentGroupTask, ContentVariants, StorageDiagnostics, StorageError,
TorrentRepository, UpsertOutcome, VerificationEnqueueOutcome, VerificationPriority,
VerificationRequest,
ContentGroupTask, ContentVariants, StorageError, TorrentRepository, UpsertOutcome,
VerificationEnqueueOutcome, VerificationPriority, VerificationRequest,
},
};
mod filter_migration;
mod lifecycle;
const DEFAULT_BLOCK_CACHE_BYTES: usize = 64 * 1024 * 1024;
type VerificationTaskEntry = (Box<[u8]>, InfoHash);
@@ -51,261 +45,6 @@ pub struct RocksTorrentRepository {
}
impl RocksTorrentRepository {
pub fn open(path: impl AsRef<Path>) -> Result<Self, StorageError> {
Self::open_with_rules(
path,
MetadataLimits::default().rule_id(),
Arc::new(ContentFilter::default()),
)
}
pub fn open_with_rejection_rule(
path: impl AsRef<Path>,
rejection_rule_id: [u8; 32],
) -> Result<Self, StorageError> {
Self::open_with_rules(path, rejection_rule_id, Arc::new(ContentFilter::default()))
}
pub fn open_with_rules(
path: impl AsRef<Path>,
rejection_rule_id: [u8; 32],
content_filter: Arc<ContentFilter>,
) -> Result<Self, StorageError> {
let mut block_options = BlockBasedOptions::default();
block_options.set_bloom_filter(10.0, false);
let block_cache = Cache::new_lru_cache(DEFAULT_BLOCK_CACHE_BYTES);
block_options.set_block_cache(&block_cache);
let mut options = Options::default();
options.create_if_missing(true);
options.set_compression_type(DBCompressionType::Lz4);
options.set_block_based_table_factory(&block_options);
options.set_prefix_extractor(SliceTransform::create_fixed_prefix(1));
options.set_max_open_files(256);
let mut repository = Self {
db: DB::open(&options, path)?,
write_lock: Mutex::new(()),
rejection_rule_id,
content_filter,
content_filter_changed: false,
};
repository.initialize_format()?;
repository.content_filter_changed = repository.synchronize_content_filter()?;
Ok(repository)
}
pub fn content_filter_changed(&self) -> bool {
self.content_filter_changed
}
pub fn diagnostics(&self) -> Result<StorageDiagnostics, StorageError> {
Ok(StorageDiagnostics {
block_cache_bytes: self.db.property_int_value("rocksdb.block-cache-usage")?,
memtable_bytes: self
.db
.property_int_value("rocksdb.cur-size-all-mem-tables")?,
pending_compaction_bytes: self
.db
.property_int_value("rocksdb.estimate-pending-compaction-bytes")?,
live_sst_bytes: self.db.property_int_value("rocksdb.live-sst-files-size")?,
running_compactions: self
.db
.property_int_value("rocksdb.num-running-compactions")?,
estimated_keys: self.db.property_int_value("rocksdb.estimate-num-keys")?,
})
}
pub fn create_checkpoint(&self, path: impl AsRef<Path>) -> Result<(), StorageError> {
Checkpoint::new(&self.db)?.create_checkpoint(path)?;
Ok(())
}
pub fn validate_checkpoint(path: impl AsRef<Path>) -> Result<CheckpointSummary, StorageError> {
let options = Options::default();
let database = DB::open_for_read_only(&options, path, false)?;
match database.get(DATABASE_FORMAT_KEY)? {
Some(value) if value.as_slice() == DATABASE_FORMAT_VALUE => {}
_ => return Err(StorageError::IncompatibleDatabaseFormat),
}
let prefix = torrent_prefix();
let iterator = database.iterator(IteratorMode::From(&prefix, Direction::Forward));
let mut records = 0_u64;
for entry in iterator {
let (key, value) = entry?;
if !key.starts_with(&prefix) {
break;
}
Self::decode(&value)?;
records = records.saturating_add(1);
}
Ok(CheckpointSummary { records })
}
fn initialize_format(&self) -> Result<(), StorageError> {
match self.db.get(DATABASE_FORMAT_KEY)? {
None => self
.db
.put(DATABASE_FORMAT_KEY, DATABASE_FORMAT_VALUE)
.map_err(Into::into),
Some(value) if value.as_slice() == DATABASE_FORMAT_VALUE => Ok(()),
Some(_) => Err(StorageError::IncompatibleDatabaseFormat),
}
}
fn synchronize_content_filter(&self) -> Result<bool, StorageError> {
const MIGRATION_BATCH_SIZE: usize = 1_000;
let fingerprint = self.content_filter.fingerprint();
if self.db.get(CONTENT_FILTER_MIGRATION_KEY)?.is_none()
&& self
.db
.get(CONTENT_FILTER_FINGERPRINT_KEY)?
.is_some_and(|stored| stored.as_slice() == fingerprint)
{
return Ok(false);
}
let _guard = self
.write_lock
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
self.db.put(CONTENT_FILTER_MIGRATION_KEY, fingerprint)?;
self.clear_prefix(content_member_prefix(), MIGRATION_BATCH_SIZE)?;
self.clear_prefix(content_group_prefix(), MIGRATION_BATCH_SIZE)?;
self.clear_prefix(pending_index_prefix(), MIGRATION_BATCH_SIZE)?;
let prefix = torrent_prefix();
let iterator = self
.db
.iterator(IteratorMode::From(&prefix, Direction::Forward));
let mut batch = WriteBatch::default();
let mut groups = BTreeMap::<[u8; 32], u64>::new();
let mut batch_records = 0_usize;
let mut records = 0_u64;
let mut hidden_files = 0_u64;
let mut hidden_records = 0_u64;
for entry in iterator {
let (key, value) = entry?;
if !key.starts_with(&prefix) {
break;
}
let mut record = Self::decode(&value)?;
let old_content_key = record.content_key;
let old_searchable = record.searchable;
let outcome = self.content_filter.apply_derivatives(&mut record)?;
records = records.saturating_add(1);
hidden_files = hidden_files.saturating_add(outcome.hidden_files as u64);
if !outcome.searchable {
hidden_records = hidden_records.saturating_add(1);
}
if old_content_key != record.content_key || old_searchable != record.searchable {
batch.put(&key, Self::encode(&record)?);
}
if record.searchable {
batch.put(
content_member_key(&record.content_key, record.info_hash),
[],
);
*groups.entry(record.content_key).or_default() += 1;
}
batch_records += 1;
if batch_records == MIGRATION_BATCH_SIZE {
self.write_filter_migration_batch(batch, &groups)?;
batch = WriteBatch::default();
groups.clear();
batch_records = 0;
}
}
if batch_records > 0 {
self.write_filter_migration_batch(batch, &groups)?;
}
self.restore_pending_markers(MIGRATION_BATCH_SIZE)?;
let mut finish = WriteBatch::default();
finish.put(CONTENT_FILTER_FINGERPRINT_KEY, fingerprint);
finish.delete(CONTENT_FILTER_MIGRATION_KEY);
self.db.write(finish)?;
tracing::info!(
records,
hidden_files,
hidden_records,
"内容过滤规则已更新并重建内容组"
);
Ok(true)
}
fn clear_prefix(&self, prefix: [u8; 1], limit: usize) -> Result<(), StorageError> {
loop {
let iterator = self
.db
.iterator(IteratorMode::From(&prefix, Direction::Forward));
let mut batch = WriteBatch::default();
let mut deleted = 0_usize;
for entry in iterator.take(limit) {
let (key, _) = entry?;
if !key.starts_with(&prefix) {
break;
}
batch.delete(key);
deleted += 1;
}
if deleted == 0 {
return Ok(());
}
self.db.write(batch)?;
}
}
fn write_filter_migration_batch(
&self,
mut batch: WriteBatch,
groups: &BTreeMap<[u8; 32], u64>,
) -> Result<(), StorageError> {
for (content_key, added) in groups {
let member_count = self
.group_state(content_key)?
.map_or(*added, |state| state.member_count.saturating_add(*added));
batch.put(
content_group_key(content_key),
Self::encode_group(ContentGroupState {
revision: 1,
member_count,
})?,
);
}
self.db.write(batch)?;
Ok(())
}
fn restore_pending_markers(&self, batch_size: usize) -> Result<(), StorageError> {
let prefix = content_group_prefix();
let iterator = self
.db
.iterator(IteratorMode::From(&prefix, Direction::Forward));
let mut batch = WriteBatch::default();
let mut batch_len = 0_usize;
for entry in iterator {
let (key, _) = entry?;
if !key.starts_with(&prefix) {
break;
}
let content_key: [u8; 32] = key[1..]
.try_into()
.map_err(|_| StorageError::CorruptContentGroup)?;
batch.put(pending_index_key(&content_key), 1_u64.to_be_bytes());
batch_len += 1;
if batch_len == batch_size {
self.db.write(batch)?;
batch = WriteBatch::default();
batch_len = 0;
}
}
if batch_len > 0 {
self.db.write(batch)?;
}
Ok(())
}
fn encode(record: &TorrentRecord) -> Result<Vec<u8>, StorageError> {
rmp_serde::to_vec_named(record).map_err(Into::into)
}
@@ -894,510 +633,4 @@ impl TorrentRepository for RocksTorrentRepository {
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use tempfile::TempDir;
use crate::domain::{
AvailabilityStatus, ContentFilter, ContentFilterConfig, FileFilterRule, FileMatchField,
FileMatchKind, FileRuleAction, TorrentFile, VerificationResult, test_record,
};
use super::*;
fn padding_filter() -> Arc<ContentFilter> {
Arc::new(
ContentFilter::compile(ContentFilterConfig {
version: 1,
file_rules: vec![FileFilterRule {
id: "padding".into(),
enabled: true,
field: FileMatchField::FileName,
match_kind: FileMatchKind::Prefix,
value: "_____padding_file_".into(),
case_sensitive: false,
action: FileRuleAction::Hide,
reason: "测试".into(),
}],
})
.unwrap(),
)
}
#[test]
fn record_survives_close_and_reopen() {
let directory = TempDir::new().unwrap();
let expected = test_record(1, 10);
{
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
assert_eq!(
repository.upsert(expected.clone()).unwrap(),
UpsertOutcome::Inserted
);
}
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
assert_eq!(repository.get(expected.info_hash).unwrap(), Some(expected));
}
#[test]
fn checkpoint_is_a_consistent_snapshot_and_can_be_opened_read_only() {
let directory = TempDir::new().unwrap();
let database_path = directory.path().join("rocksdb");
let checkpoint_path = directory.path().join("checkpoint");
let repository = RocksTorrentRepository::open(&database_path).unwrap();
let first = test_record(1, 10);
let second = test_record(2, 20);
repository.upsert(first.clone()).unwrap();
repository.create_checkpoint(&checkpoint_path).unwrap();
repository.upsert(second.clone()).unwrap();
let summary = RocksTorrentRepository::validate_checkpoint(&checkpoint_path).unwrap();
assert_eq!(summary.records, 1);
drop(repository);
let restored = RocksTorrentRepository::open(&checkpoint_path).unwrap();
assert!(restored.get(first.info_hash).unwrap().is_some());
assert!(restored.get(second.info_hash).unwrap().is_none());
}
#[test]
fn filter_keeps_raw_files_but_returns_effective_view() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open_with_rules(
directory.path(),
MetadataLimits::default().rule_id(),
padding_filter(),
)
.unwrap();
let mut record = test_record(1, 10);
record.files.push(TorrentFile {
path: "_____padding_file_1_请升级____".into(),
size: 100,
});
record.total_size += 100;
repository.upsert(record.clone()).unwrap();
let raw = repository.get(record.info_hash).unwrap().unwrap();
assert_eq!(raw.files.len(), 2);
assert_eq!(raw.total_size, 142);
let visible = repository.get_visible(record.info_hash).unwrap().unwrap();
assert_eq!(visible.files.len(), 1);
assert_eq!(visible.total_size, 42);
assert_eq!(repository.pending_index(10).unwrap().len(), 1);
}
#[test]
fn changed_filter_rebuilds_groups_without_deleting_raw_metadata() {
let directory = TempDir::new().unwrap();
let mut record = test_record(2, 10);
record.files.push(TorrentFile {
path: "_____padding_file_1_".into(),
size: 100,
});
record.total_size += 100;
let old_content_key;
{
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
repository.upsert(record.clone()).unwrap();
old_content_key = repository
.get(record.info_hash)
.unwrap()
.unwrap()
.content_key;
}
let repository = RocksTorrentRepository::open_with_rules(
directory.path(),
MetadataLimits::default().rule_id(),
padding_filter(),
)
.unwrap();
assert!(repository.content_filter_changed());
let raw = repository.get(record.info_hash).unwrap().unwrap();
assert_eq!(raw.files.len(), 2);
assert_ne!(raw.content_key, old_content_key);
assert_eq!(
repository
.content_variants(&old_content_key, 0, 10)
.unwrap()
.total,
0
);
assert_eq!(
repository
.content_variants(&raw.content_key, 0, 10)
.unwrap()
.total,
1
);
drop(repository);
let rolled_back = RocksTorrentRepository::open(directory.path()).unwrap();
assert!(rolled_back.content_filter_changed());
let restored = rolled_back.get_visible(record.info_hash).unwrap().unwrap();
assert_eq!(restored.files.len(), 2);
assert_eq!(restored.total_size, 142);
assert_eq!(restored.content_key, old_content_key);
}
#[test]
fn unfinished_filter_migration_is_rebuilt_even_when_fingerprint_matches() {
let directory = TempDir::new().unwrap();
let record = test_record(4, 10);
{
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
repository.upsert(record.clone()).unwrap();
repository
.db
.put(CONTENT_FILTER_MIGRATION_KEY, b"interrupted")
.unwrap();
repository
.db
.delete(content_member_key(&record.content_key, record.info_hash))
.unwrap();
}
let recovered = RocksTorrentRepository::open(directory.path()).unwrap();
assert!(recovered.content_filter_changed());
assert_eq!(
recovered
.content_variants(&record.content_key, 0, 10)
.unwrap()
.total,
1
);
assert!(recovered.pending_index(10).unwrap().len() == 1);
assert!(
recovered
.db
.get(CONTENT_FILTER_MIGRATION_KEY)
.unwrap()
.is_none()
);
}
#[test]
fn torrent_with_only_hidden_files_is_kept_but_not_indexed() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open_with_rules(
directory.path(),
MetadataLimits::default().rule_id(),
padding_filter(),
)
.unwrap();
let mut record = test_record(3, 10);
record.files[0].path = "_____padding_file_1_".into();
repository.upsert(record.clone()).unwrap();
let raw = repository.get(record.info_hash).unwrap().unwrap();
assert!(!raw.searchable);
assert!(repository.get_visible(record.info_hash).unwrap().is_none());
assert!(repository.pending_index(10).unwrap().is_empty());
}
#[test]
fn incompatible_database_format_is_rejected() {
let directory = TempDir::new().unwrap();
{
let mut options = Options::default();
options.create_if_missing(true);
let database = DB::open(&options, directory.path()).unwrap();
database.put(DATABASE_FORMAT_KEY, b"other-format").unwrap();
}
assert!(matches!(
RocksTorrentRepository::open(directory.path()),
Err(StorageError::IncompatibleDatabaseFormat)
));
}
#[test]
fn duplicate_updates_observation_without_creating_another_pending_item() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let first = test_record(2, 10);
let mut second = first.clone();
second.last_seen = 20;
second.source_peers = vec!["127.0.0.2:6881".into()];
assert_eq!(
repository.upsert(first.clone()).unwrap(),
UpsertOutcome::Inserted
);
assert!(repository.contains(first.info_hash).unwrap());
assert_eq!(
repository.upsert(second).unwrap(),
UpsertOutcome::Updated { seen_count: 2 }
);
let stored = repository.get(first.info_hash).unwrap().unwrap();
assert_eq!(stored.first_seen, 10);
assert_eq!(stored.last_seen, 20);
assert_eq!(stored.seen_count, 2);
let tasks = repository.pending_index(10).unwrap();
assert_eq!(tasks.len(), 1);
assert_eq!(tasks[0].content_key, first.content_key);
}
#[test]
fn batch_triage_updates_existing_and_returns_only_unknown_hashes() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let existing = test_record(1, 10);
let unknown = test_record(2, 10).info_hash;
repository.upsert(existing.clone()).unwrap();
let admitted = repository
.filter_unknown_and_observe(&[existing.info_hash, unknown], 20)
.unwrap();
assert_eq!(admitted, vec![unknown]);
let updated = repository.get(existing.info_hash).unwrap().unwrap();
assert_eq!(updated.last_seen, 20);
assert_eq!(updated.seen_count, existing.seen_count + 1);
assert_eq!(repository.pending_index(10).unwrap().len(), 1);
}
#[test]
fn existing_hash_can_be_observed_without_downloading_metadata_again() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let record = test_record(6, 10);
assert!(!repository.observe_existing(record.info_hash, 5).unwrap());
repository.upsert(record.clone()).unwrap();
assert!(repository.observe_existing(record.info_hash, 30).unwrap());
let stored = repository.get(record.info_hash).unwrap().unwrap();
assert_eq!(stored.first_seen, 10);
assert_eq!(stored.last_seen, 30);
assert_eq!(stored.seen_count, 2);
}
#[test]
fn rejection_survives_restart_and_blocks_the_same_rule() {
let directory = TempDir::new().unwrap();
let limits = MetadataLimits::default();
let info_hash = InfoHash::from_bytes([9; 20]);
{
let repository = RocksTorrentRepository::open_with_rejection_rule(
directory.path(),
limits.rule_id(),
)
.unwrap();
repository
.record_rejection(RejectedMetadata::new(
info_hash,
crate::domain::MetadataRejectionReason::TooManyFiles,
limits.rule_id(),
10,
))
.unwrap();
assert!(repository.get(info_hash).unwrap().is_none());
assert!(repository.observe_existing(info_hash, 20).unwrap());
}
let repository =
RocksTorrentRepository::open_with_rejection_rule(directory.path(), limits.rule_id())
.unwrap();
assert!(repository.observe_existing(info_hash, 30).unwrap());
let rejection = repository.rejection(info_hash).unwrap().unwrap();
assert_eq!(rejection.first_rejected_at, 10);
assert_eq!(rejection.last_seen, 30);
assert_eq!(rejection.seen_count, 3);
assert!(repository.pending_index(10).unwrap().is_empty());
}
#[test]
fn changed_rule_releases_old_rejection_and_valid_upsert_removes_it() {
let directory = TempDir::new().unwrap();
let old_limits = MetadataLimits::default();
let info_hash = InfoHash::from_bytes([10; 20]);
{
let repository = RocksTorrentRepository::open_with_rejection_rule(
directory.path(),
old_limits.rule_id(),
)
.unwrap();
repository
.record_rejection(RejectedMetadata::new(
info_hash,
crate::domain::MetadataRejectionReason::TooManyFiles,
old_limits.rule_id(),
10,
))
.unwrap();
}
let new_limits = MetadataLimits {
max_files: old_limits.max_files + 1,
..old_limits
};
let repository = RocksTorrentRepository::open_with_rejection_rule(
directory.path(),
new_limits.rule_id(),
)
.unwrap();
assert!(!repository.observe_existing(info_hash, 20).unwrap());
let mut record = test_record(10, 20);
record.info_hash = info_hash;
repository.upsert(record).unwrap();
assert!(repository.rejection(info_hash).unwrap().is_none());
assert!(repository.get(info_hash).unwrap().is_some());
}
#[test]
fn equal_content_maps_multiple_infohashes_without_merging_records() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let first = test_record(4, 10);
let mut second = test_record(5, 20);
second.content_key = first.content_key;
repository.upsert(first.clone()).unwrap();
repository.upsert(second.clone()).unwrap();
let mut hashes: Vec<_> = repository
.content_variants(&first.content_key, 0, 10)
.unwrap()
.records
.into_iter()
.map(|record| record.info_hash)
.collect();
hashes.sort_unstable_by_key(ToString::to_string);
assert_eq!(hashes, vec![first.info_hash, second.info_hash]);
assert!(repository.get(first.info_hash).unwrap().is_some());
assert!(repository.get(second.info_hash).unwrap().is_some());
}
#[test]
fn marking_indexed_is_atomic_with_removing_pending_marker() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let record = test_record(3, 10);
repository.upsert(record.clone()).unwrap();
let task = repository.pending_index(10).unwrap()[0];
assert!(
repository
.mark_indexed(&record.content_key, task.revision)
.unwrap()
);
assert!(repository.pending_index(10).unwrap().is_empty());
}
#[test]
fn stale_index_revision_cannot_clear_a_newer_update() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let record = test_record(4, 10);
repository.upsert(record.clone()).unwrap();
let stale = repository.pending_index(10).unwrap()[0];
repository.observe_existing(record.info_hash, 20).unwrap();
assert!(
!repository
.mark_indexed(&record.content_key, stale.revision)
.unwrap()
);
let current = repository.pending_index(10).unwrap()[0];
assert!(current.revision > stale.revision);
}
#[test]
fn content_variants_report_exact_group_count() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let first = test_record(4, 10);
let mut second = test_record(5, 20);
second.content_key = first.content_key;
repository.upsert(first.clone()).unwrap();
repository.upsert(second.clone()).unwrap();
let variants = repository
.content_variants(&first.content_key, 1, 10)
.unwrap();
assert_eq!(variants.total, 2);
assert_eq!(variants.records.len(), 1);
assert_eq!(variants.records[0].info_hash, second.info_hash);
}
#[test]
fn full_reindex_restores_pending_markers_for_every_record() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let first = test_record(7, 10);
let mut second = test_record(8, 10);
second.files[0].path = "different.bin".into();
repository.upsert(first.clone()).unwrap();
repository.upsert(second.clone()).unwrap();
for task in repository.pending_index(10).unwrap() {
repository
.mark_indexed(&task.content_key, task.revision)
.unwrap();
}
assert!(repository.pending_index(10).unwrap().is_empty());
assert_eq!(repository.prepare_full_reindex().unwrap(), 2);
assert_eq!(repository.pending_index(10).unwrap().len(), 2);
}
#[test]
fn verification_queue_is_persistent_prioritized_and_updates_record_atomically() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let first = test_record(1, 10);
let second = test_record(2, 10);
let high = test_record(3, 10);
repository.upsert(first.clone()).unwrap();
repository.upsert(second.clone()).unwrap();
repository.upsert(high.clone()).unwrap();
repository
.enqueue_verification(
&[first.info_hash, second.info_hash],
VerificationPriority::Normal,
100,
2,
)
.unwrap();
let outcome = repository
.enqueue_verification(&[high.info_hash], VerificationPriority::High, 101, 2)
.unwrap();
assert_eq!(outcome.accepted, 1);
assert_eq!(repository.verification_queue_len().unwrap(), 2);
drop(repository);
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let claimed = repository.claim_verification(102, 60).unwrap().unwrap();
assert_eq!(claimed.info_hash, high.info_hash);
assert_eq!(claimed.info_hash, high.info_hash);
repository
.finish_verification(
high.info_hash,
VerificationResult {
verified_at: 103,
discovered_peers: 4,
reachable_peers: 2,
},
)
.unwrap();
assert_eq!(repository.verification_queue_len().unwrap(), 1);
let stored = repository.get(high.info_hash).unwrap().unwrap();
assert_eq!(stored.availability.status, AvailabilityStatus::Active);
assert_eq!(stored.availability.reachable_peers, 2);
}
#[test]
fn expired_verification_lease_is_recovered() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let record = test_record(4, 10);
repository.upsert(record.clone()).unwrap();
repository
.enqueue_verification(&[record.info_hash], VerificationPriority::Normal, 100, 1)
.unwrap();
repository.claim_verification(100, 5).unwrap().unwrap();
let recovered = repository.claim_verification(106, 5).unwrap().unwrap();
assert_eq!(recovered.info_hash, record.info_hash);
assert_eq!(repository.verification_queue_len().unwrap(), 1);
}
}
mod tests;
src/search/src/storage/rocks/filter_migration.rs
+168
View File
@@ -0,0 +1,168 @@
// 负责内容过滤规则变更后的派生字段内容组和待索引标记重建
use super::{ContentGroupState, RocksTorrentRepository};
use crate::storage::{
StorageError,
keys::{
CONTENT_FILTER_FINGERPRINT_KEY, CONTENT_FILTER_MIGRATION_KEY, content_group_key,
content_group_prefix, content_member_key, content_member_prefix, pending_index_key,
pending_index_prefix, torrent_prefix,
},
};
use rocksdb::{Direction, IteratorMode, WriteBatch};
use std::collections::BTreeMap;
impl RocksTorrentRepository {
pub(super) fn synchronize_content_filter(&self) -> Result<bool, StorageError> {
const MIGRATION_BATCH_SIZE: usize = 1_000;
let fingerprint = self.content_filter.fingerprint();
if self.db.get(CONTENT_FILTER_MIGRATION_KEY)?.is_none()
&& self
.db
.get(CONTENT_FILTER_FINGERPRINT_KEY)?
.is_some_and(|stored| stored.as_slice() == fingerprint)
{
return Ok(false);
}
let _guard = self
.write_lock
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
self.db.put(CONTENT_FILTER_MIGRATION_KEY, fingerprint)?;
self.clear_prefix(content_member_prefix(), MIGRATION_BATCH_SIZE)?;
self.clear_prefix(content_group_prefix(), MIGRATION_BATCH_SIZE)?;
self.clear_prefix(pending_index_prefix(), MIGRATION_BATCH_SIZE)?;
let prefix = torrent_prefix();
let iterator = self
.db
.iterator(IteratorMode::From(&prefix, Direction::Forward));
let mut batch = WriteBatch::default();
let mut groups = BTreeMap::<[u8; 32], u64>::new();
let mut batch_records = 0_usize;
let mut records = 0_u64;
let mut hidden_files = 0_u64;
let mut hidden_records = 0_u64;
for entry in iterator {
let (key, value) = entry?;
if !key.starts_with(&prefix) {
break;
}
let mut record = Self::decode(&value)?;
let old_content_key = record.content_key;
let old_searchable = record.searchable;
let outcome = self.content_filter.apply_derivatives(&mut record)?;
records = records.saturating_add(1);
hidden_files = hidden_files.saturating_add(outcome.hidden_files as u64);
if !outcome.searchable {
hidden_records = hidden_records.saturating_add(1);
}
if old_content_key != record.content_key || old_searchable != record.searchable {
batch.put(&key, Self::encode(&record)?);
}
if record.searchable {
batch.put(
content_member_key(&record.content_key, record.info_hash),
[],
);
*groups.entry(record.content_key).or_default() += 1;
}
batch_records += 1;
if batch_records == MIGRATION_BATCH_SIZE {
self.write_filter_migration_batch(batch, &groups)?;
batch = WriteBatch::default();
groups.clear();
batch_records = 0;
}
}
if batch_records > 0 {
self.write_filter_migration_batch(batch, &groups)?;
}
self.restore_pending_markers(MIGRATION_BATCH_SIZE)?;
let mut finish = WriteBatch::default();
finish.put(CONTENT_FILTER_FINGERPRINT_KEY, fingerprint);
finish.delete(CONTENT_FILTER_MIGRATION_KEY);
self.db.write(finish)?;
tracing::info!(
records,
hidden_files,
hidden_records,
"内容过滤规则已更新并重建内容组"
);
Ok(true)
}
fn clear_prefix(&self, prefix: [u8; 1], limit: usize) -> Result<(), StorageError> {
loop {
let iterator = self
.db
.iterator(IteratorMode::From(&prefix, Direction::Forward));
let mut batch = WriteBatch::default();
let mut deleted = 0_usize;
for entry in iterator.take(limit) {
let (key, _) = entry?;
if !key.starts_with(&prefix) {
break;
}
batch.delete(key);
deleted += 1;
}
if deleted == 0 {
return Ok(());
}
self.db.write(batch)?;
}
}
fn write_filter_migration_batch(
&self,
mut batch: WriteBatch,
groups: &BTreeMap<[u8; 32], u64>,
) -> Result<(), StorageError> {
for (content_key, added) in groups {
let member_count = self
.group_state(content_key)?
.map_or(*added, |state| state.member_count.saturating_add(*added));
batch.put(
content_group_key(content_key),
Self::encode_group(ContentGroupState {
revision: 1,
member_count,
})?,
);
}
self.db.write(batch)?;
Ok(())
}
fn restore_pending_markers(&self, batch_size: usize) -> Result<(), StorageError> {
let prefix = content_group_prefix();
let iterator = self
.db
.iterator(IteratorMode::From(&prefix, Direction::Forward));
let mut batch = WriteBatch::default();
let mut batch_len = 0_usize;
for entry in iterator {
let (key, _) = entry?;
if !key.starts_with(&prefix) {
break;
}
let content_key: [u8; 32] = key[1..]
.try_into()
.map_err(|_| StorageError::CorruptContentGroup)?;
batch.put(pending_index_key(&content_key), 1_u64.to_be_bytes());
batch_len += 1;
if batch_len == batch_size {
self.db.write(batch)?;
batch = WriteBatch::default();
batch_len = 0;
}
}
if batch_len > 0 {
self.db.write(batch)?;
}
Ok(())
}
}
src/search/src/storage/rocks/lifecycle.rs
+120
View File
@@ -0,0 +1,120 @@
// 负责 RocksDB 打开格式校验运行诊断和检查点管理
use super::{DEFAULT_BLOCK_CACHE_BYTES, RocksTorrentRepository};
use crate::domain::{ContentFilter, MetadataLimits};
use crate::storage::{
CheckpointSummary, StorageDiagnostics, StorageError,
keys::{DATABASE_FORMAT_KEY, DATABASE_FORMAT_VALUE, torrent_prefix},
};
use rocksdb::{
BlockBasedOptions, Cache, DB, DBCompressionType, Direction, IteratorMode, Options,
SliceTransform, checkpoint::Checkpoint,
};
use std::{
path::Path,
sync::{Arc, Mutex},
};
impl RocksTorrentRepository {
pub fn open(path: impl AsRef<Path>) -> Result<Self, StorageError> {
Self::open_with_rules(
path,
MetadataLimits::default().rule_id(),
Arc::new(ContentFilter::default()),
)
}
pub fn open_with_rejection_rule(
path: impl AsRef<Path>,
rejection_rule_id: [u8; 32],
) -> Result<Self, StorageError> {
Self::open_with_rules(path, rejection_rule_id, Arc::new(ContentFilter::default()))
}
pub fn open_with_rules(
path: impl AsRef<Path>,
rejection_rule_id: [u8; 32],
content_filter: Arc<ContentFilter>,
) -> Result<Self, StorageError> {
let mut block_options = BlockBasedOptions::default();
block_options.set_bloom_filter(10.0, false);
let block_cache = Cache::new_lru_cache(DEFAULT_BLOCK_CACHE_BYTES);
block_options.set_block_cache(&block_cache);
let mut options = Options::default();
options.create_if_missing(true);
options.set_compression_type(DBCompressionType::Lz4);
options.set_block_based_table_factory(&block_options);
options.set_prefix_extractor(SliceTransform::create_fixed_prefix(1));
options.set_max_open_files(256);
let mut repository = Self {
db: DB::open(&options, path)?,
write_lock: Mutex::new(()),
rejection_rule_id,
content_filter,
content_filter_changed: false,
};
repository.initialize_format()?;
repository.content_filter_changed = repository.synchronize_content_filter()?;
Ok(repository)
}
pub fn content_filter_changed(&self) -> bool {
self.content_filter_changed
}
pub fn diagnostics(&self) -> Result<StorageDiagnostics, StorageError> {
Ok(StorageDiagnostics {
block_cache_bytes: self.db.property_int_value("rocksdb.block-cache-usage")?,
memtable_bytes: self
.db
.property_int_value("rocksdb.cur-size-all-mem-tables")?,
pending_compaction_bytes: self
.db
.property_int_value("rocksdb.estimate-pending-compaction-bytes")?,
live_sst_bytes: self.db.property_int_value("rocksdb.live-sst-files-size")?,
running_compactions: self
.db
.property_int_value("rocksdb.num-running-compactions")?,
estimated_keys: self.db.property_int_value("rocksdb.estimate-num-keys")?,
})
}
pub fn create_checkpoint(&self, path: impl AsRef<Path>) -> Result<(), StorageError> {
Checkpoint::new(&self.db)?.create_checkpoint(path)?;
Ok(())
}
pub fn validate_checkpoint(path: impl AsRef<Path>) -> Result<CheckpointSummary, StorageError> {
let options = Options::default();
let database = DB::open_for_read_only(&options, path, false)?;
match database.get(DATABASE_FORMAT_KEY)? {
Some(value) if value.as_slice() == DATABASE_FORMAT_VALUE => {}
_ => return Err(StorageError::IncompatibleDatabaseFormat),
}
let prefix = torrent_prefix();
let iterator = database.iterator(IteratorMode::From(&prefix, Direction::Forward));
let mut records = 0_u64;
for entry in iterator {
let (key, value) = entry?;
if !key.starts_with(&prefix) {
break;
}
Self::decode(&value)?;
records = records.saturating_add(1);
}
Ok(CheckpointSummary { records })
}
fn initialize_format(&self) -> Result<(), StorageError> {
match self.db.get(DATABASE_FORMAT_KEY)? {
None => self
.db
.put(DATABASE_FORMAT_KEY, DATABASE_FORMAT_VALUE)
.map_err(Into::into),
Some(value) if value.as_slice() == DATABASE_FORMAT_VALUE => Ok(()),
Some(_) => Err(StorageError::IncompatibleDatabaseFormat),
}
}
}
src/search/src/storage/rocks/tests.rs
+505
View File
@@ -0,0 +1,505 @@
// 负责验证 RocksDB 持久化去重过滤索引状态检查点和验证队列行为
use std::sync::Arc;
use tempfile::TempDir;
use crate::domain::{
AvailabilityStatus, ContentFilter, ContentFilterConfig, FileFilterRule, FileMatchField,
FileMatchKind, FileRuleAction, MetadataLimits, TorrentFile, VerificationResult, test_record,
};
use crate::storage::keys::{CONTENT_FILTER_MIGRATION_KEY, DATABASE_FORMAT_KEY};
use rocksdb::Options;
use super::*;
fn padding_filter() -> Arc<ContentFilter> {
Arc::new(
ContentFilter::compile(ContentFilterConfig {
version: 1,
file_rules: vec![FileFilterRule {
id: "padding".into(),
enabled: true,
field: FileMatchField::FileName,
match_kind: FileMatchKind::Prefix,
value: "_____padding_file_".into(),
case_sensitive: false,
action: FileRuleAction::Hide,
reason: "测试".into(),
}],
})
.unwrap(),
)
}
#[test]
fn record_survives_close_and_reopen() {
let directory = TempDir::new().unwrap();
let expected = test_record(1, 10);
{
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
assert_eq!(
repository.upsert(expected.clone()).unwrap(),
UpsertOutcome::Inserted
);
}
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
assert_eq!(repository.get(expected.info_hash).unwrap(), Some(expected));
}
#[test]
fn checkpoint_is_a_consistent_snapshot_and_can_be_opened_read_only() {
let directory = TempDir::new().unwrap();
let database_path = directory.path().join("rocksdb");
let checkpoint_path = directory.path().join("checkpoint");
let repository = RocksTorrentRepository::open(&database_path).unwrap();
let first = test_record(1, 10);
let second = test_record(2, 20);
repository.upsert(first.clone()).unwrap();
repository.create_checkpoint(&checkpoint_path).unwrap();
repository.upsert(second.clone()).unwrap();
let summary = RocksTorrentRepository::validate_checkpoint(&checkpoint_path).unwrap();
assert_eq!(summary.records, 1);
drop(repository);
let restored = RocksTorrentRepository::open(&checkpoint_path).unwrap();
assert!(restored.get(first.info_hash).unwrap().is_some());
assert!(restored.get(second.info_hash).unwrap().is_none());
}
#[test]
fn filter_keeps_raw_files_but_returns_effective_view() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open_with_rules(
directory.path(),
MetadataLimits::default().rule_id(),
padding_filter(),
)
.unwrap();
let mut record = test_record(1, 10);
record.files.push(TorrentFile {
path: "_____padding_file_1_请升级____".into(),
size: 100,
});
record.total_size += 100;
repository.upsert(record.clone()).unwrap();
let raw = repository.get(record.info_hash).unwrap().unwrap();
assert_eq!(raw.files.len(), 2);
assert_eq!(raw.total_size, 142);
let visible = repository.get_visible(record.info_hash).unwrap().unwrap();
assert_eq!(visible.files.len(), 1);
assert_eq!(visible.total_size, 42);
assert_eq!(repository.pending_index(10).unwrap().len(), 1);
}
#[test]
fn changed_filter_rebuilds_groups_without_deleting_raw_metadata() {
let directory = TempDir::new().unwrap();
let mut record = test_record(2, 10);
record.files.push(TorrentFile {
path: "_____padding_file_1_".into(),
size: 100,
});
record.total_size += 100;
let old_content_key;
{
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
repository.upsert(record.clone()).unwrap();
old_content_key = repository
.get(record.info_hash)
.unwrap()
.unwrap()
.content_key;
}
let repository = RocksTorrentRepository::open_with_rules(
directory.path(),
MetadataLimits::default().rule_id(),
padding_filter(),
)
.unwrap();
assert!(repository.content_filter_changed());
let raw = repository.get(record.info_hash).unwrap().unwrap();
assert_eq!(raw.files.len(), 2);
assert_ne!(raw.content_key, old_content_key);
assert_eq!(
repository
.content_variants(&old_content_key, 0, 10)
.unwrap()
.total,
0
);
assert_eq!(
repository
.content_variants(&raw.content_key, 0, 10)
.unwrap()
.total,
1
);
drop(repository);
let rolled_back = RocksTorrentRepository::open(directory.path()).unwrap();
assert!(rolled_back.content_filter_changed());
let restored = rolled_back.get_visible(record.info_hash).unwrap().unwrap();
assert_eq!(restored.files.len(), 2);
assert_eq!(restored.total_size, 142);
assert_eq!(restored.content_key, old_content_key);
}
#[test]
fn unfinished_filter_migration_is_rebuilt_even_when_fingerprint_matches() {
let directory = TempDir::new().unwrap();
let record = test_record(4, 10);
{
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
repository.upsert(record.clone()).unwrap();
repository
.db
.put(CONTENT_FILTER_MIGRATION_KEY, b"interrupted")
.unwrap();
repository
.db
.delete(content_member_key(&record.content_key, record.info_hash))
.unwrap();
}
let recovered = RocksTorrentRepository::open(directory.path()).unwrap();
assert!(recovered.content_filter_changed());
assert_eq!(
recovered
.content_variants(&record.content_key, 0, 10)
.unwrap()
.total,
1
);
assert!(recovered.pending_index(10).unwrap().len() == 1);
assert!(
recovered
.db
.get(CONTENT_FILTER_MIGRATION_KEY)
.unwrap()
.is_none()
);
}
#[test]
fn torrent_with_only_hidden_files_is_kept_but_not_indexed() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open_with_rules(
directory.path(),
MetadataLimits::default().rule_id(),
padding_filter(),
)
.unwrap();
let mut record = test_record(3, 10);
record.files[0].path = "_____padding_file_1_".into();
repository.upsert(record.clone()).unwrap();
let raw = repository.get(record.info_hash).unwrap().unwrap();
assert!(!raw.searchable);
assert!(repository.get_visible(record.info_hash).unwrap().is_none());
assert!(repository.pending_index(10).unwrap().is_empty());
}
#[test]
fn incompatible_database_format_is_rejected() {
let directory = TempDir::new().unwrap();
{
let mut options = Options::default();
options.create_if_missing(true);
let database = DB::open(&options, directory.path()).unwrap();
database.put(DATABASE_FORMAT_KEY, b"other-format").unwrap();
}
assert!(matches!(
RocksTorrentRepository::open(directory.path()),
Err(StorageError::IncompatibleDatabaseFormat)
));
}
#[test]
fn duplicate_updates_observation_without_creating_another_pending_item() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let first = test_record(2, 10);
let mut second = first.clone();
second.last_seen = 20;
second.source_peers = vec!["127.0.0.2:6881".into()];
assert_eq!(
repository.upsert(first.clone()).unwrap(),
UpsertOutcome::Inserted
);
assert!(repository.contains(first.info_hash).unwrap());
assert_eq!(
repository.upsert(second).unwrap(),
UpsertOutcome::Updated { seen_count: 2 }
);
let stored = repository.get(first.info_hash).unwrap().unwrap();
assert_eq!(stored.first_seen, 10);
assert_eq!(stored.last_seen, 20);
assert_eq!(stored.seen_count, 2);
let tasks = repository.pending_index(10).unwrap();
assert_eq!(tasks.len(), 1);
assert_eq!(tasks[0].content_key, first.content_key);
}
#[test]
fn batch_triage_updates_existing_and_returns_only_unknown_hashes() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let existing = test_record(1, 10);
let unknown = test_record(2, 10).info_hash;
repository.upsert(existing.clone()).unwrap();
let admitted = repository
.filter_unknown_and_observe(&[existing.info_hash, unknown], 20)
.unwrap();
assert_eq!(admitted, vec![unknown]);
let updated = repository.get(existing.info_hash).unwrap().unwrap();
assert_eq!(updated.last_seen, 20);
assert_eq!(updated.seen_count, existing.seen_count + 1);
assert_eq!(repository.pending_index(10).unwrap().len(), 1);
}
#[test]
fn existing_hash_can_be_observed_without_downloading_metadata_again() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let record = test_record(6, 10);
assert!(!repository.observe_existing(record.info_hash, 5).unwrap());
repository.upsert(record.clone()).unwrap();
assert!(repository.observe_existing(record.info_hash, 30).unwrap());
let stored = repository.get(record.info_hash).unwrap().unwrap();
assert_eq!(stored.first_seen, 10);
assert_eq!(stored.last_seen, 30);
assert_eq!(stored.seen_count, 2);
}
#[test]
fn rejection_survives_restart_and_blocks_the_same_rule() {
let directory = TempDir::new().unwrap();
let limits = MetadataLimits::default();
let info_hash = InfoHash::from_bytes([9; 20]);
{
let repository =
RocksTorrentRepository::open_with_rejection_rule(directory.path(), limits.rule_id())
.unwrap();
repository
.record_rejection(RejectedMetadata::new(
info_hash,
crate::domain::MetadataRejectionReason::TooManyFiles,
limits.rule_id(),
10,
))
.unwrap();
assert!(repository.get(info_hash).unwrap().is_none());
assert!(repository.observe_existing(info_hash, 20).unwrap());
}
let repository =
RocksTorrentRepository::open_with_rejection_rule(directory.path(), limits.rule_id())
.unwrap();
assert!(repository.observe_existing(info_hash, 30).unwrap());
let rejection = repository.rejection(info_hash).unwrap().unwrap();
assert_eq!(rejection.first_rejected_at, 10);
assert_eq!(rejection.last_seen, 30);
assert_eq!(rejection.seen_count, 3);
assert!(repository.pending_index(10).unwrap().is_empty());
}
#[test]
fn changed_rule_releases_old_rejection_and_valid_upsert_removes_it() {
let directory = TempDir::new().unwrap();
let old_limits = MetadataLimits::default();
let info_hash = InfoHash::from_bytes([10; 20]);
{
let repository = RocksTorrentRepository::open_with_rejection_rule(
directory.path(),
old_limits.rule_id(),
)
.unwrap();
repository
.record_rejection(RejectedMetadata::new(
info_hash,
crate::domain::MetadataRejectionReason::TooManyFiles,
old_limits.rule_id(),
10,
))
.unwrap();
}
let new_limits = MetadataLimits {
max_files: old_limits.max_files + 1,
..old_limits
};
let repository =
RocksTorrentRepository::open_with_rejection_rule(directory.path(), new_limits.rule_id())
.unwrap();
assert!(!repository.observe_existing(info_hash, 20).unwrap());
let mut record = test_record(10, 20);
record.info_hash = info_hash;
repository.upsert(record).unwrap();
assert!(repository.rejection(info_hash).unwrap().is_none());
assert!(repository.get(info_hash).unwrap().is_some());
}
#[test]
fn equal_content_maps_multiple_infohashes_without_merging_records() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let first = test_record(4, 10);
let mut second = test_record(5, 20);
second.content_key = first.content_key;
repository.upsert(first.clone()).unwrap();
repository.upsert(second.clone()).unwrap();
let mut hashes: Vec<_> = repository
.content_variants(&first.content_key, 0, 10)
.unwrap()
.records
.into_iter()
.map(|record| record.info_hash)
.collect();
hashes.sort_unstable_by_key(ToString::to_string);
assert_eq!(hashes, vec![first.info_hash, second.info_hash]);
assert!(repository.get(first.info_hash).unwrap().is_some());
assert!(repository.get(second.info_hash).unwrap().is_some());
}
#[test]
fn marking_indexed_is_atomic_with_removing_pending_marker() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let record = test_record(3, 10);
repository.upsert(record.clone()).unwrap();
let task = repository.pending_index(10).unwrap()[0];
assert!(
repository
.mark_indexed(&record.content_key, task.revision)
.unwrap()
);
assert!(repository.pending_index(10).unwrap().is_empty());
}
#[test]
fn stale_index_revision_cannot_clear_a_newer_update() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let record = test_record(4, 10);
repository.upsert(record.clone()).unwrap();
let stale = repository.pending_index(10).unwrap()[0];
repository.observe_existing(record.info_hash, 20).unwrap();
assert!(
!repository
.mark_indexed(&record.content_key, stale.revision)
.unwrap()
);
let current = repository.pending_index(10).unwrap()[0];
assert!(current.revision > stale.revision);
}
#[test]
fn content_variants_report_exact_group_count() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let first = test_record(4, 10);
let mut second = test_record(5, 20);
second.content_key = first.content_key;
repository.upsert(first.clone()).unwrap();
repository.upsert(second.clone()).unwrap();
let variants = repository
.content_variants(&first.content_key, 1, 10)
.unwrap();
assert_eq!(variants.total, 2);
assert_eq!(variants.records.len(), 1);
assert_eq!(variants.records[0].info_hash, second.info_hash);
}
#[test]
fn full_reindex_restores_pending_markers_for_every_record() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let first = test_record(7, 10);
let mut second = test_record(8, 10);
second.files[0].path = "different.bin".into();
repository.upsert(first.clone()).unwrap();
repository.upsert(second.clone()).unwrap();
for task in repository.pending_index(10).unwrap() {
repository
.mark_indexed(&task.content_key, task.revision)
.unwrap();
}
assert!(repository.pending_index(10).unwrap().is_empty());
assert_eq!(repository.prepare_full_reindex().unwrap(), 2);
assert_eq!(repository.pending_index(10).unwrap().len(), 2);
}
#[test]
fn verification_queue_is_persistent_prioritized_and_updates_record_atomically() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let first = test_record(1, 10);
let second = test_record(2, 10);
let high = test_record(3, 10);
repository.upsert(first.clone()).unwrap();
repository.upsert(second.clone()).unwrap();
repository.upsert(high.clone()).unwrap();
repository
.enqueue_verification(
&[first.info_hash, second.info_hash],
VerificationPriority::Normal,
100,
2,
)
.unwrap();
let outcome = repository
.enqueue_verification(&[high.info_hash], VerificationPriority::High, 101, 2)
.unwrap();
assert_eq!(outcome.accepted, 1);
assert_eq!(repository.verification_queue_len().unwrap(), 2);
drop(repository);
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let claimed = repository.claim_verification(102, 60).unwrap().unwrap();
assert_eq!(claimed.info_hash, high.info_hash);
assert_eq!(claimed.info_hash, high.info_hash);
repository
.finish_verification(
high.info_hash,
VerificationResult {
verified_at: 103,
discovered_peers: 4,
reachable_peers: 2,
},
)
.unwrap();
assert_eq!(repository.verification_queue_len().unwrap(), 1);
let stored = repository.get(high.info_hash).unwrap().unwrap();
assert_eq!(stored.availability.status, AvailabilityStatus::Active);
assert_eq!(stored.availability.reachable_peers, 2);
}
#[test]
fn expired_verification_lease_is_recovered() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let record = test_record(4, 10);
repository.upsert(record.clone()).unwrap();
repository
.enqueue_verification(&[record.info_hash], VerificationPriority::Normal, 100, 1)
.unwrap();
repository.claim_verification(100, 5).unwrap().unwrap();
let recovered = repository.claim_verification(106, 5).unwrap().unwrap();
assert_eq!(recovered.info_hash, record.info_hash);
assert_eq!(repository.verification_queue_len().unwrap(), 1);
}