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 09:39:58 +08:00
parent 5b83f64193
commit 07e44c401b
31 changed files with 2395 additions and 1980 deletions
Download Patch File Download Diff File Expand all files Collapse all files
AGENTS.md
+13
View File
@@ -77,6 +77,19 @@ Bloom Filter 只能作为前置加速结构不得作为最终去重依据
模块之间通过明确的数据结构和 trait 通信不得跨层直接访问内部实现
## 组合和文件边界
- trait 只用于存储搜索网络回调等真实替换边界 不创建只有一个调用方的抽象基类或通用 Service 层
- 应用入口只负责构造依赖启动任务选择退出原因和按顺序关闭 不承载 worker 循环和指标格式化
- 每个长期任务独立拥有状态和取消令牌 通过明确句柄组合 不共享可变全局状态
- 领域文件按 infohash Metadata 接纳内容聚合和种子状态演进划分 不按接口页面或数据库字段重复定义模型
- 搜索 schema 查询条件文档映射索引执行分别维护 Tantivy 细节不得泄漏到 API 层
- 基准和工具代码同样遵守职责拆分 不因不进入主服务而集中到单个大文件
- 文件长度不是机械拆分标准 单一状态机或单一 adapter 可以集中维护 强拆会产生私有状态穿透时应保持内聚
- 删除没有调用方的预留模块 新扩展在产生真实行为时再创建文件
测试分层位置和依赖规则统一记录在根目录 `TESTING.md`
## 资源约束
- Metadata 下载并发必须可配置
README.md
+2
View File
@@ -19,6 +19,8 @@ opencodes/ 不参与构建的参考项目
百万级性能基线和验收目标见 [`BENCHMARKS.md`](BENCHMARKS.md)
Rust 测试分层和默认验证命令见 [`TESTING.md`](TESTING.md)
应用配置模板见 [`dht-search.example.toml`](dht-search.example.toml)
应用构建运行和 API 文档见 [`dht-search/README.md`](dht-search/README.md)
TESTING.md
+50
View File
@@ -0,0 +1,50 @@
# Rust 测试分层约定
本文档说明测试应该放在哪里以及每一层允许依赖什么
## 单元测试
单个规则私有状态机编码函数和边界计算使用源码文件底部的 `#[cfg(test)] mod tests`
单元测试可以通过 `use super::*` 访问当前模块私有实现 但不应跨多个业务模块组装完整应用
当前示例包括 Metadata 限制 路径校验 内容组代表选择 响应限流和查询分位数
## 组件测试
需要模块私有装配状态的组件测试保留在对应模块中
例如 Axum router 测试需要私有 `ApiState` 和测试用验证入口 因此与 `api` 模块放在一起而不是为了目录形式公开内部 API
RocksDB adapter 测试需要验证原子批处理私有键空间租约和损坏状态 因此保留在 `storage::rocks` 内部
## 集成测试
只使用 crate 公开 API 的跨层契约放在 crate 根目录 `tests/`
`dht-search/tests/storage_search_flow.rs` 从外部组合领域模型 RocksDB repository 和 Tantivy search 验证写入索引查询关闭重开和恢复
集成测试不得依赖 `pub(crate)` 或为测试扩大生产 API 可见性
## 端到端和运行测试
真实 DHT 网络长时间运行远程部署和浏览器交互依赖外部环境 不放入默认 `cargo test`
这类测试的条件命令结果和验收结论记录到 `TODOS.md` `BENCHMARKS.md` 或部署文档
## 性能基准
可重复的规模测量使用独立 `dht-benchmark` 二进制
基准模块按参数数据集工作负载采样报告和编排拆分 单元测试只验证确定性生成分位数退避和格式化规则
性能结论必须来自 release 构建 默认 debug 运行只用于流程冒烟
## 默认验证命令
```powershell
$env:LIBCLANG_PATH = "$PWD\.tools\libclang\clang\native"
cargo fmt --all --check
cargo test --workspace --all-targets --all-features
cargo clippy --workspace --all-targets --all-features -- -D warnings
```
TODOS.md
+6
View File
@@ -311,6 +311,12 @@
## 当前下一步
- [x] 按职责拆分应用编排索引 worker 运行监控搜索查询构建和 Tantivy 文档映射
- [x] 按领域边界拆分 infohash Metadata 校验内容聚合和种子活跃状态
- [x] 将 DHT 响应限流器从服务器编排中提取为独立组合组件
- [x] 将规模基准拆分为参数数据集工作负载采样报告和编排模块
- [x] 明确单元组件集成端到端和性能测试层级并增加公开 API 集成测试
使用真实采集数据进行一小时持续运行并记录内存队列磁盘网络和验证指标
随后扩展为二十四小时持续运行并根据数据决定资源参数和正则查询优化优先级
dht-crawler/src/lib.rs
+1
View File
@@ -21,6 +21,7 @@ mod node_pool;
mod peer_lookup;
/// Serializable BEP-5 KRPC wire types.
pub mod protocol;
mod response_limiter;
mod routing_snapshot;
mod runtime_stats;
mod sample_infohashes;
dht-crawler/src/response_limiter.rs
+268
View File
@@ -0,0 +1,268 @@
// 负责限制 DHT 响应的包速率字节速率单来源速率和优先保留预算
use std::{
collections::VecDeque,
net::SocketAddr,
time::{Duration, Instant},
};
use ahash::AHashMap;
use crate::budget::RateBucket;
#[derive(Clone, Copy)]
struct SourceResponseWindow {
started_at: Instant,
last_seen: Instant,
count: u32,
}
pub(crate) struct WorkerResponseLimiter {
regular_packets: RateBucket,
regular_bytes: RateBucket,
priority_packets: RateBucket,
priority_bytes: RateBucket,
per_source_rate: u32,
sources: AHashMap<SocketAddr, SourceResponseWindow>,
source_expiry: VecDeque<(Instant, SocketAddr)>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum ResponsePermit {
Regular,
PriorityReserve,
Rejected,
}
impl WorkerResponseLimiter {
pub(crate) fn new(
packet_rate: u32,
byte_rate: u64,
per_source_rate: u32,
now: Instant,
) -> Self {
let byte_rate = byte_rate.min(u32::MAX as u64) as u32;
let priority_packet_rate = reserve_quota(packet_rate);
let priority_byte_rate = reserve_quota(byte_rate);
let packet_rate = packet_rate.saturating_sub(priority_packet_rate);
let byte_rate = byte_rate.saturating_sub(priority_byte_rate);
Self {
regular_packets: RateBucket::per_second(
packet_rate,
packet_rate.div_ceil(5).max(1),
true,
now,
),
regular_bytes: RateBucket::per_second(
byte_rate,
byte_rate.div_ceil(5).max(512),
true,
now,
),
priority_packets: RateBucket::per_second(
priority_packet_rate,
priority_packet_rate.div_ceil(5).max(1),
true,
now,
),
priority_bytes: RateBucket::per_second(
priority_byte_rate,
priority_byte_rate.div_ceil(5).max(512),
true,
now,
),
per_source_rate,
sources: AHashMap::new(),
source_expiry: VecDeque::new(),
}
}
pub(crate) fn acquire(
&mut self,
addr: SocketAddr,
encoded_len: usize,
is_priority: bool,
now: Instant,
) -> ResponsePermit {
self.expire_sources(now);
if Self::take_budget(
&mut self.regular_packets,
&mut self.regular_bytes,
encoded_len,
now,
) {
if self.acquire_source_slot(addr, now) {
return ResponsePermit::Regular;
}
Self::refund_budget(
&mut self.regular_packets,
&mut self.regular_bytes,
encoded_len,
);
return ResponsePermit::Rejected;
}
if is_priority
&& Self::take_budget(
&mut self.priority_packets,
&mut self.priority_bytes,
encoded_len,
now,
)
{
if self.acquire_source_slot(addr, now) {
return ResponsePermit::PriorityReserve;
}
Self::refund_budget(
&mut self.priority_packets,
&mut self.priority_bytes,
encoded_len,
);
}
ResponsePermit::Rejected
}
fn take_budget(
packets: &mut RateBucket,
bytes: &mut RateBucket,
encoded_len: usize,
now: Instant,
) -> bool {
if !packets.try_take_one(now) {
return false;
}
if !bytes.try_take_exact(encoded_len, now) {
packets.refund_one();
return false;
}
true
}
fn refund_budget(packets: &mut RateBucket, bytes: &mut RateBucket, encoded_len: usize) {
packets.refund_one();
bytes.refund(encoded_len);
}
fn acquire_source_slot(&mut self, addr: SocketAddr, now: Instant) -> bool {
if self.per_source_rate == 0 {
return false;
}
let entry = self.sources.entry(addr).or_insert(SourceResponseWindow {
started_at: now,
last_seen: now,
count: 0,
});
if now
.checked_duration_since(entry.started_at)
.unwrap_or_default()
>= Duration::from_secs(1)
{
entry.started_at = now;
entry.count = 0;
}
if entry.count >= self.per_source_rate {
return false;
}
entry.count += 1;
entry.last_seen = now;
self.source_expiry
.push_back((now + Duration::from_secs(60), addr));
true
}
fn expire_sources(&mut self, now: Instant) {
while let Some((deadline, addr)) = self.source_expiry.front().copied() {
if deadline > now {
break;
}
self.source_expiry.pop_front();
if self.sources.get(&addr).is_some_and(|entry| {
now.checked_duration_since(entry.last_seen)
.unwrap_or_default()
>= Duration::from_secs(60)
}) {
self.sources.remove(&addr);
}
}
}
}
pub(crate) fn split_u32_quota(total: u32, workers: usize, worker: usize) -> u32 {
let workers = workers.max(1) as u32;
total / workers + u32::from((worker as u32) < total % workers)
}
pub(crate) fn split_u64_quota(total: u64, workers: usize, worker: usize) -> u64 {
let workers = workers.max(1) as u64;
total / workers + u64::from((worker as u64) < total % workers)
}
fn reserve_quota(total: u32) -> u32 {
if total == 0 {
0
} else {
total.div_ceil(10).min(total)
}
}
#[cfg(test)]
mod tests {
use std::{
net::{IpAddr, Ipv4Addr},
time::Duration,
};
use super::*;
#[test]
fn enforces_packet_byte_and_source_limits() {
let start = Instant::now();
let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 1);
let mut limiter = WorkerResponseLimiter::new(2, 200, 1, start);
assert_eq!(
limiter.acquire(addr, 50, false, start),
ResponsePermit::Regular
);
assert_eq!(
limiter.acquire(addr, 50, false, start),
ResponsePermit::Rejected
);
assert_eq!(
limiter.acquire(addr, 50, false, start + Duration::from_secs(1)),
ResponsePermit::Regular
);
assert_eq!(
limiter.acquire(addr, 513, false, start + Duration::from_secs(2)),
ResponsePermit::Rejected
);
}
#[test]
fn priority_queries_can_use_the_reserve() {
let start = Instant::now();
let mut limiter = WorkerResponseLimiter::new(10, 1_000, 100, start);
let first = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 1);
let second = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 2);
let priority = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 3);
assert_eq!(
limiter.acquire(first, 50, false, start),
ResponsePermit::Regular
);
assert_eq!(
limiter.acquire(second, 50, false, start),
ResponsePermit::Regular
);
assert_eq!(
limiter.acquire(priority, 50, false, start),
ResponsePermit::Rejected
);
assert_eq!(
limiter.acquire(priority, 50, true, start),
ResponsePermit::PriorityReserve
);
let exhausted = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 4);
assert_eq!(
limiter.acquire(exhausted, 100, true, start),
ResponsePermit::Rejected
);
}
}
dht-crawler/src/server.rs
+5 -258
View File
@@ -1,7 +1,7 @@
// 负责组合 DHT 网络组件回调生命周期和公开运行接口
use crate::addr::is_valid_node_addr;
use crate::budget::{RateBucket, SharedRateBudget};
use crate::budget::SharedRateBudget;
use crate::crawl_config::ResolvedCrawlConfig;
use crate::crawl_engine::CrawlEngine;
use crate::error::Result;
@@ -12,6 +12,9 @@ use crate::peer_lookup::{
PeerLookupHandle, PeerLookupRuntime, is_peer_lookup_tid, spawn_peer_lookup,
};
use crate::protocol::{DhtArgs, DhtMessage};
use crate::response_limiter::{
ResponsePermit, WorkerResponseLimiter, split_u32_quota, split_u64_quota,
};
use crate::runtime_stats::{DhtRuntimeLimits, DhtRuntimeStats};
use crate::sample_infohashes::{
SampleHashAdmissionCallback, SampleInfohashesHandle, SampleInfohashesRuntime,
@@ -25,14 +28,13 @@ use crate::scheduler::{
use crate::types::{DHTOptions, MetadataFetchCompletion, NetMode, NodeTuple, TorrentInfo};
use crate::udp_buffer::UdpBufferPool;
use crate::udp_ingress::{WorkerHandle, spawn_udp_listener};
use ahash::AHashMap;
use arc_swap::ArcSwapOption;
use bytes::BytesMut;
#[cfg(feature = "metrics")]
use metrics::counter;
use rand::RngExt;
use socket2::{Domain, Protocol, Socket, Type};
use std::collections::{HashMap, VecDeque};
use std::collections::HashMap;
use std::future::Future;
use std::hash::{Hash, Hasher};
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
@@ -55,188 +57,6 @@ struct QueryResponse<'a> {
target_id: Option<&'a [u8]>,
}
#[derive(Clone, Copy)]
struct SourceResponseWindow {
started_at: Instant,
last_seen: Instant,
count: u32,
}
struct WorkerResponseLimiter {
regular_packets: RateBucket,
regular_bytes: RateBucket,
priority_packets: RateBucket,
priority_bytes: RateBucket,
per_source_rate: u32,
sources: AHashMap<SocketAddr, SourceResponseWindow>,
source_expiry: VecDeque<(Instant, SocketAddr)>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ResponsePermit {
Regular,
PriorityReserve,
Rejected,
}
fn reserve_quota(total: u32) -> u32 {
if total == 0 {
0
} else {
total.div_ceil(10).min(total)
}
}
impl WorkerResponseLimiter {
fn new(packet_rate: u32, byte_rate: u64, per_source_rate: u32, now: Instant) -> Self {
let byte_rate = byte_rate.min(u32::MAX as u64) as u32;
let priority_packet_rate = reserve_quota(packet_rate);
let priority_byte_rate = reserve_quota(byte_rate);
let packet_rate = packet_rate.saturating_sub(priority_packet_rate);
let byte_rate = byte_rate.saturating_sub(priority_byte_rate);
Self {
regular_packets: RateBucket::per_second(
packet_rate,
packet_rate.div_ceil(5).max(1),
true,
now,
),
regular_bytes: RateBucket::per_second(
byte_rate,
byte_rate.div_ceil(5).max(512),
true,
now,
),
priority_packets: RateBucket::per_second(
priority_packet_rate,
priority_packet_rate.div_ceil(5).max(1),
true,
now,
),
priority_bytes: RateBucket::per_second(
priority_byte_rate,
priority_byte_rate.div_ceil(5).max(512),
true,
now,
),
per_source_rate,
sources: AHashMap::new(),
source_expiry: VecDeque::new(),
}
}
fn acquire(
&mut self,
addr: SocketAddr,
encoded_len: usize,
is_priority: bool,
now: Instant,
) -> ResponsePermit {
self.expire_sources(now);
if Self::take_budget(
&mut self.regular_packets,
&mut self.regular_bytes,
encoded_len,
now,
) {
if self.acquire_source_slot(addr, now) {
return ResponsePermit::Regular;
}
Self::refund_budget(
&mut self.regular_packets,
&mut self.regular_bytes,
encoded_len,
);
return ResponsePermit::Rejected;
}
if is_priority
&& Self::take_budget(
&mut self.priority_packets,
&mut self.priority_bytes,
encoded_len,
now,
)
{
if self.acquire_source_slot(addr, now) {
return ResponsePermit::PriorityReserve;
}
Self::refund_budget(
&mut self.priority_packets,
&mut self.priority_bytes,
encoded_len,
);
}
ResponsePermit::Rejected
}
fn take_budget(
packets: &mut RateBucket,
bytes: &mut RateBucket,
encoded_len: usize,
now: Instant,
) -> bool {
if !packets.try_take_one(now) {
return false;
}
if !bytes.try_take_exact(encoded_len, now) {
packets.refund_one();
return false;
}
true
}
fn refund_budget(packets: &mut RateBucket, bytes: &mut RateBucket, encoded_len: usize) {
packets.refund_one();
bytes.refund(encoded_len);
}
fn acquire_source_slot(&mut self, addr: SocketAddr, now: Instant) -> bool {
if self.per_source_rate == 0 {
return false;
}
let entry = self.sources.entry(addr).or_insert(SourceResponseWindow {
started_at: now,
last_seen: now,
count: 0,
});
if now
.checked_duration_since(entry.started_at)
.unwrap_or_default()
>= Duration::from_secs(1)
{
entry.started_at = now;
entry.count = 0;
}
if entry.count >= self.per_source_rate {
return false;
}
entry.count += 1;
entry.last_seen = now;
self.source_expiry
.push_back((now + Duration::from_secs(60), addr));
true
}
fn expire_sources(&mut self, now: Instant) {
while let Some((deadline, addr)) = self.source_expiry.front().copied() {
if deadline > now {
break;
}
self.source_expiry.pop_front();
if self.sources.get(&addr).is_some_and(|entry| {
now.checked_duration_since(entry.last_seen)
.unwrap_or_default()
>= Duration::from_secs(60)
}) {
self.sources.remove(&addr);
}
}
}
}
#[derive(Debug, Clone)]
/// InfoHash and announcing Peer submitted to the Metadata scheduler.
pub struct HashDiscovered {
@@ -290,16 +110,6 @@ fn create_udp_sock(domain: Domain, ty: Type, addr: SocketAddr) -> std::io::Resul
UdpSocket::from_std(sock.into())
}
fn split_u32_quota(total: u32, workers: usize, worker: usize) -> u32 {
let workers = workers.max(1) as u32;
total / workers + u32::from((worker as u32) < total % workers)
}
fn split_u64_quota(total: u64, workers: usize, worker: usize) -> u64 {
let workers = workers.max(1) as u64;
total / workers + u64::from((worker as u64) < total % workers)
}
impl DHTServer {
/// Validates options, binds configured UDP sockets and constructs bounded pipelines.
///
@@ -996,66 +806,3 @@ impl DHTServer {
token == expected
}
}
#[cfg(test)]
mod response_limiter_tests {
use super::*;
#[test]
fn response_limiter_enforces_packet_byte_and_source_limits() {
let start = Instant::now();
let addr: SocketAddr = "8.8.8.8:6881".parse().unwrap();
let mut limiter = WorkerResponseLimiter::new(2, 200, 1, start);
assert_eq!(
limiter.acquire(addr, 100, false, start),
ResponsePermit::Regular
);
assert_eq!(
limiter.acquire(addr, 100, false, start),
ResponsePermit::Rejected
);
assert_eq!(
limiter.acquire(addr, 100, false, start + Duration::from_secs(1)),
ResponsePermit::Regular
);
assert_eq!(
limiter.acquire(
"1.1.1.1:6881".parse().unwrap(),
513,
false,
start + Duration::from_secs(2)
),
ResponsePermit::Rejected
);
}
#[test]
fn ping_and_get_peers_can_use_the_priority_reserve() {
let start = Instant::now();
let mut limiter = WorkerResponseLimiter::new(10, 1_000, 100, start);
let first: SocketAddr = "8.8.8.8:6881".parse().unwrap();
let second: SocketAddr = "1.1.1.1:6881".parse().unwrap();
let priority: SocketAddr = "9.9.9.9:6881".parse().unwrap();
assert_eq!(
limiter.acquire(first, 50, false, start),
ResponsePermit::Regular
);
assert_eq!(
limiter.acquire(second, 50, false, start),
ResponsePermit::Regular
);
assert_eq!(
limiter.acquire(priority, 50, false, start),
ResponsePermit::Rejected
);
assert_eq!(
limiter.acquire(priority, 50, true, start),
ResponsePermit::PriorityReserve
);
assert_eq!(
limiter.acquire("4.4.4.4:6881".parse().unwrap(), 50, true, start),
ResponsePermit::Rejected
);
}
}
dht-search/src/app.rs
+4 -160
View File
@@ -8,7 +8,7 @@ use std::{
use dht_crawler::DHTServer;
use dht_search::{
domain::{InfoHash, MetadataRejectionReason},
domain::InfoHash,
search::SearchEngine,
storage::{RocksTorrentRepository, TorrentRepository},
};
@@ -19,7 +19,7 @@ use crate::{
config::AppConfig,
crawler::pipeline::PersistencePipeline,
error::AppError,
shutdown, verification,
index_worker, monitor, shutdown, verification,
};
pub(crate) async fn run(config: AppConfig) -> Result<(), AppError> {
@@ -151,7 +151,7 @@ pub(crate) async fn run(config: AppConfig) -> Result<(), AppError> {
);
let monitor_cancel = CancellationToken::new();
let monitor = tokio::spawn(monitor(
let monitor = tokio::spawn(monitor::run(
server.clone(),
ingress,
config.stats_interval_secs,
@@ -159,7 +159,7 @@ pub(crate) async fn run(config: AppConfig) -> Result<(), AppError> {
));
let index_cancel = CancellationToken::new();
let (index_fatal_tx, mut index_fatal) = tokio::sync::oneshot::channel();
let index_task = tokio::spawn(run_indexer(
let index_task = tokio::spawn(index_worker::run(
repository.clone(),
search.clone(),
config.index_batch_size,
@@ -251,162 +251,6 @@ pub(crate) async fn run(config: AppConfig) -> Result<(), AppError> {
run_result
}
async fn run_indexer(
repository: Arc<RocksTorrentRepository>,
search: SearchEngine,
batch_size: usize,
interval: Duration,
cancel: CancellationToken,
fatal: tokio::sync::oneshot::Sender<String>,
) -> Result<(), String> {
let mut ticker = tokio::time::interval(interval);
ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
let mut consecutive_retries = 0_u32;
loop {
tokio::select! {
_ = cancel.cancelled() => break,
_ = ticker.tick() => {
let indexed = index_one_batch(repository.clone(), search.clone(), batch_size).await;
match indexed {
Ok(count) => {
consecutive_retries = 0;
if count > 0 {
tracing::debug!(count, "搜索索引已提交");
}
}
Err(IndexBatchError::Retryable(error)) => {
consecutive_retries = consecutive_retries.saturating_add(1);
let delay = index_retry_delay(consecutive_retries);
tracing::warn!(%error, retry = consecutive_retries, delay_ms = delay.as_millis(), "搜索索引遇到临时 I/O 错误");
tokio::select! {
_ = cancel.cancelled() => break,
_ = tokio::time::sleep(delay) => {}
}
}
Err(IndexBatchError::Fatal(error)) => {
let _ = fatal.send(error.clone());
return Err(error);
}
}
}
}
}
let mut shutdown_retries = 0_u32;
loop {
match index_one_batch(repository.clone(), search.clone(), batch_size).await {
Ok(0) => break,
Ok(_) => shutdown_retries = 0,
Err(IndexBatchError::Retryable(error)) if shutdown_retries < 3 => {
shutdown_retries += 1;
let delay = index_retry_delay(shutdown_retries);
tracing::warn!(%error, retry = shutdown_retries, delay_ms = delay.as_millis(), "关闭前提交搜索索引时遇到临时 I/O 错误");
tokio::time::sleep(delay).await;
}
Err(IndexBatchError::Retryable(error)) => {
tracing::warn!(%error, "关闭前搜索索引仍被占用 待索引状态将在下次启动恢复");
break;
}
Err(IndexBatchError::Fatal(error)) => return Err(error),
}
}
Ok(())
}
enum IndexBatchError {
Retryable(String),
Fatal(String),
}
async fn index_one_batch(
repository: Arc<RocksTorrentRepository>,
search: SearchEngine,
batch_size: usize,
) -> Result<usize, IndexBatchError> {
tokio::task::spawn_blocking(move || {
match search.index_pending(repository.as_ref(), batch_size, unix_timestamp()) {
Ok(count) => Ok(count),
Err(error) if error.is_retryable_io() => {
Err(IndexBatchError::Retryable(error.to_string()))
}
Err(error) => Err(IndexBatchError::Fatal(error.to_string())),
}
})
.await
.map_err(|error| IndexBatchError::Fatal(error.to_string()))?
}
fn index_retry_delay(attempt: u32) -> Duration {
let shift = attempt.saturating_sub(1).min(6);
Duration::from_millis((250_u64 << shift).min(10_000))
}
async fn monitor(
server: DHTServer,
ingress: crate::crawler::pipeline::PersistenceIngress,
interval_secs: u64,
cancel: CancellationToken,
) {
let mut interval = tokio::time::interval(Duration::from_secs(interval_secs));
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
let mut previous_udp_tx = 0;
let mut previous_metadata_attempts = 0;
loop {
tokio::select! {
_ = cancel.cancelled() => break,
_ = interval.tick() => {
let dht = server.runtime_stats().snapshot();
let observability = server.runtime_stats().observability_snapshot();
let storage = ingress.snapshot();
let udp_tx_per_second = observability
.udp_tx_packets
.saturating_sub(previous_udp_tx)
/ interval_secs;
let metadata_connects_per_second = dht
.metadata_peer_attempts
.saturating_sub(previous_metadata_attempts)
/ interval_secs;
previous_udp_tx = observability.udp_tx_packets;
previous_metadata_attempts = dht.metadata_peer_attempts;
tracing::info!(
nodes = dht.node_pool_size,
udp_tx = observability.udp_tx_packets,
udp_tx_per_second,
find_node_queries = dht.queries_new + dht.queries_revisit + dht.queries_bootstrap,
peer_lookup_queries = dht.peer_lookup_queries,
peer_lookup_preferred_succeeded = dht.peer_lookup_preferred_succeeded,
peer_lookup_fallbacks = dht.peer_lookup_fallbacks,
sample_queries = dht.sample_infohashes_queries,
sampled_hashes = dht.sample_infohashes_hashes_discovered,
sampled_hashes_filtered = dht.sample_infohashes_hashes_filtered,
peers = dht.peer_lookup_peers_found,
metadata_connects_per_second,
metadata_in_flight = dht.metadata_in_flight,
metadata_ok = dht.metadata_peer_succeeded,
metadata_failed = dht.metadata_peer_failed,
metadata_filtered = observability
.metadata_failure_size_limit
.saturating_add(storage.filtered.total()),
metadata_filtered_too_many_files = storage
.filtered
.count(MetadataRejectionReason::TooManyFiles),
metadata_filtered_invalid_path = storage
.filtered
.count(MetadataRejectionReason::InvalidPath),
persistence_accepted = storage.accepted,
persistence_inserted = storage.inserted,
persistence_updated = storage.updated,
persistence_rejected_full = storage.rejected_full,
persistence_invalid = storage.invalid,
persistence_failed = storage.failed,
persistence_queue = storage.queue_depth,
"运行状态"
)
}
}
}
}
fn unix_timestamp() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
dht-search/src/bin/dht-benchmark.rs
+15 -609
View File
@@ -1,616 +1,22 @@
// 负责生成可重复的大规模种子数据并测量持久化索引查询磁盘和内存表现
// 负责启动规模基准并把具体阶段委托给独立组件
use std::{
error::Error,
fs,
path::{Path, PathBuf},
time::{Duration, Instant, SystemTime, UNIX_EPOCH},
};
use std::error::Error;
use clap::Parser;
use dht_crawler::{FileInfo, TorrentInfo};
use dht_search::{
domain::{MetadataLimits, TorrentRecord},
search::{SearchEngine, SearchOptions, SearchSort},
storage::{RocksTorrentRepository, TorrentRepository},
};
use serde::Serialize;
const MAX_RECORDS: usize = 10_000_000;
const BASE_TIMESTAMP: u64 = 1_700_000_000;
#[derive(Debug, Parser)]
#[command(name = "storage-search", about = "RocksDB 和 Tantivy 端到端规模基准")]
struct Args {
#[arg(long, default_value_t = 10_000)]
records: usize,
#[arg(long, default_value_t = 1_000)]
generation_batch_size: usize,
#[arg(long, default_value_t = 1_000)]
index_batch_size: usize,
#[arg(long, default_value_t = 20)]
index_max_retries: usize,
#[arg(long, default_value_t = 50)]
query_iterations: usize,
#[arg(long, default_value_t = 5)]
query_warmup: usize,
#[arg(long, default_value_t = 10)]
duplicate_every: usize,
#[arg(long, default_value = "benchmark-data")]
output_dir: PathBuf,
#[arg(long)]
cleanup: bool,
}
#[derive(Debug, Serialize)]
struct BenchmarkReport {
generated_at: u64,
build_profile: String,
target: String,
logical_cpus: usize,
records: usize,
indexed_documents: u64,
duplicate_every: usize,
generation_seconds: f64,
rocksdb_write_seconds: f64,
rocksdb_records_per_second: f64,
index_seconds: f64,
index_documents_per_second: f64,
index_transient_retries: usize,
index_retry_wait_seconds: f64,
rocksdb_bytes: u64,
tantivy_bytes: u64,
total_bytes: u64,
rocksdb_bytes_per_record: f64,
tantivy_bytes_per_document: f64,
peak_memory_bytes: Option<u64>,
queries: Vec<QueryReport>,
}
#[derive(Debug, Serialize)]
struct QueryReport {
name: String,
iterations: usize,
result_count: usize,
mean_micros: u64,
p50_micros: u64,
p95_micros: u64,
p99_micros: u64,
}
struct QueryCase {
name: &'static str,
options: SearchOptions,
}
#[path = "dht-benchmark/config.rs"]
mod config;
#[path = "dht-benchmark/dataset.rs"]
mod dataset;
#[path = "dht-benchmark/metrics.rs"]
mod metrics;
#[path = "dht-benchmark/report.rs"]
mod report;
#[path = "dht-benchmark/runner.rs"]
mod runner;
#[path = "dht-benchmark/workload.rs"]
mod workload;
fn main() -> Result<(), Box<dyn Error>> {
let args = Args::parse();
validate_args(&args)?;
let generated_at = unix_timestamp();
let build_profile = if cfg!(debug_assertions) {
"debug"
} else {
"release"
};
let run_name = format!("records-{}-{generated_at}", args.records);
let runs_dir = args.output_dir.join("runs");
let reports_dir = args.output_dir.join("reports");
let run_dir = runs_dir.join(&run_name);
if run_dir.exists() {
return Err(format!("基准目录已经存在 {}", run_dir.display()).into());
}
fs::create_dir_all(&run_dir)?;
fs::create_dir_all(&reports_dir)?;
println!("DHT Search 规模基准");
println!("构建模式: {build_profile}");
if cfg!(debug_assertions) {
println!("警告: debug 模式仅用于流程验证 性能结论必须使用 --release");
}
println!("数据量: {}", format_integer(args.records as u64));
println!("运行目录: {}", run_dir.display());
let rocksdb_dir = run_dir.join("rocksdb");
let tantivy_dir = run_dir.join("tantivy");
let repository = RocksTorrentRepository::open(&rocksdb_dir)?;
let mut generated_duration = Duration::ZERO;
let mut write_duration = Duration::ZERO;
let progress_step = (args.records / 20).max(1);
for batch_start in (0..args.records).step_by(args.generation_batch_size) {
let batch_end = (batch_start + args.generation_batch_size).min(args.records);
let generation_started = Instant::now();
let records: Vec<_> = (batch_start..batch_end)
.map(|index| generate_record(index, args.duplicate_every))
.collect::<Result<_, _>>()?;
generated_duration += generation_started.elapsed();
let write_started = Instant::now();
for record in records {
repository.upsert(record)?;
}
write_duration += write_started.elapsed();
if batch_end == args.records || batch_end / progress_step != batch_start / progress_step {
eprintln!(
"RocksDB 写入进度: {:>3}% ({}/{})",
batch_end.saturating_mul(100) / args.records,
format_integer(batch_end as u64),
format_integer(args.records as u64)
);
}
}
let search = SearchEngine::open(&tantivy_dir)?;
let index_started = Instant::now();
let mut indexed_documents = 0_usize;
let mut index_transient_retries = 0_usize;
let mut index_retry_wait = Duration::ZERO;
let expected_documents = expected_document_count(args.records, args.duplicate_every);
let index_progress_step = (expected_documents / 20).max(1);
let mut next_index_progress = index_progress_step;
loop {
let mut consecutive_retries = 0_usize;
let indexed = loop {
match search.index_pending(
&repository,
args.index_batch_size,
BASE_TIMESTAMP.saturating_add(args.records as u64),
) {
Ok(indexed) => break indexed,
Err(error) if error.is_retryable_io() => {
consecutive_retries = consecutive_retries.saturating_add(1);
index_transient_retries = index_transient_retries.saturating_add(1);
if consecutive_retries > args.index_max_retries {
return Err(format!(
"Tantivy 临时 IO 错误连续重试超过 {} 次: {error}",
args.index_max_retries
)
.into());
}
let delay = index_retry_delay(consecutive_retries);
index_retry_wait += delay;
eprintln!(
"Tantivy 临时 IO 错误 第 {consecutive_retries} 次重试 等待 {} ms: {error}",
delay.as_millis()
);
std::thread::sleep(delay);
}
Err(error) => return Err(Box::new(error)),
}
};
if indexed == 0 {
break;
}
indexed_documents = indexed_documents.saturating_add(indexed);
if indexed_documents >= next_index_progress || indexed_documents >= expected_documents {
eprintln!(
"Tantivy 索引进度: {:>3}% ({} 个内容文档)",
indexed_documents.saturating_mul(100) / expected_documents.max(1),
format_integer(indexed_documents as u64)
);
next_index_progress = next_index_progress.saturating_add(index_progress_step);
}
}
let index_duration = index_started.elapsed();
let query_cases = query_cases(args.records, args.duplicate_every)?;
let mut query_reports = Vec::with_capacity(query_cases.len());
for case in query_cases {
let report = benchmark_query(&search, case, args.query_warmup, args.query_iterations)?;
print_query(&report);
query_reports.push(report);
}
let indexed_documents = search.num_docs();
let peak_memory_bytes = peak_memory_bytes();
drop(search);
drop(repository);
let rocksdb_bytes = directory_size(&rocksdb_dir)?;
let tantivy_bytes = directory_size(&tantivy_dir)?;
let report = BenchmarkReport {
generated_at,
build_profile: build_profile.to_owned(),
target: format!("{}-{}", std::env::consts::OS, std::env::consts::ARCH),
logical_cpus: std::thread::available_parallelism().map_or(1, usize::from),
records: args.records,
indexed_documents,
duplicate_every: args.duplicate_every,
generation_seconds: generated_duration.as_secs_f64(),
rocksdb_write_seconds: write_duration.as_secs_f64(),
rocksdb_records_per_second: rate(args.records as u64, write_duration),
index_seconds: index_duration.as_secs_f64(),
index_documents_per_second: rate(indexed_documents, index_duration),
index_transient_retries,
index_retry_wait_seconds: index_retry_wait.as_secs_f64(),
rocksdb_bytes,
tantivy_bytes,
total_bytes: rocksdb_bytes.saturating_add(tantivy_bytes),
rocksdb_bytes_per_record: ratio(rocksdb_bytes, args.records as u64),
tantivy_bytes_per_document: ratio(tantivy_bytes, indexed_documents),
peak_memory_bytes,
queries: query_reports,
};
print_summary(&report);
let report_path = reports_dir.join(format!("{run_name}.json"));
fs::write(&report_path, serde_json::to_vec_pretty(&report)?)?;
println!("报告: {}", report_path.display());
if args.cleanup {
fs::remove_dir_all(&run_dir)?;
println!("已清理本次基准数据: {}", run_dir.display());
} else {
println!("基准数据已保留 使用 --cleanup 可在完成后自动删除");
}
Ok(())
}
fn validate_args(args: &Args) -> Result<(), Box<dyn Error>> {
if !(100..=MAX_RECORDS).contains(&args.records) {
return Err(format!("records 必须在 100 到 {MAX_RECORDS} 之间").into());
}
if args.generation_batch_size == 0
|| args.index_batch_size == 0
|| args.index_max_retries == 0
|| args.query_iterations == 0
{
return Err("批量大小和查询次数必须大于零".into());
}
if args.duplicate_every == 1 {
return Err("duplicate-every 必须是零或至少为二".into());
}
Ok(())
}
fn generate_record(index: usize, duplicate_every: usize) -> Result<TorrentRecord, Box<dyn Error>> {
let content_id = content_id(index, duplicate_every);
let file_count = 1 + content_id % 4;
let mut files = Vec::with_capacity(file_count);
let main_size = 64 * 1024 * 1024 + (content_id as u64 % 8_192) * 1_048_576;
files.push(FileInfo {
path: format!("media/category_{}/item_{content_id}.mkv", content_id % 100),
size: main_size,
});
for part in 1..file_count {
files.push(FileInfo {
path: format!("docs/item_{content_id}/part_{part}.txt"),
size: 1_024 + (content_id as u64 + part as u64) % 65_536,
});
}
let total_size = files.iter().try_fold(0_u64, |total, file| {
total
.checked_add(file.size)
.ok_or("生成数据的文件总大小溢出")
})?;
let name = match content_id % 4 {
0 => format!("流浪地球 第{content_id}集 1080p"),
1 => format!("Ubuntu Linux Desktop Build {content_id}"),
2 => format!("Nature Documentary 4K Episode {content_id}"),
_ => format!("Open Source Archive Collection {content_id}"),
};
let digest = blake3::hash(&(index as u64).to_be_bytes());
let info_hash = hex::encode(&digest.as_bytes()[..20]);
TorrentRecord::try_from_with_limits(
TorrentInfo {
info_hash,
magnet_link: String::new(),
name,
total_size,
files,
piece_length: 16_384,
peers: Vec::new(),
timestamp: BASE_TIMESTAMP.saturating_add(index as u64),
},
MetadataLimits::default(),
)
.map_err(Into::into)
}
fn content_id(index: usize, duplicate_every: usize) -> usize {
if duplicate_every >= 2 && (index + 1).is_multiple_of(duplicate_every) {
index.saturating_sub(1)
} else {
index
}
}
fn expected_document_count(records: usize, duplicate_every: usize) -> usize {
if duplicate_every >= 2 {
records.saturating_sub(records / duplicate_every)
} else {
records
}
}
fn index_retry_delay(retry: usize) -> Duration {
let exponent = retry.saturating_sub(1).min(5) as u32;
Duration::from_millis(250_u64.saturating_mul(2_u64.pow(exponent))).min(Duration::from_secs(10))
}
fn query_cases(records: usize, duplicate_every: usize) -> Result<Vec<QueryCase>, Box<dyn Error>> {
let exact_index = records.saturating_sub(1).min(42);
let exact_hash = generate_record(exact_index, duplicate_every)?
.info_hash
.to_string();
Ok(vec![
QueryCase {
name: "中文关键词",
options: SearchOptions {
query: "流浪地球".into(),
limit: 20,
..SearchOptions::default()
},
},
QueryCase {
name: "英文关键词",
options: SearchOptions {
query: "ubuntu desktop".into(),
limit: 20,
..SearchOptions::default()
},
},
QueryCase {
name: "文件路径片段",
options: SearchOptions {
query: "item_42".into(),
limit: 20,
..SearchOptions::default()
},
},
QueryCase {
name: "精确 infohash",
options: SearchOptions {
query: exact_hash,
limit: 20,
..SearchOptions::default()
},
},
QueryCase {
name: "有限状态正则",
options: SearchOptions {
query: "ubuntu.*desktop".into(),
regex: true,
limit: 20,
..SearchOptions::default()
},
},
QueryCase {
name: "最近收录排序",
options: SearchOptions {
limit: 20,
sort: Some(SearchSort::Latest),
..SearchOptions::default()
},
},
QueryCase {
name: "大小扩展名过滤",
options: SearchOptions {
min_size: Some(1024 * 1024 * 1024),
extension: Some("mkv".into()),
limit: 20,
sort: Some(SearchSort::SizeDesc),
..SearchOptions::default()
},
},
])
}
fn benchmark_query(
search: &SearchEngine,
case: QueryCase,
warmup: usize,
iterations: usize,
) -> Result<QueryReport, Box<dyn Error>> {
for _ in 0..warmup {
search.search_with(case.options.clone())?;
}
let mut samples = Vec::with_capacity(iterations);
let mut result_count = 0;
for _ in 0..iterations {
let started = Instant::now();
let page = search.search_with(case.options.clone())?;
samples.push(duration_micros(started.elapsed()));
result_count = page.total;
}
samples.sort_unstable();
let mean = samples.iter().copied().sum::<u64>() / samples.len() as u64;
Ok(QueryReport {
name: case.name.to_owned(),
iterations,
result_count,
mean_micros: mean,
p50_micros: percentile(&samples, 50),
p95_micros: percentile(&samples, 95),
p99_micros: percentile(&samples, 99),
})
}
fn percentile(sorted: &[u64], percentile: usize) -> u64 {
let rank = sorted
.len()
.saturating_mul(percentile)
.div_ceil(100)
.saturating_sub(1)
.min(sorted.len().saturating_sub(1));
sorted[rank]
}
fn duration_micros(duration: Duration) -> u64 {
duration.as_micros().min(u128::from(u64::MAX)) as u64
}
fn rate(items: u64, duration: Duration) -> f64 {
if duration.is_zero() {
0.0
} else {
items as f64 / duration.as_secs_f64()
}
}
fn ratio(bytes: u64, items: u64) -> f64 {
if items == 0 {
0.0
} else {
bytes as f64 / items as f64
}
}
fn directory_size(path: &Path) -> Result<u64, std::io::Error> {
let mut total = 0_u64;
let mut pending = vec![path.to_path_buf()];
while let Some(directory) = pending.pop() {
for entry in fs::read_dir(directory)? {
let entry = entry?;
let metadata = entry.metadata()?;
if metadata.is_dir() {
pending.push(entry.path());
} else if metadata.is_file() {
total = total.saturating_add(metadata.len());
}
}
}
Ok(total)
}
fn print_query(report: &QueryReport) {
println!(
"查询 {:<16} 结果 {:>8} P50 {:>8} µs P95 {:>8} µs P99 {:>8} µs",
report.name,
format_integer(report.result_count as u64),
report.p50_micros,
report.p95_micros,
report.p99_micros
);
}
fn print_summary(report: &BenchmarkReport) {
println!();
println!("基准汇总");
println!("生成耗时: {:.3} 秒", report.generation_seconds);
println!(
"RocksDB 写入: {:.3} 秒 {:.0} 条/秒",
report.rocksdb_write_seconds, report.rocksdb_records_per_second
);
println!(
"Tantivy 索引: {:.3} 秒 {:.0} 文档/秒",
report.index_seconds, report.index_documents_per_second
);
println!(
"Tantivy 临时 IO 重试: {} 次 等待 {:.3} 秒",
report.index_transient_retries, report.index_retry_wait_seconds
);
println!(
"RocksDB: {} 平均 {:.1} 字节/记录",
format_bytes(report.rocksdb_bytes),
report.rocksdb_bytes_per_record
);
println!(
"Tantivy: {} 平均 {:.1} 字节/文档",
format_bytes(report.tantivy_bytes),
report.tantivy_bytes_per_document
);
println!("合计磁盘: {}", format_bytes(report.total_bytes));
if let Some(bytes) = report.peak_memory_bytes {
println!("进程峰值内存: {}", format_bytes(bytes));
} else {
println!("进程峰值内存: 当前平台暂不支持读取");
}
}
fn format_integer(value: u64) -> String {
let digits = value.to_string();
let mut formatted = String::with_capacity(digits.len() + digits.len() / 3);
for (index, character) in digits.chars().enumerate() {
if index > 0 && (digits.len() - index).is_multiple_of(3) {
formatted.push(',');
}
formatted.push(character);
}
formatted
}
fn format_bytes(bytes: u64) -> String {
const UNITS: &[&str] = &["B", "KiB", "MiB", "GiB", "TiB"];
let mut value = bytes as f64;
let mut unit = 0;
while value >= 1024.0 && unit + 1 < UNITS.len() {
value /= 1024.0;
unit += 1;
}
format!("{value:.2} {}", UNITS[unit])
}
fn unix_timestamp() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}
#[cfg(windows)]
fn peak_memory_bytes() -> Option<u64> {
use std::mem::{size_of, zeroed};
use windows_sys::Win32::System::{
ProcessStatus::{GetProcessMemoryInfo, PROCESS_MEMORY_COUNTERS},
Threading::GetCurrentProcess,
};
let mut counters: PROCESS_MEMORY_COUNTERS = unsafe { zeroed() };
let result = unsafe {
GetProcessMemoryInfo(
GetCurrentProcess(),
&mut counters,
size_of::<PROCESS_MEMORY_COUNTERS>() as u32,
)
};
(result != 0).then_some(counters.PeakWorkingSetSize as u64)
}
#[cfg(target_os = "linux")]
fn peak_memory_bytes() -> Option<u64> {
let status = fs::read_to_string("/proc/self/status").ok()?;
let line = status.lines().find(|line| line.starts_with("VmHWM:"))?;
let kibibytes = line.split_whitespace().nth(1)?.parse::<u64>().ok()?;
kibibytes.checked_mul(1024)
}
#[cfg(not(any(windows, target_os = "linux")))]
fn peak_memory_bytes() -> Option<u64> {
None
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn generator_is_deterministic_and_creates_requested_duplicates() {
let first = generate_record(8, 10).unwrap();
let duplicate = generate_record(9, 10).unwrap();
assert_ne!(first.info_hash, duplicate.info_hash);
assert_eq!(first.content_key, duplicate.content_key);
assert_eq!(generate_record(8, 10).unwrap(), first);
assert_eq!(expected_document_count(100, 10), 90);
}
#[test]
fn percentiles_use_nearest_rank() {
let samples: Vec<_> = (1..=100).collect();
assert_eq!(percentile(&samples, 50), 50);
assert_eq!(percentile(&samples, 95), 95);
assert_eq!(percentile(&samples, 99), 99);
}
#[test]
fn byte_and_integer_formatting_are_stable() {
assert_eq!(format_integer(1_234_567), "1,234,567");
assert_eq!(format_bytes(1024), "1.00 KiB");
}
#[test]
fn index_retry_delay_is_bounded() {
assert_eq!(index_retry_delay(1), Duration::from_millis(250));
assert_eq!(index_retry_delay(3), Duration::from_secs(1));
assert!(index_retry_delay(100) <= Duration::from_secs(10));
}
runner::run(config::Args::parse())
}
dht-search/src/bin/dht-benchmark/config.rs
+60
View File
@@ -0,0 +1,60 @@
// 负责解析和校验规模基准命令行参数
use std::{error::Error, path::PathBuf};
use clap::Parser;
const MAX_RECORDS: usize = 10_000_000;
#[derive(Debug, Parser)]
#[command(name = "dht-benchmark", about = "RocksDB 和 Tantivy 端到端规模基准")]
pub(crate) struct Args {
#[arg(long, default_value_t = 10_000)]
pub(crate) records: usize,
#[arg(long, default_value_t = 1_000)]
pub(crate) generation_batch_size: usize,
#[arg(long, default_value_t = 1_000)]
pub(crate) index_batch_size: usize,
#[arg(long, default_value_t = 20)]
pub(crate) index_max_retries: usize,
#[arg(long, default_value_t = 50)]
pub(crate) query_iterations: usize,
#[arg(long, default_value_t = 5)]
pub(crate) query_warmup: usize,
#[arg(long, default_value_t = 10)]
pub(crate) duplicate_every: usize,
#[arg(long, default_value = "benchmark-data")]
pub(crate) output_dir: PathBuf,
#[arg(long)]
pub(crate) cleanup: bool,
}
impl Args {
pub(crate) fn validate(&self) -> Result<(), Box<dyn Error>> {
if !(100..=MAX_RECORDS).contains(&self.records) {
return Err(format!("records 必须在 100 到 {MAX_RECORDS} 之间").into());
}
if self.generation_batch_size == 0
|| self.index_batch_size == 0
|| self.index_max_retries == 0
|| self.query_iterations == 0
{
return Err("批量大小重试次数和查询次数必须大于零".into());
}
if self.duplicate_every == 1 {
return Err("duplicate-every 必须是零或至少为二".into());
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn rejects_ambiguous_duplicate_interval() {
let args = Args::parse_from(["benchmark", "--duplicate-every", "1"]);
assert!(args.validate().is_err());
}
}
dht-search/src/bin/dht-benchmark/dataset.rs
+86
View File
@@ -0,0 +1,86 @@
// 负责确定性生成种子记录和可控比例的相同内容变体
use std::error::Error;
use dht_crawler::{FileInfo, TorrentInfo};
use dht_search::domain::{MetadataLimits, TorrentRecord};
pub(crate) const BASE_TIMESTAMP: u64 = 1_700_000_000;
pub(crate) fn generate_record(
index: usize,
duplicate_every: usize,
) -> Result<TorrentRecord, Box<dyn Error>> {
let content_id = content_id(index, duplicate_every);
let file_count = 1 + content_id % 4;
let mut files = Vec::with_capacity(file_count);
let main_size = 64 * 1024 * 1024 + (content_id as u64 % 8_192) * 1_048_576;
files.push(FileInfo {
path: format!("media/category_{}/item_{content_id}.mkv", content_id % 100),
size: main_size,
});
for part in 1..file_count {
files.push(FileInfo {
path: format!("docs/item_{content_id}/part_{part}.txt"),
size: 1_024 + (content_id as u64 + part as u64) % 65_536,
});
}
let total_size = files.iter().try_fold(0_u64, |total, file| {
total
.checked_add(file.size)
.ok_or("生成数据的文件总大小溢出")
})?;
let name = match content_id % 4 {
0 => format!("流浪地球 第{content_id}集 1080p"),
1 => format!("Ubuntu Linux Desktop Build {content_id}"),
2 => format!("Nature Documentary 4K Episode {content_id}"),
_ => format!("Open Source Archive Collection {content_id}"),
};
let digest = blake3::hash(&(index as u64).to_be_bytes());
let info_hash = hex::encode(&digest.as_bytes()[..20]);
TorrentRecord::try_from_with_limits(
TorrentInfo {
info_hash,
magnet_link: String::new(),
name,
total_size,
files,
piece_length: 16_384,
peers: Vec::new(),
timestamp: BASE_TIMESTAMP.saturating_add(index as u64),
},
MetadataLimits::default(),
)
.map_err(Into::into)
}
pub(crate) fn expected_document_count(records: usize, duplicate_every: usize) -> usize {
if duplicate_every >= 2 {
records.saturating_sub(records / duplicate_every)
} else {
records
}
}
fn content_id(index: usize, duplicate_every: usize) -> usize {
if duplicate_every >= 2 && (index + 1).is_multiple_of(duplicate_every) {
index.saturating_sub(1)
} else {
index
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn generator_is_deterministic_and_creates_requested_duplicates() {
let first = generate_record(8, 10).unwrap();
let duplicate = generate_record(9, 10).unwrap();
assert_ne!(first.info_hash, duplicate.info_hash);
assert_eq!(first.content_key, duplicate.content_key);
assert_eq!(generate_record(8, 10).unwrap(), first);
assert_eq!(expected_document_count(100, 10), 90);
}
}
dht-search/src/bin/dht-benchmark/metrics.rs
+146
View File
@@ -0,0 +1,146 @@
// 负责采样查询延迟计算分位数磁盘占用和进程峰值内存
use std::{
error::Error,
fs,
path::Path,
time::{Duration, Instant},
};
use dht_search::search::SearchEngine;
use serde::Serialize;
use super::workload::QueryCase;
#[derive(Debug, Serialize)]
pub(crate) struct QueryReport {
pub(crate) name: String,
pub(crate) iterations: usize,
pub(crate) result_count: usize,
pub(crate) mean_micros: u64,
pub(crate) p50_micros: u64,
pub(crate) p95_micros: u64,
pub(crate) p99_micros: u64,
}
pub(crate) fn benchmark_query(
search: &SearchEngine,
case: QueryCase,
warmup: usize,
iterations: usize,
) -> Result<QueryReport, Box<dyn Error>> {
for _ in 0..warmup {
search.search_with(case.options.clone())?;
}
let mut samples = Vec::with_capacity(iterations);
let mut result_count = 0;
for _ in 0..iterations {
let started = Instant::now();
let page = search.search_with(case.options.clone())?;
samples.push(duration_micros(started.elapsed()));
result_count = page.total;
}
samples.sort_unstable();
let mean = samples.iter().copied().sum::<u64>() / samples.len() as u64;
Ok(QueryReport {
name: case.name.to_owned(),
iterations,
result_count,
mean_micros: mean,
p50_micros: percentile(&samples, 50),
p95_micros: percentile(&samples, 95),
p99_micros: percentile(&samples, 99),
})
}
pub(crate) fn directory_size(path: &Path) -> Result<u64, std::io::Error> {
let mut total = 0_u64;
let mut pending = vec![path.to_path_buf()];
while let Some(directory) = pending.pop() {
for entry in fs::read_dir(directory)? {
let entry = entry?;
let metadata = entry.metadata()?;
if metadata.is_dir() {
pending.push(entry.path());
} else if metadata.is_file() {
total = total.saturating_add(metadata.len());
}
}
}
Ok(total)
}
pub(crate) fn rate(items: u64, duration: Duration) -> f64 {
if duration.is_zero() {
0.0
} else {
items as f64 / duration.as_secs_f64()
}
}
pub(crate) fn ratio(bytes: u64, items: u64) -> f64 {
if items == 0 {
0.0
} else {
bytes as f64 / items as f64
}
}
fn percentile(sorted: &[u64], percentile: usize) -> u64 {
let rank = sorted
.len()
.saturating_mul(percentile)
.div_ceil(100)
.saturating_sub(1)
.min(sorted.len().saturating_sub(1));
sorted[rank]
}
fn duration_micros(duration: Duration) -> u64 {
duration.as_micros().min(u128::from(u64::MAX)) as u64
}
#[cfg(windows)]
pub(crate) fn peak_memory_bytes() -> Option<u64> {
use std::mem::{size_of, zeroed};
use windows_sys::Win32::System::{
ProcessStatus::{GetProcessMemoryInfo, PROCESS_MEMORY_COUNTERS},
Threading::GetCurrentProcess,
};
let mut counters: PROCESS_MEMORY_COUNTERS = unsafe { zeroed() };
let result = unsafe {
GetProcessMemoryInfo(
GetCurrentProcess(),
&mut counters,
size_of::<PROCESS_MEMORY_COUNTERS>() as u32,
)
};
(result != 0).then_some(counters.PeakWorkingSetSize as u64)
}
#[cfg(target_os = "linux")]
pub(crate) fn peak_memory_bytes() -> Option<u64> {
let status = fs::read_to_string("/proc/self/status").ok()?;
let line = status.lines().find(|line| line.starts_with("VmHWM:"))?;
let kibibytes = line.split_whitespace().nth(1)?.parse::<u64>().ok()?;
kibibytes.checked_mul(1024)
}
#[cfg(not(any(windows, target_os = "linux")))]
pub(crate) fn peak_memory_bytes() -> Option<u64> {
None
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn percentiles_use_nearest_rank() {
let samples: Vec<_> = (1..=100).collect();
assert_eq!(percentile(&samples, 50), 50);
assert_eq!(percentile(&samples, 95), 95);
assert_eq!(percentile(&samples, 99), 99);
}
}
dht-search/src/bin/dht-benchmark/report.rs
+109
View File
@@ -0,0 +1,109 @@
// 负责序列化并展示规模基准的稳定报告格式
use serde::Serialize;
use super::metrics::QueryReport;
#[derive(Debug, Serialize)]
pub(crate) struct BenchmarkReport {
pub(crate) generated_at: u64,
pub(crate) build_profile: String,
pub(crate) target: String,
pub(crate) logical_cpus: usize,
pub(crate) records: usize,
pub(crate) indexed_documents: u64,
pub(crate) duplicate_every: usize,
pub(crate) generation_seconds: f64,
pub(crate) rocksdb_write_seconds: f64,
pub(crate) rocksdb_records_per_second: f64,
pub(crate) index_seconds: f64,
pub(crate) index_documents_per_second: f64,
pub(crate) index_transient_retries: usize,
pub(crate) index_retry_wait_seconds: f64,
pub(crate) rocksdb_bytes: u64,
pub(crate) tantivy_bytes: u64,
pub(crate) total_bytes: u64,
pub(crate) rocksdb_bytes_per_record: f64,
pub(crate) tantivy_bytes_per_document: f64,
pub(crate) peak_memory_bytes: Option<u64>,
pub(crate) queries: Vec<QueryReport>,
}
pub(crate) fn print_query(report: &QueryReport) {
println!(
"查询 {:<16} 结果 {:>8} P50 {:>8} µs P95 {:>8} µs P99 {:>8} µs",
report.name,
format_integer(report.result_count as u64),
report.p50_micros,
report.p95_micros,
report.p99_micros
);
}
pub(crate) fn print_summary(report: &BenchmarkReport) {
println!();
println!("基准汇总");
println!("生成耗时: {:.3} 秒", report.generation_seconds);
println!(
"RocksDB 写入: {:.3} 秒 {:.0} 条/秒",
report.rocksdb_write_seconds, report.rocksdb_records_per_second
);
println!(
"Tantivy 索引: {:.3} 秒 {:.0} 文档/秒",
report.index_seconds, report.index_documents_per_second
);
println!(
"Tantivy 临时 IO 重试: {} 次 等待 {:.3} 秒",
report.index_transient_retries, report.index_retry_wait_seconds
);
println!(
"RocksDB: {} 平均 {:.1} 字节/记录",
format_bytes(report.rocksdb_bytes),
report.rocksdb_bytes_per_record
);
println!(
"Tantivy: {} 平均 {:.1} 字节/文档",
format_bytes(report.tantivy_bytes),
report.tantivy_bytes_per_document
);
println!("合计磁盘: {}", format_bytes(report.total_bytes));
if let Some(bytes) = report.peak_memory_bytes {
println!("进程峰值内存: {}", format_bytes(bytes));
} else {
println!("进程峰值内存: 当前平台暂不支持读取");
}
}
pub(crate) fn format_integer(value: u64) -> String {
let digits = value.to_string();
let mut formatted = String::with_capacity(digits.len() + digits.len() / 3);
for (index, character) in digits.chars().enumerate() {
if index > 0 && (digits.len() - index).is_multiple_of(3) {
formatted.push(',');
}
formatted.push(character);
}
formatted
}
fn format_bytes(bytes: u64) -> String {
const UNITS: &[&str] = &["B", "KiB", "MiB", "GiB", "TiB"];
let mut value = bytes as f64;
let mut unit = 0;
while value >= 1024.0 && unit + 1 < UNITS.len() {
value /= 1024.0;
unit += 1;
}
format!("{value:.2} {}", UNITS[unit])
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn byte_and_integer_formatting_are_stable() {
assert_eq!(format_integer(1_234_567), "1,234,567");
assert_eq!(format_bytes(1024), "1.00 KiB");
}
}
dht-search/src/bin/dht-benchmark/runner.rs
+227
View File
@@ -0,0 +1,227 @@
// 负责按生成写入索引查询报告顺序编排一次完整规模基准
use std::{
error::Error,
fs,
time::{Duration, Instant, SystemTime, UNIX_EPOCH},
};
use dht_search::{
search::SearchEngine,
storage::{RocksTorrentRepository, TorrentRepository},
};
use super::{
config::Args,
dataset::{BASE_TIMESTAMP, expected_document_count, generate_record},
metrics::{benchmark_query, directory_size, peak_memory_bytes, rate, ratio},
report::{BenchmarkReport, format_integer, print_query, print_summary},
workload::query_cases,
};
pub(crate) fn run(args: Args) -> Result<(), Box<dyn Error>> {
args.validate()?;
let generated_at = unix_timestamp();
let build_profile = if cfg!(debug_assertions) {
"debug"
} else {
"release"
};
let run_name = format!("records-{}-{generated_at}", args.records);
let runs_dir = args.output_dir.join("runs");
let reports_dir = args.output_dir.join("reports");
let run_dir = runs_dir.join(&run_name);
if run_dir.exists() {
return Err(format!("基准目录已经存在 {}", run_dir.display()).into());
}
fs::create_dir_all(&run_dir)?;
fs::create_dir_all(&reports_dir)?;
println!("DHT Search 规模基准");
println!("构建模式: {build_profile}");
if cfg!(debug_assertions) {
println!("警告: debug 模式仅用于流程验证 性能结论必须使用 --release");
}
println!("数据量: {}", format_integer(args.records as u64));
println!("运行目录: {}", run_dir.display());
let rocksdb_dir = run_dir.join("rocksdb");
let tantivy_dir = run_dir.join("tantivy");
let repository = RocksTorrentRepository::open(&rocksdb_dir)?;
let (generation_duration, write_duration) = populate(&repository, &args)?;
let search = SearchEngine::open(&tantivy_dir)?;
let index_result = build_index(&repository, &search, &args)?;
let mut query_reports = Vec::new();
for case in query_cases(args.records, args.duplicate_every)? {
let report = benchmark_query(&search, case, args.query_warmup, args.query_iterations)?;
print_query(&report);
query_reports.push(report);
}
let indexed_documents = search.num_docs();
let peak_memory_bytes = peak_memory_bytes();
drop(search);
drop(repository);
let rocksdb_bytes = directory_size(&rocksdb_dir)?;
let tantivy_bytes = directory_size(&tantivy_dir)?;
let report = BenchmarkReport {
generated_at,
build_profile: build_profile.to_owned(),
target: format!("{}-{}", std::env::consts::OS, std::env::consts::ARCH),
logical_cpus: std::thread::available_parallelism().map_or(1, usize::from),
records: args.records,
indexed_documents,
duplicate_every: args.duplicate_every,
generation_seconds: generation_duration.as_secs_f64(),
rocksdb_write_seconds: write_duration.as_secs_f64(),
rocksdb_records_per_second: rate(args.records as u64, write_duration),
index_seconds: index_result.duration.as_secs_f64(),
index_documents_per_second: rate(indexed_documents, index_result.duration),
index_transient_retries: index_result.transient_retries,
index_retry_wait_seconds: index_result.retry_wait.as_secs_f64(),
rocksdb_bytes,
tantivy_bytes,
total_bytes: rocksdb_bytes.saturating_add(tantivy_bytes),
rocksdb_bytes_per_record: ratio(rocksdb_bytes, args.records as u64),
tantivy_bytes_per_document: ratio(tantivy_bytes, indexed_documents),
peak_memory_bytes,
queries: query_reports,
};
print_summary(&report);
let report_path = reports_dir.join(format!("{run_name}.json"));
fs::write(&report_path, serde_json::to_vec_pretty(&report)?)?;
println!("报告: {}", report_path.display());
if args.cleanup {
fs::remove_dir_all(&run_dir)?;
println!("已清理本次基准数据: {}", run_dir.display());
} else {
println!("基准数据已保留 使用 --cleanup 可在完成后自动删除");
}
Ok(())
}
fn populate(
repository: &RocksTorrentRepository,
args: &Args,
) -> Result<(Duration, Duration), Box<dyn Error>> {
let mut generation_duration = Duration::ZERO;
let mut write_duration = Duration::ZERO;
let progress_step = (args.records / 20).max(1);
for batch_start in (0..args.records).step_by(args.generation_batch_size) {
let batch_end = (batch_start + args.generation_batch_size).min(args.records);
let generation_started = Instant::now();
let records: Vec<_> = (batch_start..batch_end)
.map(|index| generate_record(index, args.duplicate_every))
.collect::<Result<_, _>>()?;
generation_duration += generation_started.elapsed();
let write_started = Instant::now();
for record in records {
repository.upsert(record)?;
}
write_duration += write_started.elapsed();
if batch_end == args.records || batch_end / progress_step != batch_start / progress_step {
eprintln!(
"RocksDB 写入进度: {:>3}% ({}/{})",
batch_end.saturating_mul(100) / args.records,
format_integer(batch_end as u64),
format_integer(args.records as u64)
);
}
}
Ok((generation_duration, write_duration))
}
struct IndexResult {
duration: Duration,
transient_retries: usize,
retry_wait: Duration,
}
fn build_index(
repository: &RocksTorrentRepository,
search: &SearchEngine,
args: &Args,
) -> Result<IndexResult, Box<dyn Error>> {
let started = Instant::now();
let mut indexed_documents = 0_usize;
let mut transient_retries = 0_usize;
let mut retry_wait = Duration::ZERO;
let expected_documents = expected_document_count(args.records, args.duplicate_every);
let progress_step = (expected_documents / 20).max(1);
let mut next_progress = progress_step;
loop {
let mut consecutive_retries = 0_usize;
let indexed = loop {
match search.index_pending(
repository,
args.index_batch_size,
BASE_TIMESTAMP.saturating_add(args.records as u64),
) {
Ok(indexed) => break indexed,
Err(error) if error.is_retryable_io() => {
consecutive_retries = consecutive_retries.saturating_add(1);
transient_retries = transient_retries.saturating_add(1);
if consecutive_retries > args.index_max_retries {
return Err(format!(
"Tantivy 临时 IO 错误连续重试超过 {} 次: {error}",
args.index_max_retries
)
.into());
}
let delay = index_retry_delay(consecutive_retries);
retry_wait += delay;
eprintln!(
"Tantivy 临时 IO 错误 第 {consecutive_retries} 次重试 等待 {} ms: {error}",
delay.as_millis()
);
std::thread::sleep(delay);
}
Err(error) => return Err(Box::new(error)),
}
};
if indexed == 0 {
break;
}
indexed_documents = indexed_documents.saturating_add(indexed);
if indexed_documents >= next_progress || indexed_documents >= expected_documents {
eprintln!(
"Tantivy 索引进度: {:>3}% ({} 个内容文档)",
indexed_documents.saturating_mul(100) / expected_documents.max(1),
format_integer(indexed_documents as u64)
);
next_progress = next_progress.saturating_add(progress_step);
}
}
Ok(IndexResult {
duration: started.elapsed(),
transient_retries,
retry_wait,
})
}
fn index_retry_delay(retry: usize) -> Duration {
let exponent = retry.saturating_sub(1).min(5) as u32;
Duration::from_millis(250_u64.saturating_mul(2_u64.pow(exponent))).min(Duration::from_secs(10))
}
fn unix_timestamp() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn index_retry_delay_is_bounded() {
assert_eq!(index_retry_delay(1), Duration::from_millis(250));
assert_eq!(index_retry_delay(3), Duration::from_secs(1));
assert!(index_retry_delay(100) <= Duration::from_secs(10));
}
}
dht-search/src/bin/dht-benchmark/workload.rs
+83
View File
@@ -0,0 +1,83 @@
// 负责定义覆盖关键词路径哈希正则排序和过滤的固定查询工作负载
use std::error::Error;
use dht_search::search::{SearchOptions, SearchSort};
use super::dataset::generate_record;
pub(crate) struct QueryCase {
pub(crate) name: &'static str,
pub(crate) options: SearchOptions,
}
pub(crate) fn query_cases(
records: usize,
duplicate_every: usize,
) -> Result<Vec<QueryCase>, Box<dyn Error>> {
let exact_index = records.saturating_sub(1).min(42);
let exact_hash = generate_record(exact_index, duplicate_every)?
.info_hash
.to_string();
Ok(vec![
QueryCase {
name: "中文关键词",
options: SearchOptions {
query: "流浪地球".into(),
limit: 20,
..SearchOptions::default()
},
},
QueryCase {
name: "英文关键词",
options: SearchOptions {
query: "ubuntu desktop".into(),
limit: 20,
..SearchOptions::default()
},
},
QueryCase {
name: "文件路径片段",
options: SearchOptions {
query: "item_42".into(),
limit: 20,
..SearchOptions::default()
},
},
QueryCase {
name: "精确 infohash",
options: SearchOptions {
query: exact_hash,
limit: 20,
..SearchOptions::default()
},
},
QueryCase {
name: "有限状态正则",
options: SearchOptions {
query: "ubuntu.*desktop".into(),
regex: true,
limit: 20,
..SearchOptions::default()
},
},
QueryCase {
name: "最近收录排序",
options: SearchOptions {
limit: 20,
sort: Some(SearchSort::Latest),
..SearchOptions::default()
},
},
QueryCase {
name: "大小扩展名过滤",
options: SearchOptions {
min_size: Some(1024 * 1024 * 1024),
extension: Some("mkv".into()),
limit: 20,
sort: Some(SearchSort::SizeDesc),
..SearchOptions::default()
},
},
])
}
dht-search/src/crawler/mod.rs
-1
View File
@@ -1,4 +1,3 @@
// 负责组合 DHT 发现 Metadata 下载和持久化提交管线
pub(crate) mod pipeline;
mod worker;
dht-search/src/crawler/worker.rs
-1
View File
@@ -1 +0,0 @@
// 负责消费 infohash 并执行受资源限制的 Metadata 下载任务
dht-search/src/domain/content_group.rs
+186
View File
@@ -0,0 +1,186 @@
// 负责把相同内容的多个 infohash 聚合为稳定且可排序的搜索内容组
use std::cmp::Ordering;
use super::{Availability, AvailabilityStatus, Heat, InfoHash, TorrentRecord};
#[derive(Debug, Clone, PartialEq)]
pub struct ContentGroup {
pub content_key: [u8; 32],
pub representative: TorrentRecord,
pub aliases: Vec<String>,
pub variant_count: u64,
pub first_seen: u64,
pub last_seen: u64,
pub seen_count: u64,
pub heat: Heat,
pub availability: Availability,
}
pub(crate) struct ContentGroupBuilder {
content_key: [u8; 32],
now: u64,
representative: Option<TorrentRecord>,
aliases: Vec<(RepresentativeRank, String)>,
variant_count: u64,
first_seen: u64,
last_seen: u64,
seen_count: u64,
heat: Heat,
availability: Availability,
availability_initialized: bool,
}
impl ContentGroupBuilder {
pub(crate) fn new(content_key: [u8; 32], now: u64) -> Self {
Self {
content_key,
now,
representative: None,
aliases: Vec::new(),
variant_count: 0,
first_seen: u64::MAX,
last_seen: 0,
seen_count: 0,
heat: Heat::from_score(0),
availability: Availability::default(),
availability_initialized: false,
}
}
pub(crate) fn push(&mut self, record: TorrentRecord) {
if record.content_key != self.content_key {
return;
}
self.variant_count = self.variant_count.saturating_add(1);
self.first_seen = self.first_seen.min(record.first_seen);
self.last_seen = self.last_seen.max(record.last_seen);
self.seen_count = self.seen_count.saturating_add(record.seen_count);
let heat = record.heat(self.now);
let rank = representative_rank(&record, heat);
if heat.score > self.heat.score {
self.heat = heat;
}
if self.availability_initialized {
merge_availability(&mut self.availability, &record.availability);
} else {
self.availability = record.availability.clone();
self.availability_initialized = true;
}
if let Some((existing_rank, _)) = self
.aliases
.iter_mut()
.find(|(_, name)| name == &record.name)
{
*existing_rank = (*existing_rank).max(rank);
} else {
self.aliases.push((rank, record.name.clone()));
}
self.aliases
.sort_unstable_by_key(|item| std::cmp::Reverse(item.0));
self.aliases.truncate(32);
let replace = self.representative.as_ref().is_none_or(|current| {
rank.cmp(&representative_rank(current, current.heat(self.now))) == Ordering::Greater
});
if replace {
self.representative = Some(record);
}
}
pub(crate) fn finish(mut self) -> Option<ContentGroup> {
let representative = self.representative?;
if !self
.aliases
.iter()
.any(|(_, name)| name == &representative.name)
{
if self.aliases.len() == 32 {
self.aliases.pop();
}
self.aliases.push((
representative_rank(&representative, representative.heat(self.now)),
representative.name.clone(),
));
}
Some(ContentGroup {
content_key: self.content_key,
representative,
aliases: self.aliases.into_iter().map(|(_, name)| name).collect(),
variant_count: self.variant_count,
first_seen: self.first_seen,
last_seen: self.last_seen,
seen_count: self.seen_count,
heat: self.heat,
availability: self.availability,
})
}
}
type RepresentativeRank = (u8, u8, u32, u64, u64, std::cmp::Reverse<InfoHash>);
fn representative_rank(record: &TorrentRecord, heat: Heat) -> RepresentativeRank {
let availability = match record.availability.status {
AvailabilityStatus::Active => 2,
AvailabilityStatus::Unknown => 1,
AvailabilityStatus::PossiblyStale => 0,
};
(
availability,
heat.score,
record.availability.reachable_peers,
record.last_seen,
record.seen_count,
std::cmp::Reverse(record.info_hash),
)
}
fn merge_availability(target: &mut Availability, candidate: &Availability) {
if availability_rank(candidate.status) > availability_rank(target.status) {
target.status = candidate.status;
}
target.last_verified_at = target.last_verified_at.max(candidate.last_verified_at);
target.last_success_at = target.last_success_at.max(candidate.last_success_at);
target.discovered_peers = target.discovered_peers.max(candidate.discovered_peers);
target.reachable_peers = target.reachable_peers.max(candidate.reachable_peers);
target.consecutive_failures = target
.consecutive_failures
.min(candidate.consecutive_failures);
target.next_check_at = target.next_check_at.max(candidate.next_check_at);
}
fn availability_rank(status: AvailabilityStatus) -> u8 {
match status {
AvailabilityStatus::Active => 2,
AvailabilityStatus::Unknown => 1,
AvailabilityStatus::PossiblyStale => 0,
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::domain::test_record;
#[test]
fn uses_best_variant_and_aggregates_activity() {
let mut stale = test_record(1, 10);
stale.name = "旧名称".into();
stale.availability.status = AvailabilityStatus::PossiblyStale;
let mut active = test_record(2, 20);
active.name = "流浪地球 S01E03".into();
active.availability.status = AvailabilityStatus::Active;
active.availability.reachable_peers = 2;
active.seen_count = 3;
let mut builder = ContentGroupBuilder::new(stale.content_key, 20);
builder.push(stale);
builder.push(active.clone());
let group = builder.finish().unwrap();
assert_eq!(group.representative.info_hash, active.info_hash);
assert_eq!(group.variant_count, 2);
assert_eq!(group.first_seen, 10);
assert_eq!(group.last_seen, 20);
assert_eq!(group.seen_count, 4);
assert_eq!(group.availability.status, AvailabilityStatus::Active);
}
}
dht-search/src/domain/fingerprint.rs
+1 -1
View File
@@ -2,7 +2,7 @@
use unicode_normalization::UnicodeNormalization;
use super::torrent::{TorrentFile, TorrentRecordError};
use super::{TorrentFile, TorrentRecordError};
const FINGERPRINT_DOMAIN: &[u8] = b"dht-search-content\0";
dht-search/src/domain/info_hash.rs
+51
View File
@@ -0,0 +1,51 @@
// 负责定义二十字节 infohash 的解析显示和稳定二进制表示
use std::{fmt, str::FromStr};
use serde::{Deserialize, Serialize};
use super::TorrentRecordError;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct InfoHash([u8; 20]);
impl InfoHash {
pub const BYTE_LEN: usize = 20;
pub fn from_bytes(bytes: [u8; Self::BYTE_LEN]) -> Self {
Self(bytes)
}
pub fn as_bytes(&self) -> &[u8; Self::BYTE_LEN] {
&self.0
}
}
impl FromStr for InfoHash {
type Err = TorrentRecordError;
fn from_str(value: &str) -> Result<Self, Self::Err> {
let decoded = hex::decode(value).map_err(|_| TorrentRecordError::InvalidInfoHash)?;
let bytes = decoded
.try_into()
.map_err(|_| TorrentRecordError::InvalidInfoHash)?;
Ok(Self(bytes))
}
}
impl fmt::Display for InfoHash {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&hex::encode(self.0))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn round_trips_as_lowercase_hex() {
let hash = InfoHash::from_str("ABABABABABABABABABABABABABABABABABABABAB").unwrap();
assert_eq!(hash.to_string(), "abababababababababababababababababababab");
}
}
dht-search/src/domain/metadata.rs
+394
View File
@@ -0,0 +1,394 @@
// 负责校验外部 Metadata 并生成领域记录和可持久化拒绝原因
use std::str::FromStr;
use dht_crawler::TorrentInfo;
use serde::{Deserialize, Serialize};
use super::{
Availability, AvailabilityStatus, InfoHash, TorrentFile, TorrentRecord, content_key,
torrent::NewTorrentRecord,
};
const MAX_STORED_PEERS: usize = 32;
const METADATA_VALIDATION_VERSION: u64 = 1;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct MetadataLimits {
pub max_files: usize,
pub max_name_bytes: usize,
pub max_path_bytes: usize,
pub max_path_depth: usize,
}
impl MetadataLimits {
pub fn rule_id(self) -> [u8; 32] {
let mut hasher = blake3::Hasher::new();
hasher.update(b"dht-search-metadata-limits\0");
hasher.update(&METADATA_VALIDATION_VERSION.to_be_bytes());
hasher.update(&(self.max_files as u64).to_be_bytes());
hasher.update(&(self.max_name_bytes as u64).to_be_bytes());
hasher.update(&(self.max_path_bytes as u64).to_be_bytes());
hasher.update(&(self.max_path_depth as u64).to_be_bytes());
*hasher.finalize().as_bytes()
}
}
impl Default for MetadataLimits {
fn default() -> Self {
Self {
max_files: 20_000,
max_name_bytes: 1_024,
max_path_bytes: 4_096,
max_path_depth: 64,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum MetadataRejectionReason {
InvalidInfoHash,
EmptyName,
NameTooLong,
InvalidName,
EmptyFileList,
TooManyFiles,
EmptyPath,
PathTooLong,
PathTooDeep,
InvalidPath,
SizeOverflow,
TotalSizeMismatch,
}
impl MetadataRejectionReason {
pub const COUNT: usize = 12;
pub const fn index(self) -> usize {
match self {
Self::InvalidInfoHash => 0,
Self::EmptyName => 1,
Self::NameTooLong => 2,
Self::InvalidName => 3,
Self::EmptyFileList => 4,
Self::TooManyFiles => 5,
Self::EmptyPath => 6,
Self::PathTooLong => 7,
Self::PathTooDeep => 8,
Self::InvalidPath => 9,
Self::SizeOverflow => 10,
Self::TotalSizeMismatch => 11,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RejectedMetadata {
pub info_hash: InfoHash,
pub reason: MetadataRejectionReason,
pub rule_id: [u8; 32],
pub first_rejected_at: u64,
pub last_seen: u64,
pub seen_count: u64,
}
impl RejectedMetadata {
pub fn new(
info_hash: InfoHash,
reason: MetadataRejectionReason,
rule_id: [u8; 32],
timestamp: u64,
) -> Self {
Self {
info_hash,
reason,
rule_id,
first_rejected_at: timestamp,
last_seen: timestamp,
seen_count: 1,
}
}
pub fn observe_again(&mut self, timestamp: u64) {
self.last_seen = self.last_seen.max(timestamp);
self.seen_count = self.seen_count.saturating_add(1);
}
}
#[derive(Debug, thiserror::Error, PartialEq, Eq)]
pub enum TorrentRecordError {
#[error("infohash 必须是二十字节的十六进制字符串")]
InvalidInfoHash,
#[error("种子名称不能为空")]
EmptyName,
#[error("种子名称长度 {actual} 字节超过上限 {limit}")]
NameTooLong { actual: usize, limit: usize },
#[error("种子名称包含控制字符")]
InvalidName,
#[error("文件列表不能为空")]
EmptyFileList,
#[error("文件数量 {actual} 超过上限 {limit}")]
TooManyFiles { actual: usize, limit: usize },
#[error("文件路径长度 {actual} 字节超过上限 {limit}")]
PathTooLong { actual: usize, limit: usize },
#[error("文件路径目录层级 {actual} 超过上限 {limit}")]
PathTooDeep { actual: usize, limit: usize },
#[error("文件路径包含空段上级目录当前目录或控制字符")]
InvalidPath,
#[error("文件总大小溢出")]
SizeOverflow,
#[error("声明大小 {declared} 与文件计算大小 {calculated} 不一致")]
TotalSizeMismatch { declared: u64, calculated: u64 },
#[error("文件路径规范化后为空")]
EmptyNormalizedPath,
}
impl TorrentRecordError {
pub fn rejection_reason(&self) -> MetadataRejectionReason {
match self {
Self::InvalidInfoHash => MetadataRejectionReason::InvalidInfoHash,
Self::EmptyName => MetadataRejectionReason::EmptyName,
Self::NameTooLong { .. } => MetadataRejectionReason::NameTooLong,
Self::InvalidName => MetadataRejectionReason::InvalidName,
Self::EmptyFileList => MetadataRejectionReason::EmptyFileList,
Self::TooManyFiles { .. } => MetadataRejectionReason::TooManyFiles,
Self::EmptyNormalizedPath => MetadataRejectionReason::EmptyPath,
Self::PathTooLong { .. } => MetadataRejectionReason::PathTooLong,
Self::PathTooDeep { .. } => MetadataRejectionReason::PathTooDeep,
Self::InvalidPath => MetadataRejectionReason::InvalidPath,
Self::SizeOverflow => MetadataRejectionReason::SizeOverflow,
Self::TotalSizeMismatch { .. } => MetadataRejectionReason::TotalSizeMismatch,
}
}
}
pub(crate) fn into_record(
info: TorrentInfo,
limits: MetadataLimits,
) -> Result<TorrentRecord, TorrentRecordError> {
let info_hash = InfoHash::from_str(&info.info_hash)?;
validate_name(&info.name, limits)?;
if info.files.is_empty() {
return Err(TorrentRecordError::EmptyFileList);
}
if info.files.len() > limits.max_files {
return Err(TorrentRecordError::TooManyFiles {
actual: info.files.len(),
limit: limits.max_files,
});
}
for file in &info.files {
validate_path(&file.path, limits)?;
}
let files: Vec<_> = info
.files
.into_iter()
.map(|file| TorrentFile {
path: file.path,
size: file.size,
})
.collect();
let calculated_size = files.iter().try_fold(0_u64, |total, file| {
total
.checked_add(file.size)
.ok_or(TorrentRecordError::SizeOverflow)
})?;
if calculated_size != info.total_size {
return Err(TorrentRecordError::TotalSizeMismatch {
declared: info.total_size,
calculated: calculated_size,
});
}
let content_key = content_key(&files)?;
let mut source_peers = info.peers;
source_peers.sort_unstable();
source_peers.dedup();
source_peers.truncate(MAX_STORED_PEERS);
let reachable_peers = source_peers.len().min(u32::MAX as usize) as u32;
let availability = if reachable_peers > 0 {
Availability {
status: AvailabilityStatus::Active,
last_verified_at: Some(info.timestamp),
last_success_at: Some(info.timestamp),
discovered_peers: reachable_peers,
reachable_peers,
consecutive_failures: 0,
next_check_at: info.timestamp.saturating_add(86_400),
}
} else {
Availability::default()
};
Ok(TorrentRecord::from_new(NewTorrentRecord {
info_hash,
name: info.name,
total_size: info.total_size,
files,
piece_length: info.piece_length,
source_peers,
content_key,
timestamp: info.timestamp,
availability,
}))
}
fn validate_name(name: &str, limits: MetadataLimits) -> Result<(), TorrentRecordError> {
if name.trim().is_empty() {
return Err(TorrentRecordError::EmptyName);
}
if name.len() > limits.max_name_bytes {
return Err(TorrentRecordError::NameTooLong {
actual: name.len(),
limit: limits.max_name_bytes,
});
}
if name.chars().any(char::is_control) {
return Err(TorrentRecordError::InvalidName);
}
Ok(())
}
fn validate_path(path: &str, limits: MetadataLimits) -> Result<(), TorrentRecordError> {
if path.is_empty() {
return Err(TorrentRecordError::EmptyNormalizedPath);
}
if path.len() > limits.max_path_bytes {
return Err(TorrentRecordError::PathTooLong {
actual: path.len(),
limit: limits.max_path_bytes,
});
}
if path.chars().any(char::is_control) {
return Err(TorrentRecordError::InvalidPath);
}
let mut depth = 0_usize;
for component in path.split(['/', '\\']) {
if component.is_empty() || matches!(component, "." | "..") {
return Err(TorrentRecordError::InvalidPath);
}
depth += 1;
}
if depth > limits.max_path_depth {
return Err(TorrentRecordError::PathTooDeep {
actual: depth,
limit: limits.max_path_depth,
});
}
Ok(())
}
impl TryFrom<TorrentInfo> for TorrentRecord {
type Error = TorrentRecordError;
fn try_from(info: TorrentInfo) -> Result<Self, Self::Error> {
into_record(info, MetadataLimits::default())
}
}
#[cfg(test)]
mod tests {
use dht_crawler::FileInfo;
use super::*;
fn torrent_info(files: Vec<FileInfo>) -> TorrentInfo {
TorrentInfo {
info_hash: "0101010101010101010101010101010101010101".into(),
magnet_link: String::new(),
name: "Example".into(),
total_size: files
.iter()
.fold(0_u64, |total, file| total.saturating_add(file.size)),
files,
piece_length: 16_384,
peers: Vec::new(),
timestamp: 1,
}
}
#[test]
fn limits_accept_boundary_and_reject_excess() {
let limits = MetadataLimits {
max_files: 1,
max_name_bytes: 7,
max_path_bytes: 8,
max_path_depth: 2,
};
let accepted = torrent_info(vec![FileInfo {
path: "dir/a.rs".into(),
size: 1,
}]);
assert!(into_record(accepted, limits).is_ok());
let too_many = torrent_info(vec![
FileInfo {
path: "a".into(),
size: 1,
},
FileInfo {
path: "b".into(),
size: 1,
},
]);
assert_eq!(
into_record(too_many, limits)
.unwrap_err()
.rejection_reason(),
MetadataRejectionReason::TooManyFiles
);
}
#[test]
fn unsafe_and_deep_paths_are_rejected_by_reason() {
let limits = MetadataLimits {
max_path_depth: 2,
..MetadataLimits::default()
};
for (path, reason) in [
("dir/../file", MetadataRejectionReason::InvalidPath),
("dir//file", MetadataRejectionReason::InvalidPath),
("a/b/c", MetadataRejectionReason::PathTooDeep),
("a\0b", MetadataRejectionReason::InvalidPath),
] {
let info = torrent_info(vec![FileInfo {
path: path.into(),
size: 1,
}]);
assert_eq!(
into_record(info, limits).unwrap_err().rejection_reason(),
reason
);
}
}
#[test]
fn size_overflow_is_rejected_without_panicking() {
let mut info = torrent_info(vec![
FileInfo {
path: "a".into(),
size: u64::MAX,
},
FileInfo {
path: "b".into(),
size: 1,
},
]);
info.total_size = u64::MAX;
assert_eq!(
into_record(info, MetadataLimits::default())
.unwrap_err()
.rejection_reason(),
MetadataRejectionReason::SizeOverflow
);
}
#[test]
fn rule_id_changes_when_a_limit_changes() {
let defaults = MetadataLimits::default();
let changed = MetadataLimits {
max_files: defaults.max_files - 1,
..defaults
};
assert_ne!(defaults.rule_id(), changed.rule_id());
}
}
dht-search/src/domain/mod.rs
+8 -3
View File
@@ -1,14 +1,19 @@
// 负责导出不依赖存储搜索和传输实现的核心领域模型
mod content_group;
mod fingerprint;
mod info_hash;
mod metadata;
mod torrent;
pub use content_group::ContentGroup;
pub(crate) use content_group::ContentGroupBuilder;
pub use fingerprint::content_key;
pub(crate) use torrent::ContentGroupBuilder;
pub use info_hash::InfoHash;
pub use metadata::{MetadataLimits, MetadataRejectionReason, RejectedMetadata, TorrentRecordError};
#[cfg(test)]
pub(crate) use torrent::test_record;
pub use torrent::{
Availability, AvailabilityStatus, ContentGroup, Heat, HeatLevel, InfoHash, MetadataLimits,
MetadataRejectionReason, RejectedMetadata, TorrentFile, TorrentRecord, TorrentRecordError,
Availability, AvailabilityStatus, Heat, HeatLevel, TorrentFile, TorrentRecord,
VerificationResult,
};
dht-search/src/domain/torrent.rs
+40 -591
View File
@@ -1,151 +1,15 @@
// 负责定义种子元数据文件条目发现状态和索引状态模型
// 负责定义种子记录活跃度可用性和重复发现时的状态演进
use dht_crawler::TorrentInfo;
use serde::{Deserialize, Serialize};
use std::{cmp::Ordering, fmt, str::FromStr};
use super::fingerprint::content_key;
#[cfg(test)]
use super::content_key;
use super::{InfoHash, MetadataLimits, TorrentRecordError};
const MAX_STORED_PEERS: usize = 32;
const ACTIVITY_SCALE: u64 = 1_000;
const ACTIVITY_HALF_LIFE_SECS: f64 = 86_400.0;
const METADATA_VALIDATION_VERSION: u64 = 1;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct MetadataLimits {
pub max_files: usize,
pub max_name_bytes: usize,
pub max_path_bytes: usize,
pub max_path_depth: usize,
}
impl MetadataLimits {
pub fn rule_id(self) -> [u8; 32] {
let mut hasher = blake3::Hasher::new();
hasher.update(b"dht-search-metadata-limits\0");
hasher.update(&METADATA_VALIDATION_VERSION.to_be_bytes());
hasher.update(&(self.max_files as u64).to_be_bytes());
hasher.update(&(self.max_name_bytes as u64).to_be_bytes());
hasher.update(&(self.max_path_bytes as u64).to_be_bytes());
hasher.update(&(self.max_path_depth as u64).to_be_bytes());
*hasher.finalize().as_bytes()
}
}
impl Default for MetadataLimits {
fn default() -> Self {
Self {
max_files: 20_000,
max_name_bytes: 1_024,
max_path_bytes: 4_096,
max_path_depth: 64,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum MetadataRejectionReason {
InvalidInfoHash,
EmptyName,
NameTooLong,
InvalidName,
EmptyFileList,
TooManyFiles,
EmptyPath,
PathTooLong,
PathTooDeep,
InvalidPath,
SizeOverflow,
TotalSizeMismatch,
}
impl MetadataRejectionReason {
pub const COUNT: usize = 12;
pub const fn index(self) -> usize {
match self {
Self::InvalidInfoHash => 0,
Self::EmptyName => 1,
Self::NameTooLong => 2,
Self::InvalidName => 3,
Self::EmptyFileList => 4,
Self::TooManyFiles => 5,
Self::EmptyPath => 6,
Self::PathTooLong => 7,
Self::PathTooDeep => 8,
Self::InvalidPath => 9,
Self::SizeOverflow => 10,
Self::TotalSizeMismatch => 11,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RejectedMetadata {
pub info_hash: InfoHash,
pub reason: MetadataRejectionReason,
pub rule_id: [u8; 32],
pub first_rejected_at: u64,
pub last_seen: u64,
pub seen_count: u64,
}
impl RejectedMetadata {
pub fn new(
info_hash: InfoHash,
reason: MetadataRejectionReason,
rule_id: [u8; 32],
timestamp: u64,
) -> Self {
Self {
info_hash,
reason,
rule_id,
first_rejected_at: timestamp,
last_seen: timestamp,
seen_count: 1,
}
}
pub fn observe_again(&mut self, timestamp: u64) {
self.last_seen = self.last_seen.max(timestamp);
self.seen_count = self.seen_count.saturating_add(1);
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct InfoHash([u8; 20]);
impl InfoHash {
pub const BYTE_LEN: usize = 20;
pub fn from_bytes(bytes: [u8; Self::BYTE_LEN]) -> Self {
Self(bytes)
}
pub fn as_bytes(&self) -> &[u8; Self::BYTE_LEN] {
&self.0
}
}
impl FromStr for InfoHash {
type Err = TorrentRecordError;
fn from_str(value: &str) -> Result<Self, Self::Err> {
let decoded = hex::decode(value).map_err(|_| TorrentRecordError::InvalidInfoHash)?;
let bytes = decoded
.try_into()
.map_err(|_| TorrentRecordError::InvalidInfoHash)?;
Ok(Self(bytes))
}
}
impl fmt::Display for InfoHash {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&hex::encode(self.0))
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct TorrentFile {
@@ -195,160 +59,6 @@ pub struct Heat {
pub level: HeatLevel,
}
#[derive(Debug, Clone, PartialEq)]
pub struct ContentGroup {
pub content_key: [u8; 32],
pub representative: TorrentRecord,
pub aliases: Vec<String>,
pub variant_count: u64,
pub first_seen: u64,
pub last_seen: u64,
pub seen_count: u64,
pub heat: Heat,
pub availability: Availability,
}
pub(crate) struct ContentGroupBuilder {
content_key: [u8; 32],
now: u64,
representative: Option<TorrentRecord>,
aliases: Vec<(RepresentativeRank, String)>,
variant_count: u64,
first_seen: u64,
last_seen: u64,
seen_count: u64,
heat: Heat,
availability: Availability,
availability_initialized: bool,
}
impl ContentGroupBuilder {
pub(crate) fn new(content_key: [u8; 32], now: u64) -> Self {
Self {
content_key,
now,
representative: None,
aliases: Vec::new(),
variant_count: 0,
first_seen: u64::MAX,
last_seen: 0,
seen_count: 0,
heat: Heat::from_score(0),
availability: Availability::default(),
availability_initialized: false,
}
}
pub(crate) fn push(&mut self, record: TorrentRecord) {
if record.content_key != self.content_key {
return;
}
self.variant_count = self.variant_count.saturating_add(1);
self.first_seen = self.first_seen.min(record.first_seen);
self.last_seen = self.last_seen.max(record.last_seen);
self.seen_count = self.seen_count.saturating_add(record.seen_count);
let heat = record.heat(self.now);
let rank = representative_rank(&record, heat);
if heat.score > self.heat.score {
self.heat = heat;
}
if self.availability_initialized {
merge_availability(&mut self.availability, &record.availability);
} else {
self.availability = record.availability.clone();
self.availability_initialized = true;
}
if let Some((existing_rank, _)) = self
.aliases
.iter_mut()
.find(|(_, name)| name == &record.name)
{
*existing_rank = (*existing_rank).max(rank);
} else {
self.aliases.push((rank, record.name.clone()));
}
self.aliases
.sort_unstable_by_key(|item| std::cmp::Reverse(item.0));
self.aliases.truncate(32);
let replace = self.representative.as_ref().is_none_or(|current| {
rank.cmp(&representative_rank(current, current.heat(self.now))) == Ordering::Greater
});
if replace {
self.representative = Some(record);
}
}
pub(crate) fn finish(mut self) -> Option<ContentGroup> {
let representative = self.representative?;
if !self
.aliases
.iter()
.any(|(_, name)| name == &representative.name)
{
if self.aliases.len() == 32 {
self.aliases.pop();
}
self.aliases.push((
representative_rank(&representative, representative.heat(self.now)),
representative.name.clone(),
));
}
Some(ContentGroup {
content_key: self.content_key,
representative,
aliases: self.aliases.into_iter().map(|(_, name)| name).collect(),
variant_count: self.variant_count,
first_seen: self.first_seen,
last_seen: self.last_seen,
seen_count: self.seen_count,
heat: self.heat,
availability: self.availability,
})
}
}
type RepresentativeRank = (u8, u8, u32, u64, u64, std::cmp::Reverse<InfoHash>);
fn representative_rank(record: &TorrentRecord, heat: Heat) -> RepresentativeRank {
let availability = match record.availability.status {
AvailabilityStatus::Active => 2,
AvailabilityStatus::Unknown => 1,
AvailabilityStatus::PossiblyStale => 0,
};
(
availability,
heat.score,
record.availability.reachable_peers,
record.last_seen,
record.seen_count,
std::cmp::Reverse(record.info_hash),
)
}
fn merge_availability(target: &mut Availability, candidate: &Availability) {
let target_rank = availability_rank(target.status);
let candidate_rank = availability_rank(candidate.status);
if candidate_rank > target_rank {
target.status = candidate.status;
}
target.last_verified_at = target.last_verified_at.max(candidate.last_verified_at);
target.last_success_at = target.last_success_at.max(candidate.last_success_at);
target.discovered_peers = target.discovered_peers.max(candidate.discovered_peers);
target.reachable_peers = target.reachable_peers.max(candidate.reachable_peers);
target.consecutive_failures = target
.consecutive_failures
.min(candidate.consecutive_failures);
target.next_check_at = target.next_check_at.max(candidate.next_check_at);
}
fn availability_rank(status: AvailabilityStatus) -> u8 {
match status {
AvailabilityStatus::Active => 2,
AvailabilityStatus::Unknown => 1,
AvailabilityStatus::PossiblyStale => 0,
}
}
impl Heat {
pub fn from_score(score: u8) -> Self {
let level = match score {
@@ -381,90 +91,42 @@ pub struct TorrentRecord {
pub activity_updated_at: u64,
}
pub(crate) struct NewTorrentRecord {
pub(crate) info_hash: InfoHash,
pub(crate) name: String,
pub(crate) total_size: u64,
pub(crate) files: Vec<TorrentFile>,
pub(crate) piece_length: u64,
pub(crate) source_peers: Vec<String>,
pub(crate) content_key: [u8; 32],
pub(crate) timestamp: u64,
pub(crate) availability: Availability,
}
impl TorrentRecord {
pub(crate) fn from_new(new: NewTorrentRecord) -> Self {
Self {
info_hash: new.info_hash,
name: new.name,
total_size: new.total_size,
files: new.files,
piece_length: new.piece_length,
source_peers: new.source_peers,
content_key: new.content_key,
first_seen: new.timestamp,
last_seen: new.timestamp,
seen_count: 1,
availability: new.availability,
activity_score_millis: ACTIVITY_SCALE,
activity_updated_at: new.timestamp,
}
}
pub fn try_from_with_limits(
info: TorrentInfo,
limits: MetadataLimits,
) -> Result<Self, TorrentRecordError> {
let info_hash = InfoHash::from_str(&info.info_hash)?;
if info.name.trim().is_empty() {
return Err(TorrentRecordError::EmptyName);
}
if info.name.len() > limits.max_name_bytes {
return Err(TorrentRecordError::NameTooLong {
actual: info.name.len(),
limit: limits.max_name_bytes,
});
}
if info.name.chars().any(char::is_control) {
return Err(TorrentRecordError::InvalidName);
}
if info.files.is_empty() {
return Err(TorrentRecordError::EmptyFileList);
}
if info.files.len() > limits.max_files {
return Err(TorrentRecordError::TooManyFiles {
actual: info.files.len(),
limit: limits.max_files,
});
}
for file in &info.files {
validate_path(&file.path, limits)?;
}
let files: Vec<_> = info
.files
.into_iter()
.map(|file| TorrentFile {
path: file.path,
size: file.size,
})
.collect();
let calculated_size = files.iter().try_fold(0_u64, |total, file| {
total
.checked_add(file.size)
.ok_or(TorrentRecordError::SizeOverflow)
})?;
if calculated_size != info.total_size {
return Err(TorrentRecordError::TotalSizeMismatch {
declared: info.total_size,
calculated: calculated_size,
});
}
let content_key = content_key(&files)?;
let mut source_peers = info.peers;
source_peers.sort_unstable();
source_peers.dedup();
source_peers.truncate(MAX_STORED_PEERS);
let reachable_peers = source_peers.len().min(u32::MAX as usize) as u32;
let availability = if reachable_peers > 0 {
Availability {
status: AvailabilityStatus::Active,
last_verified_at: Some(info.timestamp),
last_success_at: Some(info.timestamp),
discovered_peers: reachable_peers,
reachable_peers,
consecutive_failures: 0,
next_check_at: info.timestamp.saturating_add(86_400),
}
} else {
Availability::default()
};
Ok(Self {
info_hash,
name: info.name,
total_size: info.total_size,
files,
piece_length: info.piece_length,
source_peers,
content_key,
first_seen: info.timestamp,
last_seen: info.timestamp,
seen_count: 1,
availability,
activity_score_millis: ACTIVITY_SCALE,
activity_updated_at: info.timestamp,
})
super::metadata::into_record(info, limits)
}
pub fn observe_again(&mut self, timestamp: u64, peers: &[String]) {
@@ -554,90 +216,6 @@ fn failure_retry_secs(failures: u32) -> u64 {
}
}
impl TryFrom<TorrentInfo> for TorrentRecord {
type Error = TorrentRecordError;
fn try_from(info: TorrentInfo) -> Result<Self, Self::Error> {
Self::try_from_with_limits(info, MetadataLimits::default())
}
}
fn validate_path(path: &str, limits: MetadataLimits) -> Result<(), TorrentRecordError> {
if path.is_empty() {
return Err(TorrentRecordError::EmptyNormalizedPath);
}
if path.len() > limits.max_path_bytes {
return Err(TorrentRecordError::PathTooLong {
actual: path.len(),
limit: limits.max_path_bytes,
});
}
if path.chars().any(char::is_control) {
return Err(TorrentRecordError::InvalidPath);
}
let mut depth = 0_usize;
for component in path.split(['/', '\\']) {
if component.is_empty() || matches!(component, "." | "..") {
return Err(TorrentRecordError::InvalidPath);
}
depth += 1;
}
if depth > limits.max_path_depth {
return Err(TorrentRecordError::PathTooDeep {
actual: depth,
limit: limits.max_path_depth,
});
}
Ok(())
}
#[derive(Debug, thiserror::Error, PartialEq, Eq)]
pub enum TorrentRecordError {
#[error("infohash 必须是二十字节的十六进制字符串")]
InvalidInfoHash,
#[error("种子名称不能为空")]
EmptyName,
#[error("种子名称长度 {actual} 字节超过上限 {limit}")]
NameTooLong { actual: usize, limit: usize },
#[error("种子名称包含控制字符")]
InvalidName,
#[error("文件列表不能为空")]
EmptyFileList,
#[error("文件数量 {actual} 超过上限 {limit}")]
TooManyFiles { actual: usize, limit: usize },
#[error("文件路径长度 {actual} 字节超过上限 {limit}")]
PathTooLong { actual: usize, limit: usize },
#[error("文件路径目录层级 {actual} 超过上限 {limit}")]
PathTooDeep { actual: usize, limit: usize },
#[error("文件路径包含空段上级目录当前目录或控制字符")]
InvalidPath,
#[error("文件总大小溢出")]
SizeOverflow,
#[error("声明大小 {declared} 与文件计算大小 {calculated} 不一致")]
TotalSizeMismatch { declared: u64, calculated: u64 },
#[error("文件路径规范化后为空")]
EmptyNormalizedPath,
}
impl TorrentRecordError {
pub fn rejection_reason(&self) -> MetadataRejectionReason {
match self {
Self::InvalidInfoHash => MetadataRejectionReason::InvalidInfoHash,
Self::EmptyName => MetadataRejectionReason::EmptyName,
Self::NameTooLong { .. } => MetadataRejectionReason::NameTooLong,
Self::InvalidName => MetadataRejectionReason::InvalidName,
Self::EmptyFileList => MetadataRejectionReason::EmptyFileList,
Self::TooManyFiles { .. } => MetadataRejectionReason::TooManyFiles,
Self::EmptyNormalizedPath => MetadataRejectionReason::EmptyPath,
Self::PathTooLong { .. } => MetadataRejectionReason::PathTooLong,
Self::PathTooDeep { .. } => MetadataRejectionReason::PathTooDeep,
Self::InvalidPath => MetadataRejectionReason::InvalidPath,
Self::SizeOverflow => MetadataRejectionReason::SizeOverflow,
Self::TotalSizeMismatch { .. } => MetadataRejectionReason::TotalSizeMismatch,
}
}
}
#[cfg(test)]
pub(crate) fn test_record(hash_byte: u8, timestamp: u64) -> TorrentRecord {
let files = vec![TorrentFile {
@@ -663,112 +241,9 @@ pub(crate) fn test_record(hash_byte: u8, timestamp: u64) -> TorrentRecord {
#[cfg(test)]
mod tests {
use super::*;
use dht_crawler::FileInfo;
fn torrent_info(files: Vec<FileInfo>) -> TorrentInfo {
TorrentInfo {
info_hash: "0101010101010101010101010101010101010101".into(),
magnet_link: String::new(),
name: "Example".into(),
total_size: files
.iter()
.fold(0_u64, |total, file| total.saturating_add(file.size)),
files,
piece_length: 16_384,
peers: Vec::new(),
timestamp: 1,
}
}
#[test]
fn infohash_round_trips_as_lowercase_hex() {
let hash = InfoHash::from_str("ABABABABABABABABABABABABABABABABABABABAB").unwrap();
assert_eq!(hash.to_string(), "abababababababababababababababababababab");
}
#[test]
fn metadata_limits_accept_boundary_and_reject_excess() {
let limits = MetadataLimits {
max_files: 1,
max_name_bytes: 7,
max_path_bytes: 8,
max_path_depth: 2,
};
let accepted = torrent_info(vec![FileInfo {
path: "dir/a.rs".into(),
size: 1,
}]);
assert!(TorrentRecord::try_from_with_limits(accepted, limits).is_ok());
let too_many = torrent_info(vec![
FileInfo {
path: "a".into(),
size: 1,
},
FileInfo {
path: "b".into(),
size: 1,
},
]);
let error = TorrentRecord::try_from_with_limits(too_many, limits).unwrap_err();
assert_eq!(
error.rejection_reason(),
MetadataRejectionReason::TooManyFiles
);
}
#[test]
fn unsafe_and_deep_paths_are_rejected_by_reason() {
let limits = MetadataLimits {
max_path_depth: 2,
..MetadataLimits::default()
};
for (path, reason) in [
("dir/../file", MetadataRejectionReason::InvalidPath),
("dir//file", MetadataRejectionReason::InvalidPath),
("a/b/c", MetadataRejectionReason::PathTooDeep),
("a\0b", MetadataRejectionReason::InvalidPath),
] {
let info = torrent_info(vec![FileInfo {
path: path.into(),
size: 1,
}]);
let error = TorrentRecord::try_from_with_limits(info, limits).unwrap_err();
assert_eq!(error.rejection_reason(), reason, "path={path:?}");
}
}
#[test]
fn size_overflow_is_rejected_without_panicking() {
let mut info = torrent_info(vec![
FileInfo {
path: "a".into(),
size: u64::MAX,
},
FileInfo {
path: "b".into(),
size: 1,
},
]);
info.total_size = u64::MAX;
let error =
TorrentRecord::try_from_with_limits(info, MetadataLimits::default()).unwrap_err();
assert_eq!(
error.rejection_reason(),
MetadataRejectionReason::SizeOverflow
);
}
#[test]
fn rule_id_changes_when_a_limit_changes() {
let defaults = MetadataLimits::default();
let changed = MetadataLimits {
max_files: defaults.max_files - 1,
..defaults
};
assert_ne!(defaults.rule_id(), changed.rule_id());
}
use super::*;
#[test]
fn repeated_observation_updates_time_count_and_unique_peers() {
@@ -783,10 +258,10 @@ mod tests {
#[test]
fn activity_decays_and_recent_observation_increases_heat() {
let mut record = test_record(1, 10);
let old_heat = record.heat(10 + 30 * 86_400).score;
record.observe_again(10 + 30 * 86_400, &[]);
let new_heat = record.heat(10 + 30 * 86_400).score;
assert!(new_heat > old_heat);
let now = 10 + 30 * 86_400;
let old_heat = record.heat(now).score;
record.observe_again(now, &[]);
assert!(record.heat(now).score > old_heat);
}
#[test]
@@ -799,7 +274,6 @@ mod tests {
});
assert_eq!(record.availability.status, AvailabilityStatus::Active);
assert_eq!(record.availability.next_check_at, 86_500);
record.apply_verification(VerificationResult {
verified_at: 200,
discovered_peers: 0,
@@ -809,7 +283,6 @@ mod tests {
record.availability.status,
AvailabilityStatus::PossiblyStale
);
assert_eq!(record.availability.consecutive_failures, 1);
assert_eq!(record.availability.next_check_at, 3_800);
}
@@ -834,28 +307,4 @@ mod tests {
assert_eq!(record.availability.last_verified_at, Some(100));
assert_eq!(record.availability.next_check_at, 86_500);
}
#[test]
fn content_group_uses_best_variant_and_aggregates_activity() {
let mut stale = test_record(1, 10);
stale.name = "旧名称".into();
stale.availability.status = AvailabilityStatus::PossiblyStale;
let mut active = test_record(2, 20);
active.name = "流浪地球 S01E03".into();
active.availability.status = AvailabilityStatus::Active;
active.availability.reachable_peers = 2;
active.seen_count = 3;
let mut builder = ContentGroupBuilder::new(stale.content_key, 20);
builder.push(stale);
builder.push(active.clone());
let group = builder.finish().unwrap();
assert_eq!(group.representative.info_hash, active.info_hash);
assert_eq!(group.variant_count, 2);
assert_eq!(group.first_seen, 10);
assert_eq!(group.last_seen, 20);
assert_eq!(group.seen_count, 4);
assert_eq!(group.availability.status, AvailabilityStatus::Active);
}
}
dht-search/src/index_worker.rs
+121
View File
@@ -0,0 +1,121 @@
// 负责批量提交搜索索引并处理临时错误和关闭前排空
use std::{sync::Arc, time::Duration};
use dht_search::{search::SearchEngine, storage::RocksTorrentRepository};
use tokio_util::sync::CancellationToken;
pub(crate) async fn run(
repository: Arc<RocksTorrentRepository>,
search: SearchEngine,
batch_size: usize,
interval: Duration,
cancel: CancellationToken,
fatal: tokio::sync::oneshot::Sender<String>,
) -> Result<(), String> {
let mut ticker = tokio::time::interval(interval);
ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
let mut consecutive_retries = 0_u32;
loop {
tokio::select! {
_ = cancel.cancelled() => break,
_ = ticker.tick() => {
match index_one_batch(repository.clone(), search.clone(), batch_size).await {
Ok(count) => {
consecutive_retries = 0;
if count > 0 {
tracing::debug!(count, "搜索索引已提交");
}
}
Err(IndexBatchError::Retryable(error)) => {
consecutive_retries = consecutive_retries.saturating_add(1);
let delay = retry_delay(consecutive_retries);
tracing::warn!(%error, retry = consecutive_retries, delay_ms = delay.as_millis(), "搜索索引遇到临时 I/O 错误");
tokio::select! {
_ = cancel.cancelled() => break,
_ = tokio::time::sleep(delay) => {}
}
}
Err(IndexBatchError::Fatal(error)) => {
let _ = fatal.send(error.clone());
return Err(error);
}
}
}
}
}
drain_before_shutdown(repository, search, batch_size).await
}
async fn drain_before_shutdown(
repository: Arc<RocksTorrentRepository>,
search: SearchEngine,
batch_size: usize,
) -> Result<(), String> {
let mut retries = 0_u32;
loop {
match index_one_batch(repository.clone(), search.clone(), batch_size).await {
Ok(0) => return Ok(()),
Ok(_) => retries = 0,
Err(IndexBatchError::Retryable(error)) if retries < 3 => {
retries += 1;
let delay = retry_delay(retries);
tracing::warn!(%error, retry = retries, delay_ms = delay.as_millis(), "关闭前提交搜索索引时遇到临时 I/O 错误");
tokio::time::sleep(delay).await;
}
Err(IndexBatchError::Retryable(error)) => {
tracing::warn!(%error, "关闭前搜索索引仍被占用 待索引状态将在下次启动恢复");
return Ok(());
}
Err(IndexBatchError::Fatal(error)) => return Err(error),
}
}
}
enum IndexBatchError {
Retryable(String),
Fatal(String),
}
async fn index_one_batch(
repository: Arc<RocksTorrentRepository>,
search: SearchEngine,
batch_size: usize,
) -> Result<usize, IndexBatchError> {
tokio::task::spawn_blocking(move || {
match search.index_pending(repository.as_ref(), batch_size, unix_timestamp()) {
Ok(count) => Ok(count),
Err(error) if error.is_retryable_io() => {
Err(IndexBatchError::Retryable(error.to_string()))
}
Err(error) => Err(IndexBatchError::Fatal(error.to_string())),
}
})
.await
.map_err(|error| IndexBatchError::Fatal(error.to_string()))?
}
fn retry_delay(attempt: u32) -> Duration {
let shift = attempt.saturating_sub(1).min(6);
Duration::from_millis((250_u64 << shift).min(10_000))
}
fn unix_timestamp() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn retry_delay_is_exponential_and_bounded() {
assert_eq!(retry_delay(1), Duration::from_millis(250));
assert_eq!(retry_delay(3), Duration::from_secs(1));
assert_eq!(retry_delay(100), Duration::from_secs(10));
}
}
dht-search/src/main.rs
+2
View File
@@ -7,6 +7,8 @@ mod app;
mod config;
mod crawler;
mod error;
mod index_worker;
mod monitor;
mod shutdown;
mod telemetry;
mod verification;
dht-search/src/monitor.rs
+75
View File
@@ -0,0 +1,75 @@
// 负责定期汇总采集持久化和过滤指标并输出结构化运行状态
use std::time::Duration;
use dht_crawler::DHTServer;
use dht_search::domain::MetadataRejectionReason;
use tokio_util::sync::CancellationToken;
use crate::crawler::pipeline::PersistenceIngress;
pub(crate) async fn run(
server: DHTServer,
ingress: PersistenceIngress,
interval_secs: u64,
cancel: CancellationToken,
) {
let mut interval = tokio::time::interval(Duration::from_secs(interval_secs));
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
let mut previous_udp_tx = 0;
let mut previous_metadata_attempts = 0;
loop {
tokio::select! {
_ = cancel.cancelled() => break,
_ = interval.tick() => {
let dht = server.runtime_stats().snapshot();
let observability = server.runtime_stats().observability_snapshot();
let storage = ingress.snapshot();
let udp_tx_per_second = observability
.udp_tx_packets
.saturating_sub(previous_udp_tx)
/ interval_secs;
let metadata_connects_per_second = dht
.metadata_peer_attempts
.saturating_sub(previous_metadata_attempts)
/ interval_secs;
previous_udp_tx = observability.udp_tx_packets;
previous_metadata_attempts = dht.metadata_peer_attempts;
tracing::info!(
nodes = dht.node_pool_size,
udp_tx = observability.udp_tx_packets,
udp_tx_per_second,
find_node_queries = dht.queries_new + dht.queries_revisit + dht.queries_bootstrap,
peer_lookup_queries = dht.peer_lookup_queries,
peer_lookup_preferred_succeeded = dht.peer_lookup_preferred_succeeded,
peer_lookup_fallbacks = dht.peer_lookup_fallbacks,
sample_queries = dht.sample_infohashes_queries,
sampled_hashes = dht.sample_infohashes_hashes_discovered,
sampled_hashes_filtered = dht.sample_infohashes_hashes_filtered,
peers = dht.peer_lookup_peers_found,
metadata_connects_per_second,
metadata_in_flight = dht.metadata_in_flight,
metadata_ok = dht.metadata_peer_succeeded,
metadata_failed = dht.metadata_peer_failed,
metadata_filtered = observability
.metadata_failure_size_limit
.saturating_add(storage.filtered.total()),
metadata_filtered_too_many_files = storage
.filtered
.count(MetadataRejectionReason::TooManyFiles),
metadata_filtered_invalid_path = storage
.filtered
.count(MetadataRejectionReason::InvalidPath),
persistence_accepted = storage.accepted,
persistence_inserted = storage.inserted,
persistence_updated = storage.updated,
persistence_rejected_full = storage.rejected_full,
persistence_invalid = storage.invalid,
persistence_failed = storage.failed,
persistence_queue = storage.queue_depth,
"运行状态"
)
}
}
}
}
dht-search/src/search/document.rs
+158
View File
@@ -0,0 +1,158 @@
// 负责在领域内容组和 Tantivy 文档之间执行双向字段映射
use std::{collections::BTreeSet, path::Path};
use tantivy::{TantivyDocument, schema::Value};
use unicode_normalization::UnicodeNormalization;
use crate::domain::{AvailabilityStatus, ContentGroup, Heat, TorrentRecord};
use super::{
SearchError,
query::{AvailabilitySummary, SearchHit},
schema::SearchFields,
};
pub(crate) fn from_group(group: &ContentGroup, fields: SearchFields) -> TantivyDocument {
const MAX_INDEXED_FILES: usize = 512;
const MAX_PATH_TEXT_BYTES: usize = 32 * 1024;
let record = &group.representative;
let mut document = TantivyDocument::default();
document.add_text(fields.info_hash, record.info_hash.to_string());
document.add_text(fields.name, normalize_bounded(&record.name, 512));
document.add_text(fields.regex_text, normalize_bounded(&record.name, 512));
document.add_text(fields.display_name, &record.name);
for alias in &group.aliases {
let alias = normalize_bounded(alias, 512);
document.add_text(fields.aliases, &alias);
document.add_text(fields.regex_text, alias);
}
let mut indexed_path_bytes = 0_usize;
for file in record.files.iter().take(MAX_INDEXED_FILES) {
let path = normalize_bounded(&file.path, 512);
if indexed_path_bytes.saturating_add(path.len()) > MAX_PATH_TEXT_BYTES {
break;
}
indexed_path_bytes += path.len();
document.add_text(fields.files_text, &path);
document.add_text(fields.regex_text, path);
}
for extension in extensions(record) {
document.add_text(fields.extensions, extension);
}
document.add_u64(fields.total_size, record.total_size);
document.add_u64(fields.file_count, record.files.len() as u64);
document.add_u64(fields.first_seen, group.first_seen);
document.add_u64(fields.last_seen, group.last_seen);
document.add_u64(fields.seen_count, group.seen_count);
document.add_text(fields.content_key, hex::encode(group.content_key));
document.add_u64(
fields.availability_status,
availability_number(group.availability.status),
);
document.add_u64(
fields.reachable_peers,
u64::from(group.availability.reachable_peers),
);
document.add_u64(
fields.last_verified_at,
group.availability.last_verified_at.unwrap_or(0),
);
document.add_u64(fields.heat_score, u64::from(group.heat.score));
document.add_u64(fields.variant_count, group.variant_count);
document
}
pub(crate) fn to_hit(
document: &TantivyDocument,
fields: SearchFields,
score: f32,
) -> Result<SearchHit, SearchError> {
Ok(SearchHit {
info_hash: text(document, fields.info_hash, "info_hash")?,
name: text(document, fields.display_name, "display_name")?,
total_size: number(document, fields.total_size, "total_size")?,
file_count: number(document, fields.file_count, "file_count")?,
first_seen: number(document, fields.first_seen, "first_seen")?,
last_seen: number(document, fields.last_seen, "last_seen")?,
seen_count: number(document, fields.seen_count, "seen_count")?,
content_key: text(document, fields.content_key, "content_key")?,
variant_count: number(document, fields.variant_count, "variant_count")?,
score,
heat: Heat::from_score(number(document, fields.heat_score, "heat_score")?.min(100) as u8),
availability: AvailabilitySummary {
status: availability_status(number(
document,
fields.availability_status,
"availability_status",
)?),
last_verified_at: match number(document, fields.last_verified_at, "last_verified_at")? {
0 => None,
value => Some(value),
},
reachable_peers: number(document, fields.reachable_peers, "reachable_peers")?
.min(u64::from(u32::MAX)) as u32,
},
})
}
pub(crate) fn availability_number(status: AvailabilityStatus) -> u64 {
match status {
AvailabilityStatus::Unknown => 0,
AvailabilityStatus::Active => 1,
AvailabilityStatus::PossiblyStale => 2,
}
}
fn availability_status(value: u64) -> AvailabilityStatus {
match value {
1 => AvailabilityStatus::Active,
2 => AvailabilityStatus::PossiblyStale,
_ => AvailabilityStatus::Unknown,
}
}
fn normalize_bounded(value: &str, max_chars: usize) -> String {
value
.chars()
.take(max_chars)
.collect::<String>()
.nfkc()
.collect::<String>()
.to_lowercase()
}
fn extensions(record: &TorrentRecord) -> BTreeSet<String> {
record
.files
.iter()
.take(512)
.filter_map(|file| Path::new(&file.path).extension())
.filter_map(|extension| extension.to_str())
.map(str::to_lowercase)
.collect()
}
fn text(
document: &TantivyDocument,
field: tantivy::schema::Field,
name: &'static str,
) -> Result<String, SearchError> {
document
.get_first(field)
.and_then(|value| value.as_str())
.map(str::to_owned)
.ok_or(SearchError::MissingField(name))
}
fn number(
document: &TantivyDocument,
field: tantivy::schema::Field,
name: &'static str,
) -> Result<u64, SearchError> {
document
.get_first(field)
.and_then(|value| value.as_u64())
.ok_or(SearchError::MissingField(name))
}
dht-search/src/search/filter.rs
+222
View File
@@ -0,0 +1,222 @@
// 负责将搜索选项组合为 Tantivy 查询条件但不执行查询
use std::ops::Bound;
use tantivy::{
Term,
query::{AllQuery, BooleanQuery, BoostQuery, Occur, Query, RangeQuery, RegexQuery, TermQuery},
schema::IndexRecordOption,
};
use unicode_normalization::UnicodeNormalization;
use super::{
SearchError,
document::availability_number,
query::{SearchOptions, SearchSort},
schema::SearchFields,
};
pub(crate) struct PreparedQuery {
pub(crate) query: Box<dyn Query>,
pub(crate) sort: SearchSort,
}
pub(crate) fn prepare(
options: &SearchOptions,
fields: SearchFields,
) -> Result<PreparedQuery, SearchError> {
let mut clauses: Vec<Box<dyn Query>> = Vec::new();
let query_text = options.query.trim();
if query_text.is_empty() || query_text == "*" {
clauses.push(Box::new(AllQuery));
} else if options.regex {
clauses.push(regex_query(query_text, fields)?);
} else {
clauses.push(text_query(query_text, fields));
}
if let Some(content_key) = options.content_key {
clauses.push(Box::new(TermQuery::new(
Term::from_field_text(fields.content_key, &hex::encode(content_key)),
IndexRecordOption::Basic,
)));
}
add_range(
&mut clauses,
fields.total_size,
options.min_size,
options.max_size,
);
add_range(
&mut clauses,
fields.file_count,
options.min_files,
options.max_files,
);
add_range(
&mut clauses,
fields.first_seen,
options.first_seen_after,
options.first_seen_before,
);
add_range(
&mut clauses,
fields.last_seen,
options.last_seen_after,
options.last_seen_before,
);
if let Some(extension) = &options.extension {
let extension = extension.trim().trim_start_matches('.').to_lowercase();
if !extension.is_empty() {
clauses.push(Box::new(TermQuery::new(
Term::from_field_text(fields.extensions, &extension),
IndexRecordOption::Basic,
)));
}
}
if let Some(status) = options.availability {
let value = availability_number(status);
add_range(
&mut clauses,
fields.availability_status,
Some(value),
Some(value),
);
}
if let Some(level) = options.heat {
let (min, max) = match level {
crate::domain::HeatLevel::Hot => (75, 100),
crate::domain::HeatLevel::Active => (50, 74),
crate::domain::HeatLevel::Normal => (25, 49),
crate::domain::HeatLevel::Cold => (0, 24),
};
add_range(&mut clauses, fields.heat_score, Some(min), Some(max));
}
let query = if clauses.len() == 1 {
clauses.pop().expect("one query clause exists")
} else {
Box::new(BooleanQuery::intersection(clauses))
};
let sort = options.sort.unwrap_or_else(|| {
if query_text.is_empty() || query_text == "*" {
SearchSort::Latest
} else {
SearchSort::Relevance
}
});
Ok(PreparedQuery { query, sort })
}
fn regex_query(pattern: &str, fields: SearchFields) -> Result<Box<dyn Query>, SearchError> {
let mut pattern = pattern.to_lowercase();
let anchored_start = pattern.starts_with('^');
if anchored_start {
pattern.remove(0);
}
let anchored_end = pattern.ends_with('$')
&& pattern[..pattern.len() - 1]
.chars()
.rev()
.take_while(|character| *character == '\\')
.count()
.is_multiple_of(2);
if anchored_end {
pattern.pop();
}
let prefix = if anchored_start { "" } else { ".*" };
let suffix = if anchored_end { "" } else { ".*" };
let contains_pattern = format!("{prefix}({pattern}){suffix}");
Ok(Box::new(RegexQuery::from_pattern(
&contains_pattern,
fields.regex_text,
)?))
}
fn text_query(query: &str, fields: SearchFields) -> Box<dyn Query> {
let normalized = normalize_text(query.trim());
if normalized.len() == 40 && normalized.bytes().all(|byte| byte.is_ascii_hexdigit()) {
return Box::new(TermQuery::new(
Term::from_field_text(fields.info_hash, &normalized),
IndexRecordOption::Basic,
));
}
let terms = query_terms(query);
if terms.is_empty() {
return Box::new(AllQuery);
}
let mut required = Vec::with_capacity(terms.len());
for term in terms {
let alternatives: Vec<(Occur, Box<dyn Query>)> = vec![
(
Occur::Should,
Box::new(BoostQuery::new(
Box::new(TermQuery::new(
Term::from_field_text(fields.name, &term),
IndexRecordOption::WithFreqs,
)),
3.0,
)),
),
(
Occur::Should,
Box::new(BoostQuery::new(
Box::new(TermQuery::new(
Term::from_field_text(fields.aliases, &term),
IndexRecordOption::WithFreqs,
)),
2.0,
)),
),
(
Occur::Should,
Box::new(TermQuery::new(
Term::from_field_text(fields.files_text, &term),
IndexRecordOption::WithFreqs,
)),
),
];
required.push((
Occur::Must,
Box::new(BooleanQuery::new(alternatives)) as Box<dyn Query>,
));
}
Box::new(BooleanQuery::new(required))
}
fn query_terms(query: &str) -> Vec<String> {
normalize_text(query)
.split_whitespace()
.flat_map(|part| {
let chars: Vec<_> = part.chars().collect();
if chars.len() <= 20 {
vec![part.to_owned()]
} else {
chars
.windows(20)
.map(|window| window.iter().collect())
.collect()
}
})
.collect()
}
fn normalize_text(value: &str) -> String {
value.nfkc().collect::<String>().to_lowercase()
}
fn add_range(
clauses: &mut Vec<Box<dyn Query>>,
field: tantivy::schema::Field,
min: Option<u64>,
max: Option<u64>,
) {
if min.is_none() && max.is_none() {
return;
}
let lower = min
.map(|value| Bound::Included(Term::from_field_u64(field, value)))
.unwrap_or(Bound::Unbounded);
let upper = max
.map(|value| Bound::Included(Term::from_field_u64(field, value)))
.unwrap_or(Bound::Unbounded);
clauses.push(Box::new(RangeQuery::new(lower, upper)));
}
dht-search/src/search/indexer.rs
+12 -356
View File
@@ -1,29 +1,24 @@
// 负责批量写入删除提交和从权威存储重建 Tantivy 索引
use std::{
collections::BTreeSet,
ops::Bound,
path::Path,
sync::{Arc, Mutex},
};
use crate::domain::ContentGroup;
use crate::storage::TorrentRepository;
use tantivy::{
DocAddress, Index, IndexReader, IndexWriter, Order, ReloadPolicy, Searcher, TantivyDocument,
Term,
collector::{Count, TopDocs},
directory::MmapDirectory,
query::{AllQuery, BooleanQuery, BoostQuery, Occur, Query, RangeQuery, RegexQuery, TermQuery},
schema::{IndexRecordOption, Value},
query::Query,
tokenizer::{LowerCaser, NgramTokenizer, TextAnalyzer},
};
use unicode_normalization::UnicodeNormalization;
use crate::domain::{AvailabilityStatus, ContentGroup, Heat, HeatLevel, TorrentRecord};
use crate::storage::TorrentRepository;
use super::{
IndexingError, SearchError,
query::{AvailabilitySummary, SearchHit, SearchOptions, SearchPage, SearchSort},
query::{SearchOptions, SearchPage, SearchSort},
schema::{MIXED_NGRAM_TOKENIZER, SearchFields, build_schema},
};
@@ -111,7 +106,7 @@ impl SearchEngine {
fields.content_key,
&hex::encode(group.content_key),
));
writer.add_document(document(group, fields))?;
writer.add_document(super::document::from_group(group, fields))?;
}
writer.commit()?;
self.inner.reader.reload()?;
@@ -160,85 +155,10 @@ impl SearchEngine {
let offset = options.offset.min(MAX_OFFSET);
let limit = options.limit.clamp(1, MAX_PAGE_SIZE);
let fields = self.inner.fields;
let mut clauses: Vec<Box<dyn Query>> = Vec::new();
let query_text = options.query.trim();
if query_text.is_empty() || query_text == "*" {
clauses.push(Box::new(AllQuery));
} else if options.regex {
clauses.push(regex_query(query_text, fields)?);
} else {
clauses.push(text_query(query_text, fields));
}
if let Some(content_key) = options.content_key {
clauses.push(Box::new(TermQuery::new(
Term::from_field_text(fields.content_key, &hex::encode(content_key)),
IndexRecordOption::Basic,
)));
}
add_range(
&mut clauses,
fields.total_size,
options.min_size,
options.max_size,
);
add_range(
&mut clauses,
fields.file_count,
options.min_files,
options.max_files,
);
add_range(
&mut clauses,
fields.first_seen,
options.first_seen_after,
options.first_seen_before,
);
add_range(
&mut clauses,
fields.last_seen,
options.last_seen_after,
options.last_seen_before,
);
if let Some(extension) = options.extension {
let extension = extension.trim().trim_start_matches('.').to_lowercase();
if !extension.is_empty() {
clauses.push(Box::new(TermQuery::new(
Term::from_field_text(fields.extensions, &extension),
IndexRecordOption::Basic,
)));
}
}
if let Some(status) = options.availability {
let value = availability_number(status);
add_range(
&mut clauses,
fields.availability_status,
Some(value),
Some(value),
);
}
if let Some(level) = options.heat {
let (min, max) = match level {
HeatLevel::Hot => (75, 100),
HeatLevel::Active => (50, 74),
HeatLevel::Normal => (25, 49),
HeatLevel::Cold => (0, 24),
};
add_range(&mut clauses, fields.heat_score, Some(min), Some(max));
}
let query: Box<dyn Query> = if clauses.len() == 1 {
clauses.pop().expect("one query clause exists")
} else {
Box::new(BooleanQuery::intersection(clauses))
};
let prepared = super::filter::prepare(&options, fields)?;
let query = prepared.query;
let searcher = self.inner.reader.searcher();
let sort = options.sort.unwrap_or_else(|| {
if query_text.is_empty() || query_text == "*" {
SearchSort::Latest
} else {
SearchSort::Relevance
}
});
let sort = prepared.sort;
let total = searcher.search(query.as_ref(), &Count)?;
let documents = match sort {
SearchSort::Relevance => searcher
@@ -302,38 +222,7 @@ impl SearchEngine {
let mut hits = Vec::with_capacity(documents.len());
for (score, address) in documents {
let document: TantivyDocument = searcher.doc(address)?;
hits.push(SearchHit {
info_hash: text(&document, fields.info_hash, "info_hash")?,
name: text(&document, fields.display_name, "display_name")?,
total_size: number(&document, fields.total_size, "total_size")?,
file_count: number(&document, fields.file_count, "file_count")?,
first_seen: number(&document, fields.first_seen, "first_seen")?,
last_seen: number(&document, fields.last_seen, "last_seen")?,
seen_count: number(&document, fields.seen_count, "seen_count")?,
content_key: text(&document, fields.content_key, "content_key")?,
variant_count: number(&document, fields.variant_count, "variant_count")?,
score,
heat: Heat::from_score(
number(&document, fields.heat_score, "heat_score")?.min(100) as u8,
),
availability: AvailabilitySummary {
status: availability_status(number(
&document,
fields.availability_status,
"availability_status",
)?),
last_verified_at: match number(
&document,
fields.last_verified_at,
"last_verified_at",
)? {
0 => None,
value => Some(value),
},
reachable_peers: number(&document, fields.reachable_peers, "reachable_peers")?
.min(u64::from(u32::MAX)) as u32,
},
});
hits.push(super::document::to_hit(&document, fields, score)?);
}
Ok(SearchPage {
total,
@@ -345,191 +234,6 @@ impl SearchEngine {
}
}
fn document(group: &ContentGroup, fields: SearchFields) -> TantivyDocument {
const MAX_INDEXED_FILES: usize = 512;
const MAX_PATH_TEXT_BYTES: usize = 32 * 1024;
let record = &group.representative;
let mut document = TantivyDocument::default();
document.add_text(fields.info_hash, record.info_hash.to_string());
document.add_text(fields.name, normalize_bounded(&record.name, 512));
document.add_text(fields.regex_text, normalize_bounded(&record.name, 512));
document.add_text(fields.display_name, &record.name);
for alias in &group.aliases {
let alias = normalize_bounded(alias, 512);
document.add_text(fields.aliases, &alias);
document.add_text(fields.regex_text, alias);
}
let mut indexed_path_bytes = 0_usize;
for file in record.files.iter().take(MAX_INDEXED_FILES) {
let path = normalize_bounded(&file.path, 512);
if indexed_path_bytes.saturating_add(path.len()) > MAX_PATH_TEXT_BYTES {
break;
}
indexed_path_bytes += path.len();
document.add_text(fields.files_text, &path);
document.add_text(fields.regex_text, path);
}
for extension in extensions(record) {
document.add_text(fields.extensions, extension);
}
document.add_u64(fields.total_size, record.total_size);
document.add_u64(fields.file_count, record.files.len() as u64);
document.add_u64(fields.first_seen, group.first_seen);
document.add_u64(fields.last_seen, group.last_seen);
document.add_u64(fields.seen_count, group.seen_count);
document.add_text(fields.content_key, hex::encode(group.content_key));
document.add_u64(
fields.availability_status,
availability_number(group.availability.status),
);
document.add_u64(
fields.reachable_peers,
u64::from(group.availability.reachable_peers),
);
document.add_u64(
fields.last_verified_at,
group.availability.last_verified_at.unwrap_or(0),
);
document.add_u64(fields.heat_score, u64::from(group.heat.score));
document.add_u64(fields.variant_count, group.variant_count);
document
}
fn regex_query(pattern: &str, fields: SearchFields) -> Result<Box<dyn Query>, SearchError> {
let mut pattern = pattern.to_lowercase();
let anchored_start = pattern.starts_with('^');
if anchored_start {
pattern.remove(0);
}
let anchored_end = pattern.ends_with('$')
&& pattern[..pattern.len() - 1]
.chars()
.rev()
.take_while(|character| *character == '\\')
.count()
.is_multiple_of(2);
if anchored_end {
pattern.pop();
}
let prefix = if anchored_start { "" } else { ".*" };
let suffix = if anchored_end { "" } else { ".*" };
let contains_pattern = format!("{prefix}({pattern}){suffix}");
Ok(Box::new(RegexQuery::from_pattern(
&contains_pattern,
fields.regex_text,
)?))
}
fn text_query(query: &str, fields: SearchFields) -> Box<dyn Query> {
let normalized = normalize_text(query.trim());
if normalized.len() == 40 && normalized.bytes().all(|byte| byte.is_ascii_hexdigit()) {
return Box::new(TermQuery::new(
Term::from_field_text(fields.info_hash, &normalized),
IndexRecordOption::Basic,
));
}
let terms = query_terms(query);
if terms.is_empty() {
return Box::new(AllQuery);
}
let mut required = Vec::with_capacity(terms.len());
for term in terms {
let mut alternatives: Vec<(Occur, Box<dyn Query>)> = vec![
(
Occur::Should,
Box::new(BoostQuery::new(
Box::new(TermQuery::new(
Term::from_field_text(fields.name, &term),
IndexRecordOption::WithFreqs,
)),
3.0,
)),
),
(
Occur::Should,
Box::new(BoostQuery::new(
Box::new(TermQuery::new(
Term::from_field_text(fields.aliases, &term),
IndexRecordOption::WithFreqs,
)),
2.0,
)),
),
(
Occur::Should,
Box::new(TermQuery::new(
Term::from_field_text(fields.files_text, &term),
IndexRecordOption::WithFreqs,
)),
),
];
if term.len() == 40 && term.bytes().all(|byte| byte.is_ascii_hexdigit()) {
alternatives.push((
Occur::Should,
Box::new(TermQuery::new(
Term::from_field_text(fields.info_hash, &term),
IndexRecordOption::Basic,
)),
));
}
required.push((
Occur::Must,
Box::new(BooleanQuery::new(alternatives)) as Box<dyn Query>,
));
}
Box::new(BooleanQuery::new(required))
}
fn query_terms(query: &str) -> Vec<String> {
normalize_text(query)
.split_whitespace()
.flat_map(|part| {
let chars: Vec<_> = part.chars().collect();
if chars.len() <= 20 {
vec![part.to_owned()]
} else {
chars
.windows(20)
.map(|window| window.iter().collect())
.collect()
}
})
.collect()
}
fn normalize_text(value: &str) -> String {
value.nfkc().collect::<String>().to_lowercase()
}
fn normalize_bounded(value: &str, max_chars: usize) -> String {
value
.chars()
.take(max_chars)
.collect::<String>()
.nfkc()
.collect::<String>()
.to_lowercase()
}
fn add_range(
clauses: &mut Vec<Box<dyn Query>>,
field: tantivy::schema::Field,
min: Option<u64>,
max: Option<u64>,
) {
if min.is_none() && max.is_none() {
return;
}
let lower = min
.map(|value| Bound::Included(Term::from_field_u64(field, value)))
.unwrap_or(Bound::Unbounded);
let upper = max
.map(|value| Bound::Included(Term::from_field_u64(field, value)))
.unwrap_or(Bound::Unbounded);
clauses.push(Box::new(RangeQuery::new(lower, upper)));
}
fn sorted_documents(
searcher: &Searcher,
query: &dyn Query,
@@ -550,22 +254,6 @@ fn sorted_documents(
.collect())
}
fn availability_number(status: AvailabilityStatus) -> u64 {
match status {
AvailabilityStatus::Unknown => 0,
AvailabilityStatus::Active => 1,
AvailabilityStatus::PossiblyStale => 2,
}
}
fn availability_status(value: u64) -> AvailabilityStatus {
match value {
1 => AvailabilityStatus::Active,
2 => AvailabilityStatus::PossiblyStale,
_ => AvailabilityStatus::Unknown,
}
}
#[cfg(test)]
fn unix_timestamp() -> u64 {
std::time::SystemTime::now()
@@ -574,47 +262,15 @@ fn unix_timestamp() -> u64 {
.as_secs()
}
fn extensions(record: &TorrentRecord) -> BTreeSet<String> {
record
.files
.iter()
.take(512)
.filter_map(|file| Path::new(&file.path).extension())
.filter_map(|extension| extension.to_str())
.map(str::to_lowercase)
.collect()
}
fn text(
document: &TantivyDocument,
field: tantivy::schema::Field,
name: &'static str,
) -> Result<String, SearchError> {
document
.get_first(field)
.and_then(|value| value.as_str())
.map(str::to_owned)
.ok_or(SearchError::MissingField(name))
}
fn number(
document: &TantivyDocument,
field: tantivy::schema::Field,
name: &'static str,
) -> Result<u64, SearchError> {
document
.get_first(field)
.and_then(|value| value.as_u64())
.ok_or(SearchError::MissingField(name))
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use tempfile::TempDir;
use crate::domain::{ContentGroupBuilder, InfoHash, TorrentFile, TorrentRecord};
use crate::domain::{
AvailabilityStatus, ContentGroupBuilder, InfoHash, TorrentFile, TorrentRecord,
};
use crate::storage::{RocksTorrentRepository, TorrentRepository};
use super::*;
dht-search/src/search/mod.rs
+2
View File
@@ -1,5 +1,7 @@
// 负责暴露全文搜索抽象并隐藏 Tantivy 的具体实现细节
mod document;
mod filter;
mod indexer;
mod query;
mod schema;
dht-search/tests/storage_search_flow.rs
+48
View File
@@ -0,0 +1,48 @@
// 负责从 crate 外部验证持久化索引搜索和重启恢复的公开组合契约
#![cfg(feature = "rocksdb-storage")]
use dht_crawler::{FileInfo, TorrentInfo};
use dht_search::{
domain::TorrentRecord,
search::SearchEngine,
storage::{RocksTorrentRepository, TorrentRepository, UpsertOutcome},
};
use tempfile::TempDir;
#[test]
fn public_components_compose_into_a_restart_safe_search_flow() {
let directory = TempDir::new().unwrap();
let rocksdb = directory.path().join("rocksdb");
let tantivy = directory.path().join("tantivy");
let record = TorrentRecord::try_from(TorrentInfo {
info_hash: "1212121212121212121212121212121212121212".into(),
magnet_link: String::new(),
name: "Public API 测试资源".into(),
total_size: 42,
files: vec![FileInfo {
path: "docs/public-api.txt".into(),
size: 42,
}],
piece_length: 16_384,
peers: Vec::new(),
timestamp: 100,
})
.unwrap();
{
let repository = RocksTorrentRepository::open(&rocksdb).unwrap();
assert_eq!(
repository.upsert(record.clone()).unwrap(),
UpsertOutcome::Inserted
);
let search = SearchEngine::open(&tantivy).unwrap();
assert_eq!(search.index_pending(&repository, 100, 100).unwrap(), 1);
let page = search.search("public-api", 0, 10).unwrap();
assert_eq!(page.total, 1);
assert_eq!(page.hits[0].info_hash, record.info_hash.to_string());
}
let reopened = RocksTorrentRepository::open(&rocksdb).unwrap();
assert_eq!(reopened.get(record.info_hash).unwrap(), Some(record));
}