feat: 实现持久化 DHT 搜索服务

This commit is contained in:
chuan
2026-08-06 22:09:21 +08:00
parent 2e8595575c
commit d25003d446
40 changed files with 3505 additions and 69 deletions
+6 -1
View File
@@ -1,5 +1,10 @@
# Rust
/target/
/.tools/
/.run-data/
/.remote-data/
/data/
/dht-search.toml
Cargo.lock
**/*.rs.bk
@@ -19,4 +24,4 @@ torrents/
# 个人脚本(不提交到仓库)
scripts/
opencodes/
opencodes/
+4
View File
@@ -90,6 +90,10 @@ Bloom Filter 只能作为前置加速结构不得作为最终去重依据
新业务代码写入 `dht-search`
项目阶段任务完成状态和验收标准统一维护在根目录 `TODOS.md`
需求实现或技术决策发生变化时必须同步更新 `TODOS.md`
`dht-crawler` 只接受可复用的 DHT 基础能力不得包含数据库搜索接口或部署逻辑
每个 Rust 文件顶部必须使用中文行注释描述该文件的功能边界且注释行尾不添加标点
+7 -1
View File
@@ -10,6 +10,12 @@ dht-crawler/ 可独立复用的 DHT 协议与 Metadata 获取基础库
opencodes/ 不参与构建的参考项目
```
当前应用层仅完成目录和依赖准备尚未实现业务逻辑
当前已完成 DHT 基础能力和 RocksDB 本地持久化基础下一阶段将接通真实采集写入管线
基础库的使用方式和指标说明见 [`dht-crawler/README.md`](dht-crawler/README.md)
当前实施阶段和后续计划见 [`TODOS.md`](TODOS.md)
应用配置模板见 [`dht-search.example.toml`](dht-search.example.toml)
应用构建运行和 API 文档见 [`dht-search/README.md`](dht-search/README.md)
+272
View File
@@ -0,0 +1,272 @@
# DHT 元数据搜索服务计划
本文档记录项目当前规划实施顺序和完成状态
它是随需求实现结果性能数据和部署条件持续调整的活文档
## 维护规则
- 已经通过验收的任务使用 `[x]` 标记
- 正在规划但尚未完成的任务使用 `[ ]` 标记
- 需求变化时允许新增删除拆分合并或调整阶段顺序
- 调整计划时同步修改任务说明依赖关系和验收标准
- 不因代码已经存在就标记完成必须满足对应验收标准
- 发现原方案不合适时记录新决策并更新后续阶段
- 每次完成一个可交付功能时同步更新本文档
## 当前技术方向
- `dht-crawler` 负责可复用的 DHT 协议节点发现 Peer 查找和 Metadata 下载
- `dht-search` 负责持久化去重索引搜索接口配置和运行生命周期
- RocksDB 保存权威数据去重信息和任务状态
- Tantivy 保存可以从 RocksDB 重建的搜索索引
- Axum 提供搜索详情统计和健康检查接口
- 所有长期任务通过有界队列和背压控制资源占用
如果实际运行证明 RocksDB 的构建部署或资源成本不合适可以重新评估 redb SQLite 或其他存储方案
## 阶段零 项目基础
### 目标
建立清晰的 workspace 边界开发规则和可持续验证的基础库
### 任务
- [x] 将 workspace 扁平化为 `dht-crawler``dht-search`
- [x] 使用当前 Git 配置统一作者仓库许可证和 edition 元数据
- [x] 编写 `AGENTS.md` 记录架构边界和开发约定
- [x] 将最终应用与可复用 DHT 基础库分离
- [x] 实现 BEP-51 `sample_infohashes` 主动发现
- [x] 实现主动 Peer 查找和 Metadata 获取
- [x] 验证远程公网设备能够持续获取 Metadata
- [x] 确认 Xray 全局代理会影响 Metadata TCP 连接并完成旁路验证
- [x] 保持基础库测试通过
### 验收标准
- [x] `cargo check --workspace --all-targets` 通过
- [x] `dht-crawler` 单元测试通过
- [x] `opencodes` 参考项目不参与 workspace 构建
## 阶段一 本地持久化基础
### 目标
建立跨重启保留的权威数据源并完成精确去重和内容聚合基础
### 任务
- [x] 定义二十字节 `InfoHash` 类型和十六进制转换
- [x] 定义 `TorrentRecord` `TorrentFile``IndexState`
- [x] 校验名称文件列表文件总大小和 infohash
- [x] 使用 BLAKE3 计算版本化内容指纹
- [x] 规范化 Unicode 路径分隔符大小写和文件顺序
- [x] 保留真实子目录避免内容指纹碰撞
- [x] 定义版本化 RocksDB 二进制键空间
- [x] 实现数据库 schema 版本检查
- [x] 实现 infohash 精确查询和存在性判断
- [x] 实现新记录 WriteBatch 原子写入
- [x] 实现重复 infohash 的 `last_seen` `seen_count` 和 Peer 更新
- [x] 实现相同内容不同 infohash 的聚合映射
- [x] 实现待索引记录查询和索引完成标记
- [x] 配置 Bloom Filter LZ4 压缩和有限 block cache
- [x] 准备 Windows 本地 RocksDB 构建所需的 libclang
- [x] 将本地构建工具目录排除出 Git
### 验收标准
- [x] 数据库关闭并重新打开后记录仍可读取
- [x] 重复写入不会创建第二条 torrent 记录
- [x] 重复写入会正确增加发现次数
- [x] 相同内容的不同 infohash 可以独立保存并聚合查询
- [x] Metadata 主体内容映射和待索引标记原子写入
- [x] RocksDB 功能测试通过
- [x] `dht-search` Clippy `-D warnings` 通过
## 阶段二 采集持久化闭环
### 目标
让 DHT 获取的真实 Metadata 自动进入有界持久化管线并支持安全停止和重新启动
### 任务
- [x] 定义应用配置结构和默认配置文件
- [x] 支持通过配置指定固定数据目录
- [x] 支持配置 DHT 端口并发队列容量和 Metadata 限制
- [x] 初始化 RocksDB repository 并处理启动错误
- [x]`TorrentInfo` callback 转换为 `TorrentRecord`
- [x] 建立有界持久化队列并实现背压
- [x] 使用专用阻塞任务执行 RocksDB 操作避免阻塞 Tokio worker
- [x] 在 Metadata 下载前查询持久化 infohash 状态减少重复下载
- [x] 在 BEP-51 Peer Lookup 前批量查询 RocksDB 并更新已有 infohash 发现状态
- [x] 将已存在记录更新为再次发现而不是重复创建
- [x] 增加接收写入重复拒绝失败和队列深度指标
- [x] 实现 `Ctrl+C` `SIGINT``SIGTERM` 优雅退出
- [x] 退出时停止接收新任务并排空或持久化剩余任务
- [x] 支持重新启动后继续使用原数据库
- [x] 将 example 运行方式替换为正式 `dht-search` 二进制
- [x] 支持通过运行时长参数进行间歇运行
### 验收标准
- [x] 本地运行可以持续向 RocksDB 写入真实 Metadata
- [x] 停止并重启后旧 infohash 不会作为新记录重复写入
- [x] 队列达到容量时内存不继续无界增长
- [x] 正常退出后已接受的任务不会静默丢失
- [ ] 远程设备运行一小时没有持续内存增长
- [x] 记录采集速度重复率数据库增长和写入延迟
## 阶段三 Tantivy 搜索索引
### 目标
让持久化 Metadata 支持快速全文搜索过滤排序和索引恢复
### 任务
- [x] 定义 Tantivy schema 和索引版本
- [x] 索引名称文件路径扩展名 infohash 和内容指纹
- [x] 将大小文件数时间和发现次数定义为 fast fields
- [ ] 设计中英文数字和文件名 tokenizer
- [x] 实现待索引任务批量消费
- [x] 实现按数量和时间间隔批量 commit
- [x] commit 成功后原子更新 RocksDB 索引状态
- [x] 实现关键词短语和精确 infohash 查询
- [ ] 实现大小时间扩展名和文件数过滤
- [x] 实现大小范围和扩展名过滤
- [ ] 实现相关性时间热度和大小排序
- [x] 实现分页并限制最大翻页成本
- [ ] 实现相同 `content_key` 结果折叠
- [x] 实现从 RocksDB 全量重建 Tantivy 索引
- [ ] 支持索引 schema 不兼容时安全重建
### 验收标准
- [x] 新写入记录在目标延迟内可搜索
- [ ] 搜索索引删除后可以从 RocksDB 完整重建
- [x] 索引过程中异常退出不会永久丢失文档
- [ ] 百万级测试数据常用查询延迟达到约定目标
## 阶段四 HTTP 搜索服务
### 目标
提供稳定可验证并且资源受限的搜索和详情接口
### 任务
- [x] 使用 Axum 建立 HTTP 服务
- [x] 实现 `/health``/ready` 接口
- [x] 实现 `/stats` 运行状态接口
- [x] 实现 `/search` 搜索过滤和分页接口
- [x] 实现 `/torrents/{infohash}` 详情接口
- [x] 定义统一错误响应
- [x] 限制查询长度分页大小和最大 offset
- [ ] 增加请求延迟错误率和并发指标
- [ ] 增加 API 单元测试和端到端测试
### 验收标准
- [x] API 能搜索真实采集数据
- [x] 非法参数返回稳定的客户端错误
- [x] 搜索查询在独立阻塞任务执行不会阻塞异步 worker
- [x] 健康检查能区分进程存活和服务可用
## 阶段五 质量过滤和重复内容控制
### 目标
减少垃圾数据和重复展示同时避免不可恢复的误删
### 任务
- [ ] 统计真实数据的 infohash 重复率和内容重复率
- [ ] 定义可配置的名称路径扩展名和大小过滤规则
- [ ] 定义 Metadata 最大大小文件数和路径长度限制
- [ ] 识别空名称异常路径大小溢出和文件数量攻击
- [ ] 设计可解释的名称标准化规则
- [ ] 为模糊相似结果生成聚合候选但不自动删除
- [ ] 支持黑名单规则版本和命中原因
- [ ] 保留被过滤记录的计数指标但避免保存大内容
- [ ] 增加误判测试和边界数据集
### 验收标准
- [ ] 精确重复不会重复下载和重复展示
- [ ] 内容重复可以折叠并保留全部 infohash
- [ ] 过滤规则可以配置更新和回滚
- [ ] 模糊去重不会直接造成数据丢失
## 阶段六 性能资源和长期运行
### 目标
以真实数据验证持续运行时的吞吐延迟磁盘放大和资源上限
### 任务
- [x] 增加 `find_node` Peer Lookup 新目标和 Metadata 建连的显式配置
- [x] 为主动 `find_node` `get_peers``sample_infohashes` 增加共享 UDP 查询总预算
- [x] 为 Metadata TCP 建连增加独立每秒速率限制
- [x] 将桌面默认配置调整为保守网络预算
- [x] 根据本机首次验证将主动 UDP 从 `40/s` 下调至 `10/s` 并将 Metadata 建连从 `5/s` 下调至 `2/s`
- [x] 验证极保守配置运行三分钟不影响同机代理网络并安全退出
- [x] 将 BEP-51 采样准入压力反向传递到采样查询调度
- [x] 实现样本来源节点单点 `get_peers` 优先和失败后有限递归降级
- [x] 完成首轮三分钟对比并验证 Peer Lookup UDP 从 `278` 降至 `254` 且网络稳定
- [ ] 通过多轮或更长时间运行评估随机 DHT 样本下的 Metadata 成功率
- [ ] 根据公网设备长期实测设计超时率自动降速
- [ ] 建立采集存储索引和查询基准测试
- [ ] 记录每条元数据和每个索引文档的平均磁盘占用
- [ ] 记录 RocksDB block cache memtable 和 compaction 指标
- [ ] 记录 Tantivy IndexWriter 内存和 commit 延迟
- [ ] 根据实测调整批量大小队列容量和并发
- [ ] 增加磁盘剩余空间保护和只读降级策略
- [ ] 增加数据库备份检查点和恢复验证
- [ ] 增加日志轮转和保留策略
- [x] 验证间歇运行和正常退出恢复
- [ ] 验证二十四小时和七天连续运行
- [ ] 根据规模决定是否继续使用 RocksDB
### 验收标准
- [ ] 内存使用在目标上限内稳定
- [ ] 队列和缓存不会随运行时间无限增长
- [ ] 磁盘不足时能够安全停止写入
- [ ] 备份可以在独立目录恢复并搜索
- [ ] 连续运行期间没有数据格式损坏和不可恢复任务
## 阶段七 部署和运维
### 目标
让应用可以在公网 Linux 设备上重复构建部署监控停止和恢复
### 任务
- [x] 固化 Linux 目标构建方式和 RocksDB 构建依赖
- [x] 生成 release 二进制并使用 SHA-256 校验部署
- [ ] 定义配置数据日志和索引目录布局
- [ ] 编写 systemd service
- [ ] 编写 systemd timer 支持间歇运行
- [x] 使用专用低权限 UID 运行验证
- [x] 固化 Xray 环境下的最小范围网络旁路
- [x] 验证高并发运行需要 `LimitNOFILE=65536`
- [ ] 实现启动前数据目录权限检查
- [ ] 实现优雅升级和回滚流程
- [ ] 编写备份恢复和故障排查文档
### 验收标准
- [ ] 新设备可以按文档完成部署
- [ ] 服务重启不会丢失已提交数据
- [ ] Xray 旁路只影响爬虫进程
- [ ] 更新失败时可以恢复上一版本二进制和数据
## 当前下一步
继续完善阶段三和阶段四
下一步实现时间文件数过滤排序策略内容聚合展示以及自动化 API 端到端测试
+10 -5
View File
@@ -156,13 +156,14 @@ let options = DHTOptions {
metadata: MetadataOptions {
timeout_secs: 5,
max_queue_size: 20_000,
max_worker_count: 256,
max_worker_count: 8,
max_connects_per_second: 2,
..Default::default()
},
crawl: CrawlOptions {
rate_limit: RateLimitOptions {
max_find_node_rate_per_sec: 300,
max_in_flight: 768,
max_find_node_rate_per_sec: 6,
max_in_flight: 12,
..Default::default()
},
..Default::default()
@@ -175,8 +176,8 @@ let options = DHTOptions {
| 类型 | 控制内容 |
|---|---|
| `DHTOptions` | 监听端口、网络模式顶层队列 |
| `MetadataOptions` | 下载超时、队列、并发和失败 Peer 缓存 |
| `DHTOptions` | 监听端口、网络模式顶层队列和主动 UDP 查询总预算 |
| `MetadataOptions` | 下载超时、队列、并发、每秒 TCP 建连和失败 Peer 缓存 |
| `PeerLookupOptions` | 主动 `get_peers` 的速率与并发 |
| `SampleInfohashesOptions` | BEP-51 采样速率、并发、超时、退避和 Hash 去重容量 |
| `RateLimitOptions` | `find_node`、在途请求和 UDP 回复预算 |
@@ -187,6 +188,10 @@ let options = DHTOptions {
完整字段和默认值以 [docs.rs API 文档](https://docs.rs/dht-crawler) 为准。需要注意:
- BEP-51 采样 hash 可以通过 `DHTServer::on_sampled_hashes` 批量异步准入
- 采样准入队列有固定容量并在压力升高时暂停新的 BEP-51 查询
- 带首选节点的 Peer Lookup 先执行单点查询只有失败后才进入有限迭代查找
- `DHTOptions::default()` 使用 `Ipv4Only`
- `NetMode::DualStack` 会分别绑定 IPv4 和 IPv6 Socket
- `DHTServer::new()` 会立即在所有可用接口上绑定配置的 UDP 端口;
+49 -1
View File
@@ -1,4 +1,7 @@
use std::time::{Duration, Instant};
use std::{
sync::{Arc, Mutex},
time::{Duration, Instant},
};
/// Single-owner token bucket. It deliberately contains no atomics or locks.
pub(crate) struct RateBucket {
@@ -8,6 +11,38 @@ pub(crate) struct RateBucket {
last_refill: Instant,
}
#[derive(Clone)]
pub(crate) struct SharedRateBudget {
bucket: Arc<Mutex<RateBucket>>,
}
impl SharedRateBudget {
pub(crate) fn per_second(rate_per_sec: u32, burst: u32, initially_full: bool) -> Self {
Self {
bucket: Arc::new(Mutex::new(RateBucket::per_second(
rate_per_sec,
burst,
initially_full,
Instant::now(),
))),
}
}
pub(crate) fn try_take_one(&self, now: Instant) -> bool {
self.bucket
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.try_take_one(now)
}
pub(crate) fn refund_one(&self) {
self.bucket
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.refund_one();
}
}
impl RateBucket {
pub(crate) fn per_second(
rate_per_sec: u32,
@@ -129,4 +164,17 @@ mod tests {
let mut bucket = RateBucket::per_minute(600, 10, false, start);
assert_eq!(bucket.try_take(100, start + Duration::from_millis(500)), 5);
}
#[test]
fn shared_budget_is_global_across_clones() {
let budget = SharedRateBudget::per_second(10, 2, true);
let clone = budget.clone();
let now = Instant::now();
assert!(budget.try_take_one(now));
assert!(clone.try_take_one(now));
assert!(!budget.try_take_one(now));
clone.refund_one();
assert!(budget.try_take_one(now));
}
}
+32 -4
View File
@@ -1,5 +1,5 @@
use crate::bootstrap::{BootstrapGate, BootstrapSourcePool, resolve_bootstrap_nodes};
use crate::budget::RateBucket;
use crate::budget::{RateBucket, SharedRateBudget};
use crate::crawl_config::ResolvedCrawlConfig;
use crate::krpc::{for_each_response_node, send_find_node_query};
use crate::node_id::{
@@ -77,10 +77,15 @@ pub(crate) struct CrawlEngine {
pub(crate) snapshot: Arc<ArcSwap<RoutingSnapshot>>,
pub(crate) node_count: Arc<AtomicUsize>,
runtime_stats: DhtRuntimeStats,
outbound_query_budget: SharedRateBudget,
}
impl CrawlEngine {
pub(crate) fn new(config: ResolvedCrawlConfig, runtime_stats: DhtRuntimeStats) -> Self {
pub(crate) fn new(
config: ResolvedCrawlConfig,
runtime_stats: DhtRuntimeStats,
outbound_query_budget: SharedRateBudget,
) -> Self {
let (priority_tx, priority_rx) = mpsc::channel(config.priority_event_channel_capacity);
let (discovery_tx, discovery_rx) = mpsc::channel(config.discovery_event_channel_capacity);
Self {
@@ -91,6 +96,7 @@ impl CrawlEngine {
snapshot: Arc::new(ArcSwap::from_pointee(RoutingSnapshot::default())),
node_count: Arc::new(AtomicUsize::new(0)),
runtime_stats,
outbound_query_budget,
}
}
@@ -181,6 +187,7 @@ impl CrawlEngine {
metadata_queue_len,
max_metadata_queue_size,
runtime_stats: self.runtime_stats.clone(),
outbound_query_budget: self.outbound_query_budget.clone(),
shutdown,
});
tokio::spawn(actor.run());
@@ -276,6 +283,7 @@ struct CrawlActorInit {
metadata_queue_len: Arc<AtomicUsize>,
max_metadata_queue_size: usize,
runtime_stats: DhtRuntimeStats,
outbound_query_budget: SharedRateBudget,
shutdown: CancellationToken,
}
@@ -305,6 +313,7 @@ struct CrawlActor {
next_tid: u64,
metrics: ActorMetrics,
runtime_stats: DhtRuntimeStats,
outbound_query_budget: SharedRateBudget,
shutdown: CancellationToken,
}
@@ -324,6 +333,7 @@ impl CrawlActor {
metadata_queue_len,
max_metadata_queue_size,
runtime_stats,
outbound_query_budget,
shutdown,
} = init;
let now = Instant::now();
@@ -373,6 +383,7 @@ impl CrawlActor {
next_tid: 1,
metrics: ActorMetrics::default(),
runtime_stats,
outbound_query_budget,
shutdown,
}
}
@@ -546,6 +557,17 @@ impl CrawlActor {
}
fn schedule_one(&mut self, now: Instant) -> bool {
if !self.outbound_query_budget.try_take_one(now) {
return false;
}
let scheduled = self.schedule_one_with_budget(now);
if !scheduled {
self.outbound_query_budget.refund_one();
}
scheduled
}
fn schedule_one_with_budget(&mut self, now: Instant) -> bool {
if self.pool.len() < self.config.low_watermark
&& let Some(addr) = self.bootstrap_queue.front().copied()
{
@@ -855,6 +877,7 @@ mod tests {
metadata_queue_len: Arc::new(AtomicUsize::new(0)),
max_metadata_queue_size: 100_000,
runtime_stats: runtime_stats.clone(),
outbound_query_budget: SharedRateBudget::per_second(10_000, 10_000, true),
shutdown: CancellationToken::new(),
});
(actor, egress_rx, runtime_stats)
@@ -864,6 +887,7 @@ mod tests {
fn saturated_head_subnet_does_not_block_later_node() {
let config = ResolvedCrawlConfig::from_options(&CrawlOptions::default());
let max_per_subnet = config.max_in_flight_per_subnet;
let max_in_flight = config.max_in_flight;
let (mut actor, mut egress_rx, runtime_stats) = test_actor(config);
let now = Instant::now() + Duration::from_secs(1);
let blocked = node(1, "8.8.8.8:6881");
@@ -886,7 +910,7 @@ mod tests {
let snapshot = runtime_stats.snapshot();
assert_eq!(snapshot.queries_new, 1);
assert_eq!(snapshot.find_node_in_flight, 1);
assert_eq!(snapshot.find_node_in_flight_max, 512);
assert_eq!(snapshot.find_node_in_flight_max, max_in_flight);
}
#[test]
@@ -905,7 +929,11 @@ mod tests {
crawl_priority_queue: config.priority_event_channel_capacity,
crawl_discovery_queue: config.discovery_event_channel_capacity,
});
let engine = CrawlEngine::new(config, stats.clone());
let engine = CrawlEngine::new(
config,
stats.clone(),
SharedRateBudget::per_second(10_000, 10_000, true),
);
engine.route_discovered(node(1, "8.8.8.8:1"));
engine.route_discovered(node(2, "1.1.1.1:2"));
+45 -2
View File
@@ -41,6 +41,36 @@ enum PeerFailureReason {
ConnectFailed,
}
struct ConnectRateLimiter {
interval: Duration,
next_start: Mutex<Instant>,
}
impl ConnectRateLimiter {
fn per_second(rate: u32) -> Self {
Self {
interval: Duration::from_secs_f64(1.0 / f64::from(rate.max(1))),
next_start: Mutex::new(Instant::now()),
}
}
async fn acquire(&self) {
let delay = {
let mut next_start = self
.next_start
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let now = Instant::now();
let reserved = (*next_start).max(now);
*next_start = reserved + self.interval;
reserved.saturating_duration_since(now)
};
if !delay.is_zero() {
tokio::time::sleep(delay).await;
}
}
}
#[cfg(feature = "metrics")]
impl PeerFailureReason {
fn as_str(self) -> &'static str {
@@ -162,6 +192,7 @@ pub struct RbitFetcher {
total_timeout: Duration,
runtime_stats: DhtRuntimeStats,
peer_failure_cache: Arc<PeerFailureCache>,
connect_rate_limiter: Arc<ConnectRateLimiter>,
}
impl RbitFetcher {
@@ -170,11 +201,12 @@ impl RbitFetcher {
/// [`DHTServer`](crate::DHTServer) normally constructs this component from
/// [`MetadataOptions`](crate::MetadataOptions).
pub fn new(timeout_secs: u64) -> Self {
Self::new_with_runtime_stats(timeout_secs, 200_000, 60, DhtRuntimeStats::default())
Self::new_with_runtime_stats(timeout_secs, 32, 200_000, 60, DhtRuntimeStats::default())
}
pub(crate) fn new_with_runtime_stats(
timeout_secs: u64,
max_connects_per_second: u32,
peer_failure_cache_capacity: usize,
peer_failure_ttl_secs: u64,
runtime_stats: DhtRuntimeStats,
@@ -186,6 +218,7 @@ impl RbitFetcher {
peer_failure_cache_capacity,
Duration::from_secs(peer_failure_ttl_secs),
)),
connect_rate_limiter: Arc::new(ConnectRateLimiter::per_second(max_connects_per_second)),
}
}
@@ -229,6 +262,7 @@ impl RbitFetcher {
return MetadataFetchOutcome::SkippedCached;
}
self.connect_rate_limiter.acquire().await;
on_attempt();
self.runtime_stats.metadata_peer_attempt();
#[cfg(feature = "metrics")]
@@ -559,6 +593,15 @@ mod tests {
use super::*;
use tokio::net::TcpListener;
#[tokio::test]
async fn connection_rate_limiter_spaces_attempts() {
let limiter = ConnectRateLimiter::per_second(20);
limiter.acquire().await;
let started = Instant::now();
limiter.acquire().await;
assert!(started.elapsed() >= Duration::from_millis(40));
}
#[test]
fn peer_failure_cache_is_socket_specific_and_expires() {
let start = Instant::now();
@@ -599,7 +642,7 @@ mod tests {
});
let stats = DhtRuntimeStats::default();
let fetcher = RbitFetcher::new_with_runtime_stats(1, 10, 60, stats.clone());
let fetcher = RbitFetcher::new_with_runtime_stats(1, 10, 10, 60, stats.clone());
let started = Instant::now();
assert!(matches!(
fetcher.fetch(&[7; 20], addr).await,
+153 -15
View File
@@ -1,4 +1,4 @@
use crate::budget::RateBucket;
use crate::budget::{RateBucket, SharedRateBudget};
use crate::krpc::{encode_get_peers_query, for_each_response_node, for_each_response_peer};
use crate::node_id::TransactionId;
use crate::protocol::DhtResponse;
@@ -40,6 +40,7 @@ struct PendingKey {
struct PendingQuery {
lookup_id: u64,
deadline: Instant,
preferred_phase: bool,
}
struct LookupState {
@@ -52,6 +53,7 @@ struct LookupState {
queried: usize,
outstanding: usize,
deadline: Instant,
preferred_phase: bool,
}
impl LookupState {
@@ -107,6 +109,7 @@ pub(crate) struct PeerLookupHandle {
pub(crate) struct PeerLookupRuntime {
pub(crate) options: PeerLookupOptions,
pub(crate) stats: DhtRuntimeStats,
pub(crate) outbound_query_budget: SharedRateBudget,
pub(crate) shutdown: CancellationToken,
}
@@ -153,6 +156,7 @@ pub(crate) fn spawn_peer_lookup(
let PeerLookupRuntime {
options,
stats,
outbound_query_budget,
shutdown,
} = runtime;
let (request_tx, request_rx) = mpsc::channel(REQUEST_CHANNEL_CAPACITY);
@@ -189,6 +193,7 @@ pub(crate) fn spawn_peer_lookup(
next_lookup_id: 1,
next_tid: 1,
runtime_stats: stats.clone(),
outbound_query_budget,
shutdown,
};
tokio::spawn(actor.run());
@@ -220,6 +225,7 @@ struct PeerLookupActor {
next_lookup_id: u64,
next_tid: u64,
runtime_stats: DhtRuntimeStats,
outbound_query_budget: SharedRateBudget,
shutdown: CancellationToken,
}
@@ -319,6 +325,7 @@ impl PeerLookupActor {
queried: 0,
outstanding: 0,
deadline: now + LOOKUP_TIMEOUT,
preferred_phase: preferred.is_some(),
},
);
self.active_hashes.insert(info_hash);
@@ -330,18 +337,28 @@ impl PeerLookupActor {
}
async fn dispatch_more(&mut self, lookup_id: u64, now: Instant) {
while let Some((node, info_hash)) = self.active.get_mut(&lookup_id).and_then(|state| {
if state.deadline <= now
|| state.peers.len() >= MAX_PEERS_PER_LOOKUP
|| state.queried >= MAX_QUERIES_PER_LOOKUP
|| state.outstanding >= MAX_CONCURRENT_QUERIES_PER_LOOKUP
{
return None;
loop {
if !self.outbound_query_budget.try_take_one(now) {
break;
}
let node = state.pop_closest()?;
state.queried += 1;
Some((node, state.info_hash))
}) {
let next = self.active.get_mut(&lookup_id).and_then(|state| {
if state.deadline <= now
|| state.peers.len() >= MAX_PEERS_PER_LOOKUP
|| state.queried >= MAX_QUERIES_PER_LOOKUP
|| state.outstanding >= MAX_CONCURRENT_QUERIES_PER_LOOKUP
|| (state.preferred_phase && state.outstanding > 0)
{
return None;
}
let preferred_phase = state.preferred_phase;
let node = state.pop_closest()?;
state.queried += 1;
Some((node, state.info_hash, preferred_phase))
});
let Some((node, info_hash, preferred_phase)) = next else {
self.outbound_query_budget.refund_one();
break;
};
let tid = self.next_transaction_id();
let key = PendingKey {
addr: node.addr,
@@ -359,7 +376,13 @@ impl PeerLookupActor {
None => false,
};
if !sent {
self.outbound_query_budget.refund_one();
self.runtime_stats.peer_lookup_send_failed();
if preferred_phase && let Some(state) = self.active.get_mut(&lookup_id) {
state.preferred_phase = false;
state.deadline = now + LOOKUP_TIMEOUT;
self.runtime_stats.peer_lookup_fallback();
}
continue;
}
@@ -369,6 +392,7 @@ impl PeerLookupActor {
PendingQuery {
lookup_id,
deadline,
preferred_phase,
},
);
self.pending_expiry.push_back((deadline, key));
@@ -395,6 +419,10 @@ impl PeerLookupActor {
return;
};
state.outstanding = state.outstanding.saturating_sub(1);
if pending.preferred_phase {
state.preferred_phase = false;
state.deadline = now + LOOKUP_TIMEOUT;
}
self.runtime_stats.peer_lookup_response();
let mut discovered = Vec::new();
@@ -417,6 +445,7 @@ impl PeerLookupActor {
}
let hash = state.info_hash_hex.clone();
let lookup_id = pending.lookup_id;
let preferred_succeeded = pending.preferred_phase && !discovered.is_empty();
let _ = state;
for peer in discovered {
@@ -436,6 +465,14 @@ impl PeerLookupActor {
.increment(1);
}
}
if preferred_succeeded {
self.runtime_stats.peer_lookup_preferred_succeeded();
self.finish_lookup(lookup_id);
return;
}
if pending.preferred_phase {
self.runtime_stats.peer_lookup_fallback();
}
self.dispatch_more(lookup_id, now).await;
}
@@ -456,6 +493,11 @@ impl PeerLookupActor {
let pending = self.pending.remove(&key).expect("pending lookup exists");
if let Some(state) = self.active.get_mut(&pending.lookup_id) {
state.outstanding = state.outstanding.saturating_sub(1);
if pending.preferred_phase {
state.preferred_phase = false;
state.deadline = now + LOOKUP_TIMEOUT;
self.runtime_stats.peer_lookup_fallback();
}
affected.insert(pending.lookup_id);
}
self.runtime_stats.peer_lookup_timeout();
@@ -516,13 +558,15 @@ mod tests {
async fn lookup_response_feeds_discovered_peer_back_to_metadata_scheduler() {
let local_socket = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
let remote_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
let fallback_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
let remote_addr = remote_socket.local_addr().unwrap();
let fallback_addr = fallback_socket.local_addr().unwrap();
let local_addr = local_socket.local_addr().unwrap();
let sockets = std::collections::HashMap::from([(local_addr, local_socket)]);
let snapshot = Arc::new(ArcSwap::from_pointee(RoutingSnapshot::from_nodes(
vec![NodeTuple {
id: [9; 20],
addr: remote_addr,
id: [8; 20],
addr: fallback_addr,
}],
1,
)));
@@ -538,12 +582,19 @@ mod tests {
PeerLookupRuntime {
options: PeerLookupOptions::default(),
stats: stats.clone(),
outbound_query_budget: SharedRateBudget::per_second(10_000, 10_000, true),
shutdown: shutdown.clone(),
},
);
handle
.request_tx
.send(PeerLookupRequest::new([3; 20]))
.send(PeerLookupRequest {
info_hash: [3; 20],
preferred_node: Some(NodeTuple {
id: [9; 20],
addr: remote_addr,
}),
})
.await
.unwrap();
@@ -557,6 +608,9 @@ mod tests {
let query: DhtMessage = serde_bencode::from_bytes(&buffer[..len]).unwrap();
assert_eq!(query.q.as_deref(), Some("get_peers"));
let tid: TransactionId = query.t.as_ref().try_into().unwrap();
tokio::time::sleep(Duration::from_millis(100)).await;
assert_eq!(stats.snapshot().peer_lookup_queries, 1);
assert!(fallback_socket.try_recv_from(&mut buffer).is_err());
handle.route_response(
remote_addr,
tid,
@@ -584,6 +638,90 @@ mod tests {
assert_eq!(snapshot.peer_lookup_queries, 1);
assert_eq!(snapshot.peer_lookup_responses, 1);
assert_eq!(snapshot.peer_lookup_peers_found, 1);
assert_eq!(snapshot.peer_lookup_preferred_succeeded, 1);
assert_eq!(snapshot.peer_lookup_fallbacks, 0);
shutdown.cancel();
}
#[tokio::test]
async fn preferred_node_without_peers_falls_back_to_iterative_lookup() {
let local_socket = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
let preferred_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
let fallback_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
let local_addr = local_socket.local_addr().unwrap();
let preferred_addr = preferred_socket.local_addr().unwrap();
let fallback_addr = fallback_socket.local_addr().unwrap();
let sockets = std::collections::HashMap::from([(local_addr, local_socket)]);
let snapshot = Arc::new(ArcSwap::from_pointee(RoutingSnapshot::from_nodes(
vec![NodeTuple {
id: [8; 20],
addr: fallback_addr,
}],
1,
)));
let (hash_tx, _hash_rx) = mpsc::channel(4);
let stats = DhtRuntimeStats::default();
let shutdown = CancellationToken::new();
let handle = spawn_peer_lookup(
NetMode::Ipv4Only,
[7; 20],
&sockets,
snapshot,
hash_tx,
PeerLookupRuntime {
options: PeerLookupOptions::default(),
stats: stats.clone(),
outbound_query_budget: SharedRateBudget::per_second(10_000, 10_000, true),
shutdown: shutdown.clone(),
},
);
handle
.request_tx
.send(PeerLookupRequest {
info_hash: [3; 20],
preferred_node: Some(NodeTuple {
id: [9; 20],
addr: preferred_addr,
}),
})
.await
.unwrap();
let mut buffer = [0u8; 512];
let (len, _) = tokio::time::timeout(
Duration::from_secs(1),
preferred_socket.recv_from(&mut buffer),
)
.await
.unwrap()
.unwrap();
let query: DhtMessage = serde_bencode::from_bytes(&buffer[..len]).unwrap();
let tid: TransactionId = query.t.as_ref().try_into().unwrap();
handle.route_response(
preferred_addr,
tid,
DhtResponse {
id: Some(serde_bytes::ByteBuf::from(vec![9; 20])),
nodes: None,
nodes6: None,
values: None,
samples: None,
num: None,
interval: None,
},
);
tokio::time::timeout(
Duration::from_secs(1),
fallback_socket.recv_from(&mut buffer),
)
.await
.unwrap()
.unwrap();
let snapshot = stats.snapshot();
assert_eq!(snapshot.peer_lookup_queries, 2);
assert_eq!(snapshot.peer_lookup_preferred_succeeded, 0);
assert_eq!(snapshot.peer_lookup_fallbacks, 1);
shutdown.cancel();
}
}
+43
View File
@@ -246,6 +246,10 @@ pub struct DhtRuntimeSnapshot {
pub peer_lookup_response_dropped: u64,
/// Discovered Peer endpoints dropped because Hash ingress was full.
pub peer_lookup_output_dropped: u64,
/// Preferred BEP-51 responder queries that returned at least one Peer.
pub peer_lookup_preferred_succeeded: u64,
/// Preferred-node attempts that fell back to iterative lookup.
pub peer_lookup_fallbacks: u64,
/// BEP-51 requests sent.
pub sample_infohashes_queries: u64,
/// Matched BEP-51 responses.
@@ -258,6 +262,8 @@ pub struct DhtRuntimeSnapshot {
pub sample_infohashes_response_dropped: u64,
/// New sampled hashes accepted for Peer lookup.
pub sample_infohashes_hashes_discovered: u64,
/// Sampled hashes rejected by the application admission callback.
pub sample_infohashes_hashes_filtered: u64,
/// Sampled hashes rejected by the bounded deduplicator.
pub sample_infohashes_hashes_duplicate: u64,
/// New sampled hashes dropped because Peer lookup ingress was full.
@@ -356,12 +362,15 @@ struct DhtRuntimeStatsInner {
peer_lookup_peers_found: AtomicU64,
peer_lookup_response_dropped: AtomicU64,
peer_lookup_output_dropped: AtomicU64,
peer_lookup_preferred_succeeded: AtomicU64,
peer_lookup_fallbacks: AtomicU64,
sample_infohashes_queries: AtomicU64,
sample_infohashes_responses: AtomicU64,
sample_infohashes_timeouts: AtomicU64,
sample_infohashes_send_failures: AtomicU64,
sample_infohashes_response_dropped: AtomicU64,
sample_infohashes_hashes_discovered: AtomicU64,
sample_infohashes_hashes_filtered: AtomicU64,
sample_infohashes_hashes_duplicate: AtomicU64,
sample_infohashes_hashes_dropped: AtomicU64,
metadata_peer_attempts: AtomicU64,
@@ -466,12 +475,15 @@ impl Default for DhtRuntimeStatsInner {
peer_lookup_peers_found: AtomicU64::new(0),
peer_lookup_response_dropped: AtomicU64::new(0),
peer_lookup_output_dropped: AtomicU64::new(0),
peer_lookup_preferred_succeeded: AtomicU64::new(0),
peer_lookup_fallbacks: AtomicU64::new(0),
sample_infohashes_queries: AtomicU64::new(0),
sample_infohashes_responses: AtomicU64::new(0),
sample_infohashes_timeouts: AtomicU64::new(0),
sample_infohashes_send_failures: AtomicU64::new(0),
sample_infohashes_response_dropped: AtomicU64::new(0),
sample_infohashes_hashes_discovered: AtomicU64::new(0),
sample_infohashes_hashes_filtered: AtomicU64::new(0),
sample_infohashes_hashes_duplicate: AtomicU64::new(0),
sample_infohashes_hashes_dropped: AtomicU64::new(0),
metadata_peer_attempts: AtomicU64::new(0),
@@ -601,6 +613,10 @@ impl DhtRuntimeStats {
.peer_lookup_response_dropped
.load(Ordering::Relaxed),
peer_lookup_output_dropped: inner.peer_lookup_output_dropped.load(Ordering::Relaxed),
peer_lookup_preferred_succeeded: inner
.peer_lookup_preferred_succeeded
.load(Ordering::Relaxed),
peer_lookup_fallbacks: inner.peer_lookup_fallbacks.load(Ordering::Relaxed),
sample_infohashes_queries: inner.sample_infohashes_queries.load(Ordering::Relaxed),
sample_infohashes_responses: inner.sample_infohashes_responses.load(Ordering::Relaxed),
sample_infohashes_timeouts: inner.sample_infohashes_timeouts.load(Ordering::Relaxed),
@@ -613,6 +629,9 @@ impl DhtRuntimeStats {
sample_infohashes_hashes_discovered: inner
.sample_infohashes_hashes_discovered
.load(Ordering::Relaxed),
sample_infohashes_hashes_filtered: inner
.sample_infohashes_hashes_filtered
.load(Ordering::Relaxed),
sample_infohashes_hashes_duplicate: inner
.sample_infohashes_hashes_duplicate
.load(Ordering::Relaxed),
@@ -857,6 +876,18 @@ impl DhtRuntimeStats {
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn peer_lookup_preferred_succeeded(&self) {
self.inner
.peer_lookup_preferred_succeeded
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn peer_lookup_fallback(&self) {
self.inner
.peer_lookup_fallbacks
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn sample_query(&self) {
self.inner
.sample_infohashes_queries
@@ -893,6 +924,12 @@ impl DhtRuntimeStats {
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn sample_hash_filtered(&self, count: usize) {
self.inner
.sample_infohashes_hashes_filtered
.fetch_add(count as u64, Ordering::Relaxed);
}
pub(crate) fn sample_hash_duplicate(&self) {
self.inner
.sample_infohashes_hashes_duplicate
@@ -1231,12 +1268,15 @@ mod tests {
writer.peer_lookup_peer_found();
writer.peer_lookup_response_dropped();
writer.peer_lookup_output_dropped();
writer.peer_lookup_preferred_succeeded();
writer.peer_lookup_fallback();
writer.sample_query();
writer.sample_response();
writer.sample_timeout();
writer.sample_send_failed();
writer.sample_response_dropped();
writer.sample_hash_discovered();
writer.sample_hash_filtered(1);
writer.sample_hash_duplicate();
writer.sample_hash_dropped();
writer.metadata_peer_attempt();
@@ -1298,12 +1338,15 @@ mod tests {
peer_lookup_peers_found: 1,
peer_lookup_response_dropped: 1,
peer_lookup_output_dropped: 1,
peer_lookup_preferred_succeeded: 1,
peer_lookup_fallbacks: 1,
sample_infohashes_queries: 1,
sample_infohashes_responses: 1,
sample_infohashes_timeouts: 1,
sample_infohashes_send_failures: 1,
sample_infohashes_response_dropped: 1,
sample_infohashes_hashes_discovered: 1,
sample_infohashes_hashes_filtered: 1,
sample_infohashes_hashes_duplicate: 1,
sample_infohashes_hashes_dropped: 1,
metadata_peer_attempts: 1,
+153 -19
View File
@@ -1,4 +1,4 @@
use crate::budget::RateBucket;
use crate::budget::{RateBucket, SharedRateBudget};
use crate::crawl_engine::CrawlEngine;
use crate::krpc::{encode_sample_infohashes_query, for_each_response_node};
use crate::node_id::{TransactionId, random_node_id};
@@ -8,11 +8,12 @@ use crate::routing_snapshot::RoutingSnapshot;
use crate::runtime_stats::DhtRuntimeStats;
use crate::types::{NetMode, NodeTuple, SampleInfohashesOptions};
use ahash::{AHashMap, AHashSet};
use arc_swap::ArcSwap;
use arc_swap::{ArcSwap, ArcSwapOption};
use bytes::BytesMut;
#[cfg(feature = "metrics")]
use metrics::{counter, gauge};
use std::collections::VecDeque;
use std::future::Future;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::{Duration, Instant};
@@ -22,6 +23,7 @@ use tokio_util::sync::CancellationToken;
const SAMPLE_TID_TAG: u8 = 0x51;
const RESPONSE_CHANNEL_CAPACITY: usize = 4_096;
const ADMISSION_BATCH_CHANNEL_CAPACITY: usize = 64;
const MAINTENANCE_INTERVAL: Duration = Duration::from_millis(25);
const PRODUCTIVE_REVISIT: Duration = Duration::from_secs(60);
const MAX_PROTOCOL_INTERVAL: Duration = Duration::from_secs(6 * 60 * 60);
@@ -43,6 +45,18 @@ struct SampleResponse {
response: DhtResponse,
}
struct SampleAdmissionBatch {
preferred_node: NodeTuple,
hashes: Vec<[u8; 20]>,
}
pub(crate) type SampleHashAdmissionCallback = Box<
dyn Fn(Vec<[u8; 20]>) -> std::pin::Pin<Box<dyn Future<Output = Vec<[u8; 20]>> + Send>>
+ Send
+ Sync
+ 'static,
>;
#[derive(Clone)]
pub(crate) struct SampleInfohashesHandle {
response_tx: mpsc::Sender<SampleResponse>,
@@ -80,6 +94,8 @@ pub(crate) fn is_sample_infohashes_tid(tid: &TransactionId) -> bool {
pub(crate) struct SampleInfohashesRuntime {
pub(crate) options: SampleInfohashesOptions,
pub(crate) stats: DhtRuntimeStats,
pub(crate) outbound_query_budget: SharedRateBudget,
pub(crate) hash_admission: Arc<ArcSwapOption<SampleHashAdmissionCallback>>,
pub(crate) shutdown: CancellationToken,
}
@@ -95,9 +111,12 @@ pub(crate) fn spawn_sample_infohashes(
let SampleInfohashesRuntime {
options,
stats,
outbound_query_budget,
hash_admission,
shutdown,
} = runtime;
let (response_tx, response_rx) = mpsc::channel(RESPONSE_CHANNEL_CAPACITY);
let (admission_tx, admission_rx) = mpsc::channel(ADMISSION_BATCH_CHANNEL_CAPACITY);
let socket_v4 = sockets
.iter()
.find_map(|(addr, socket)| addr.is_ipv4().then(|| socket.clone()));
@@ -105,6 +124,13 @@ pub(crate) fn spawn_sample_infohashes(
.iter()
.find_map(|(addr, socket)| addr.is_ipv6().then(|| socket.clone()));
let now = Instant::now();
spawn_sample_admission(
admission_rx,
hash_admission,
peer_lookup.request_sender(),
stats.clone(),
shutdown.clone(),
);
let actor = SampleInfohashesActor {
netmode,
local_id,
@@ -112,7 +138,7 @@ pub(crate) fn spawn_sample_infohashes(
socket_v6,
snapshot,
crawl_engine,
peer_lookup,
admission_tx,
response_rx,
query_budget: RateBucket::per_second(
options.max_queries_per_second,
@@ -132,6 +158,7 @@ pub(crate) fn spawn_sample_infohashes(
pending_expiry: VecDeque::new(),
next_tid: 1,
runtime_stats: stats.clone(),
outbound_query_budget,
shutdown,
};
tokio::spawn(actor.run());
@@ -148,7 +175,7 @@ struct SampleInfohashesActor {
socket_v6: Option<Arc<UdpSocket>>,
snapshot: Arc<ArcSwap<RoutingSnapshot>>,
crawl_engine: Arc<CrawlEngine>,
peer_lookup: PeerLookupHandle,
admission_tx: mpsc::Sender<SampleAdmissionBatch>,
response_rx: mpsc::Receiver<SampleResponse>,
query_budget: RateBucket,
max_in_flight: usize,
@@ -163,9 +190,50 @@ struct SampleInfohashesActor {
pending_expiry: VecDeque<(Instant, PendingKey)>,
next_tid: u64,
runtime_stats: DhtRuntimeStats,
outbound_query_budget: SharedRateBudget,
shutdown: CancellationToken,
}
fn spawn_sample_admission(
mut receiver: mpsc::Receiver<SampleAdmissionBatch>,
hash_admission: Arc<ArcSwapOption<SampleHashAdmissionCallback>>,
peer_lookup_tx: mpsc::Sender<PeerLookupRequest>,
runtime_stats: DhtRuntimeStats,
shutdown: CancellationToken,
) {
tokio::spawn(async move {
loop {
let batch = tokio::select! {
_ = shutdown.cancelled() => break,
batch = receiver.recv() => {
let Some(batch) = batch else { break };
batch
}
};
let input_len = batch.hashes.len();
let admitted = match hash_admission.load_full() {
Some(callback) => callback(batch.hashes).await,
None => batch.hashes,
};
runtime_stats.sample_hash_filtered(input_len.saturating_sub(admitted.len()));
for info_hash in admitted {
let request = PeerLookupRequest {
info_hash,
preferred_node: Some(batch.preferred_node),
};
let sent = tokio::select! {
_ = shutdown.cancelled() => false,
result = peer_lookup_tx.send(request) => result.is_ok(),
};
if !sent {
return;
}
runtime_stats.sample_hash_discovered();
}
}
});
}
impl SampleInfohashesActor {
async fn run(mut self) {
let mut maintenance = tokio::time::interval(MAINTENANCE_INTERVAL);
@@ -188,7 +256,10 @@ impl SampleInfohashesActor {
}
async fn dispatch(&mut self, now: Instant) {
if self.max_in_flight == 0 || self.pending.len() >= self.max_in_flight {
if self.max_in_flight == 0
|| self.pending.len() >= self.max_in_flight
|| self.admission_tx.capacity() == 0
{
return;
}
let available = self.max_in_flight.saturating_sub(self.pending.len());
@@ -236,10 +307,14 @@ impl SampleInfohashesActor {
let Some(socket) = socket else {
return false;
};
if !self.outbound_query_budget.try_take_one(now) {
return false;
}
let tid = self.next_transaction_id();
let mut buffer = BytesMut::with_capacity(128);
encode_sample_infohashes_query(&mut buffer, &tid, &random_node_id(), &self.local_id);
if socket.send_to(&buffer, node.addr).await.is_err() {
self.outbound_query_budget.refund_one();
self.runtime_stats.sample_send_failed();
self.next_allowed
.insert(node.addr, now + self.unsupported_backoff);
@@ -288,7 +363,8 @@ impl SampleInfohashesActor {
id: responder_id,
addr: event.remote_addr,
};
let mut discovered = 0usize;
let mut hashes = Vec::new();
let mut batch_hashes = AHashSet::new();
if let Some(samples) = event.response.samples.as_deref()
&& samples.len() % 20 == 0
{
@@ -298,23 +374,41 @@ impl SampleInfohashesActor {
self.runtime_stats.sample_hash_duplicate();
continue;
}
let request = PeerLookupRequest {
info_hash: hash,
preferred_node: Some(preferred_node),
};
if self.peer_lookup.request_sender().try_send(request).is_ok() {
self.remember_hash(hash);
discovered += 1;
self.runtime_stats.sample_hash_discovered();
} else {
self.runtime_stats.sample_hash_dropped();
#[cfg(feature = "metrics")]
counter!("dht_sample_infohashes_dropped_total", "reason" => "peer_lookup_queue_full")
.increment(1);
if batch_hashes.insert(hash) {
hashes.push(hash);
}
}
}
let candidate_count = hashes.len();
let admitted_to_triage = if hashes.is_empty() {
false
} else {
let hashes_to_remember = hashes.clone();
if self
.admission_tx
.try_send(SampleAdmissionBatch {
preferred_node,
hashes,
})
.is_ok()
{
for hash in hashes_to_remember {
self.remember_hash(hash);
}
true
} else {
for _ in 0..candidate_count {
self.runtime_stats.sample_hash_dropped();
}
#[cfg(feature = "metrics")]
counter!("dht_sample_infohashes_dropped_total", "reason" => "admission_queue_full")
.increment(candidate_count as u64);
false
}
};
let discovered = usize::from(admitted_to_triage) * candidate_count;
let protocol_interval = Duration::from_secs(event.response.interval.unwrap_or(300))
.clamp(Duration::from_secs(10), MAX_PROTOCOL_INTERVAL);
let delay = if discovered > 0 {
@@ -407,4 +501,44 @@ mod tests {
assert_eq!(response.num, Some(2));
assert_eq!(response.samples.unwrap().len(), 40);
}
#[tokio::test]
async fn batch_admission_only_forwards_application_approved_hashes() {
let (batch_tx, batch_rx) = mpsc::channel(1);
let (lookup_tx, mut lookup_rx) = mpsc::channel(2);
let admission = Arc::new(ArcSwapOption::empty());
let callback: Arc<SampleHashAdmissionCallback> = Arc::new(Box::new(|hashes| {
Box::pin(async move { hashes.into_iter().filter(|hash| *hash == [2; 20]).collect() })
}));
admission.store(Some(callback));
let stats = DhtRuntimeStats::default();
let shutdown = CancellationToken::new();
spawn_sample_admission(
batch_rx,
admission,
lookup_tx,
stats.clone(),
shutdown.clone(),
);
batch_tx
.send(SampleAdmissionBatch {
preferred_node: NodeTuple {
id: [9; 20],
addr: "127.0.0.1:6881".parse().unwrap(),
},
hashes: vec![[1; 20], [2; 20]],
})
.await
.unwrap();
let request = tokio::time::timeout(Duration::from_secs(1), lookup_rx.recv())
.await
.unwrap()
.unwrap();
assert_eq!(request.info_hash, [2; 20]);
let snapshot = stats.snapshot();
assert_eq!(snapshot.sample_infohashes_hashes_filtered, 1);
assert_eq!(snapshot.sample_infohashes_hashes_discovered, 1);
shutdown.cancel();
}
}
+30 -3
View File
@@ -1,5 +1,5 @@
use crate::addr::is_valid_node_addr;
use crate::budget::RateBucket;
use crate::budget::{RateBucket, SharedRateBudget};
use crate::crawl_config::ResolvedCrawlConfig;
use crate::crawl_engine::CrawlEngine;
use crate::error::Result;
@@ -12,8 +12,8 @@ use crate::peer_lookup::{
use crate::protocol::{DhtArgs, DhtMessage};
use crate::runtime_stats::{DhtRuntimeLimits, DhtRuntimeStats};
use crate::sample_infohashes::{
SampleInfohashesHandle, SampleInfohashesRuntime, is_sample_infohashes_tid,
spawn_sample_infohashes,
SampleHashAdmissionCallback, SampleInfohashesHandle, SampleInfohashesRuntime,
is_sample_infohashes_tid, spawn_sample_infohashes,
};
use crate::scheduler::{
MetadataCompletionCallback, MetadataFetchCallback, MetadataScheduler,
@@ -257,6 +257,7 @@ pub struct DHTServer {
torrent_callback: Arc<ArcSwapOption<TorrentAckCallback>>,
hash_filter: Arc<ArcSwapOption<FilterCallback>>,
on_metadata_fetch: Arc<ArcSwapOption<MetadataFetchCallback>>,
sample_hash_admission: Arc<ArcSwapOption<SampleHashAdmissionCallback>>,
metadata_completion_callback: Arc<ArcSwapOption<MetadataCompletionCallback>>,
error_callback: Arc<ArcSwapOption<ErrorCallback>>,
crawl_engine: Arc<CrawlEngine>,
@@ -353,18 +354,26 @@ impl DHTServer {
mpsc::channel::<HashDiscovered>(options.hash_queue_capacity);
let fetcher = Arc::new(RbitFetcher::new_with_runtime_stats(
options.metadata.timeout_secs,
options.metadata.max_connects_per_second,
options.metadata.peer_failure_cache_capacity,
options.metadata.peer_failure_ttl_secs,
runtime_stats.clone(),
));
let torrent_callback = Arc::new(ArcSwapOption::empty());
let on_metadata_fetch = Arc::new(ArcSwapOption::empty());
let sample_hash_admission = Arc::new(ArcSwapOption::empty());
let metadata_completion_callback = Arc::new(ArcSwapOption::empty());
let metadata_queue_len = Arc::new(AtomicUsize::new(0));
let shutdown = CancellationToken::new();
let outbound_query_budget = SharedRateBudget::per_second(
options.max_outbound_queries_per_second,
options.outbound_query_burst,
true,
);
let crawl_engine = Arc::new(CrawlEngine::new(
crawl_config.clone(),
runtime_stats.clone(),
outbound_query_budget.clone(),
));
let peer_lookup = spawn_peer_lookup(
options.netmode,
@@ -375,6 +384,7 @@ impl DHTServer {
PeerLookupRuntime {
options: options.peer_lookup.clone(),
stats: runtime_stats.clone(),
outbound_query_budget: outbound_query_budget.clone(),
shutdown: shutdown.clone(),
},
);
@@ -388,6 +398,8 @@ impl DHTServer {
SampleInfohashesRuntime {
options: options.sample_infohashes.clone(),
stats: runtime_stats.clone(),
outbound_query_budget,
hash_admission: sample_hash_admission.clone(),
shutdown: shutdown.clone(),
},
);
@@ -423,6 +435,7 @@ impl DHTServer {
torrent_callback,
hash_filter: Arc::new(ArcSwapOption::empty()),
on_metadata_fetch,
sample_hash_admission,
metadata_completion_callback,
error_callback: Arc::new(ArcSwapOption::empty()),
crawl_engine,
@@ -450,6 +463,20 @@ impl DHTServer {
self.on_metadata_fetch.store(Some(callback));
}
/// Registers an asynchronous batch admission callback for BEP-51 sampled InfoHashes.
///
/// Only hashes returned by the callback proceed to active Peer lookup. The callback should
/// preserve hashes it cannot classify so transient application errors fail open.
pub fn on_sampled_hashes<F, Fut>(&self, callback: F)
where
F: Fn(Vec<[u8; 20]>) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Vec<[u8; 20]>> + Send + 'static,
{
let callback: Arc<SampleHashAdmissionCallback> =
Arc::new(Box::new(move |hashes| Box::pin(callback(hashes))));
self.sample_hash_admission.store(Some(callback));
}
/// Registers a torrent callback whose return is implicitly treated as accepted delivery.
///
/// Registering a new torrent callback replaces the previous one.
+20 -11
View File
@@ -115,6 +115,10 @@ pub struct DHTOptions {
pub netmode: NetMode,
/// Capacity between announce processing and the Metadata scheduler.
pub hash_queue_capacity: usize,
/// Maximum combined active find_node get_peers and sample_infohashes queries per second.
pub max_outbound_queries_per_second: u32,
/// Maximum shared outbound query budget consumed immediately after an idle period.
pub outbound_query_burst: u32,
/// Metadata download and Peer-cache limits.
pub metadata: MetadataOptions,
/// Active get_peers lookup rate and concurrency limits.
@@ -134,6 +138,8 @@ pub struct MetadataOptions {
pub max_queue_size: usize,
/// Maximum number of concurrent Metadata jobs.
pub max_worker_count: usize,
/// Maximum real TCP connection attempts started per second.
pub max_connects_per_second: u32,
/// Maximum number of cached bad Peer socket addresses.
pub peer_failure_cache_capacity: usize,
/// Timeout/connect failure cache lifetime in seconds.
@@ -274,6 +280,8 @@ impl Default for DHTOptions {
port: 6881,
netmode: NetMode::Ipv4Only,
hash_queue_capacity: 10_000,
max_outbound_queries_per_second: 10,
outbound_query_burst: 2,
metadata: MetadataOptions::default(),
peer_lookup: PeerLookupOptions::default(),
sample_infohashes: SampleInfohashesOptions::default(),
@@ -287,7 +295,8 @@ impl Default for MetadataOptions {
Self {
timeout_secs: 4,
max_queue_size: 10_000,
max_worker_count: 256,
max_worker_count: 8,
max_connects_per_second: 2,
peer_failure_cache_capacity: 200_000,
peer_failure_ttl_secs: 60,
}
@@ -297,9 +306,9 @@ impl Default for MetadataOptions {
impl Default for PeerLookupOptions {
fn default() -> Self {
Self {
max_lookups_per_second: 128,
burst: 128,
max_active_lookups: 256,
max_lookups_per_second: 1,
burst: 1,
max_active_lookups: 4,
}
}
}
@@ -307,9 +316,9 @@ impl Default for PeerLookupOptions {
impl Default for SampleInfohashesOptions {
fn default() -> Self {
Self {
max_queries_per_second: 16,
burst: 16,
max_in_flight: 64,
max_queries_per_second: 1,
burst: 1,
max_in_flight: 4,
request_timeout_millis: 1_500,
unsupported_backoff_secs: 300,
dedup_capacity: 1_000_000,
@@ -320,11 +329,11 @@ impl Default for SampleInfohashesOptions {
impl Default for RateLimitOptions {
fn default() -> Self {
Self {
max_find_node_rate_per_sec: 200,
burst: 40,
max_in_flight: 512,
max_find_node_rate_per_sec: 6,
burst: 2,
max_in_flight: 12,
request_timeout_secs: 2,
max_new_destinations_per_minute: 10_000,
max_new_destinations_per_minute: 60,
max_response_rate_per_sec: 500,
max_response_bytes_per_sec: 1024 * 1024,
max_response_rate_per_source: 40,
+27
View File
@@ -0,0 +1,27 @@
# 定义 dht-search 的推荐起始配置并作为用户配置模板
data_dir = "data"
persistence_queue_capacity = 4096
stats_interval_secs = 10
index_batch_size = 512
index_interval_millis = 5000
[dht]
port = 12313
netmode = "ipv4-only"
hash_queue_capacity = 10000
max_outbound_queries_per_second = 10
outbound_query_burst = 2
metadata_timeout_secs = 4
metadata_queue_capacity = 10000
metadata_workers = 8
metadata_connects_per_second = 2
sample_queries_per_second = 1
peer_lookups_per_second = 1
peer_lookup_max_active = 4
find_node_queries_per_second = 6
find_node_max_in_flight = 12
new_destinations_per_minute = 60
[http]
listen = "127.0.0.1:8080"
+10 -2
View File
@@ -10,19 +10,27 @@ repository.workspace = true
publish = false
[features]
default = []
default = ["rocksdb-storage"]
rocksdb-storage = ["dep:rocksdb"]
[dependencies]
axum = "0.8.9"
blake3 = "1.8.5"
clap = { version = "4.5", features = ["derive"] }
dht-crawler = { path = "../dht-crawler", features = ["metrics"] }
rocksdb = { version = "0.24.0", optional = true }
hex = "0.4"
rocksdb = { version = "0.24.0", default-features = false, features = ["bindgen-runtime", "lz4"], optional = true }
rmp-serde = "1.3"
serde.workspace = true
serde_json = "1.0"
tantivy = "0.26.1"
thiserror.workspace = true
tokio = { workspace = true, features = ["macros", "rt-multi-thread", "signal", "sync", "time"] }
tokio-util.workspace = true
toml = "0.9"
tracing.workspace = true
tracing-subscriber = { workspace = true, features = ["env-filter", "fmt"] }
unicode-normalization = "0.1"
[dev-dependencies]
tempfile = "3.27"
+109
View File
@@ -0,0 +1,109 @@
# dht-search
`dht-search` 是集 DHT Metadata 采集 RocksDB 持久化 Tantivy 搜索索引和 HTTP API 于一体的应用
## 构建准备
RocksDB 包含 C++ 代码并在构建时使用 bindgen 因此需要 libclang
Windows 可以把 libclang 安装到工作区本地目录
```powershell
python -m pip install --target .tools\libclang libclang
$env:LIBCLANG_PATH = "$PWD\.tools\libclang\clang\native"
cargo build -p dht-search --release
```
`.tools` 只用于本地构建不会部署到运行设备
## 配置
复制根目录的配置模板
```powershell
Copy-Item dht-search.example.toml dht-search.toml
```
相对 `data_dir` 以配置文件所在目录为基准解析
也可以通过命令行覆盖数据目录和本次运行时长
```powershell
cargo run -p dht-search -- --data-dir D:\data\dht-search --run-duration-secs 3600
```
不设置 `run-duration-secs` 时服务持续运行直到收到 Ctrl+C SIGINT 或 SIGTERM
### 网络保护配置
主动 DHT 查询共享 `max_outbound_queries_per_second` 总预算,因此 `find_node` `get_peers``sample_infohashes` 的总发送速率不会各自叠加后失控
| 配置项 | 保守默认值 | 作用 |
|---|---:|---|
| `max_outbound_queries_per_second` | `10` | 三类主动 DHT UDP 查询的合计每秒速率 |
| `outbound_query_burst` | `2` | 空闲后允许立即消费的 UDP 查询数 |
| `find_node_queries_per_second` | `6` | `find_node` 自身速率上限 |
| `find_node_max_in_flight` | `12` | 同时等待响应的 `find_node` 数量 |
| `new_destinations_per_minute` | `60` | 每分钟首次探测的新 UDP 目标数量 |
| `peer_lookups_per_second` | `1` | 每秒启动的 infohash Peer 查找数量 |
| `peer_lookup_max_active` | `4` | 同时运行的 Peer 查找数量 |
| `sample_queries_per_second` | `1` | BEP-51 采样查询速率 |
| `metadata_workers` | `8` | 同时处理的 Metadata 任务数量 |
| `metadata_connects_per_second` | `2` | 每秒真正开始的 Peer TCP 连接数量 |
桌面网络不要在不了解路由器 NAT 和代理容量时大幅提高这些值
### 采样去重和 Peer 查找
BEP-51 返回的 infohash 会先进入有界批量准入队列并由 RocksDB 精确判断
已有 infohash 只更新最后发现时间和发现次数不会再次执行 Peer Lookup
未知 infohash 首先只向返回样本的 DHT 节点查询一次 `get_peers` 只有单点查询没有返回 Peer 时才降级为有限迭代查找
`/stats` 中的 `sampled_hashes_filtered` `peer_lookup_preferred_succeeded``peer_lookup_fallbacks` 用于观察提前去重和单点优先效果
## Linux 资源限制
生产环境仍可能同时使用较多 TCP socket
Linux 生产运行必须把文件描述符上限提高到至少 65536
```bash
prlimit --nofile=65536:65536 -- \
/opt/dht-search/dht-search --config /opt/dht-search/dht-search.toml
```
systemd 服务需要设置
```ini
[Service]
LimitNOFILE=65536
```
文件描述符上限过低时 DHT 连接会挤占 Tantivy 和 RocksDB 打开文件所需的描述符并导致服务安全停止
## API
默认只监听 `127.0.0.1:8080`
```text
GET /health
GET /ready
GET /stats
GET /search?q=ubuntu&offset=0&limit=20
GET /search?q=&min_size=1048576&max_size=10737418240&extension=mkv
GET /torrents/{infohash}
```
`limit` 被限制在 1 到 100 之间且 `offset` 最大为 10000
搜索和大小扩展名过滤由 Tantivy 索引执行不会把全部记录加载到内存过滤
## 数据恢复
RocksDB 是权威数据源而 Tantivy 是可重建索引
当 Tantivy 目录不存在时应用会把全部 RocksDB 记录重新标记为待索引并自动完成全量重建
正常退出会先停止 DHT 再排空持久化队列提交剩余索引最后关闭 HTTP 服务
+141
View File
@@ -1 +1,142 @@
// 负责处理搜索详情统计和健康检查请求
use std::str::FromStr;
use axum::{
Json,
extract::{Path, Query, State},
http::StatusCode,
response::{IntoResponse, Response},
};
use dht_search::{
domain::InfoHash,
search::{SearchOptions, SearchPage},
};
use super::{
ApiState,
request::SearchRequest,
response::{ErrorResponse, StatsResponse, StatusResponse, TorrentResponse},
};
pub(crate) async fn health() -> Json<StatusResponse> {
Json(StatusResponse { status: "ok" })
}
pub(crate) async fn ready() -> Json<StatusResponse> {
Json(StatusResponse { status: "ready" })
}
pub(crate) async fn stats(State(state): State<ApiState>) -> Json<StatsResponse> {
let dht = state.dht_stats.snapshot();
let observability = state.dht_stats.observability_snapshot();
let persistence = state.persistence.snapshot();
Json(StatsResponse {
nodes: dht.node_pool_size,
udp_tx_packets: observability.udp_tx_packets,
find_node_queries: dht
.queries_new
.saturating_add(dht.queries_revisit)
.saturating_add(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,
metadata_peer_attempts: dht.metadata_peer_attempts,
metadata_in_flight: dht.metadata_in_flight,
metadata_ok: dht.metadata_peer_succeeded,
metadata_failed: dht.metadata_peer_failed,
persistence_accepted: persistence.accepted,
persistence_inserted: persistence.inserted,
persistence_updated: persistence.updated,
persistence_rejected_full: persistence.rejected_full,
persistence_queue: persistence.queue_depth,
indexed_documents: state.search.num_docs(),
})
}
pub(crate) async fn search(
State(state): State<ApiState>,
Query(request): Query<SearchRequest>,
) -> Result<Json<SearchPage>, ApiError> {
if request.q.len() > 512 {
return Err(ApiError::bad_request("查询文本不能超过 512 字节"));
}
if request
.min_size
.zip(request.max_size)
.is_some_and(|(min, max)| min > max)
{
return Err(ApiError::bad_request("min_size 不能大于 max_size"));
}
let page = tokio::task::spawn_blocking(move || {
state.search.search_with(SearchOptions {
query: request.q,
offset: request.offset,
limit: request.limit,
min_size: request.min_size,
max_size: request.max_size,
extension: request.extension,
})
})
.await
.map_err(|error| ApiError::internal(error.to_string()))?
.map_err(|error| ApiError::bad_request(error.to_string()))?;
Ok(Json(page))
}
pub(crate) async fn torrent(
State(state): State<ApiState>,
Path(info_hash): Path<String>,
) -> Result<Json<TorrentResponse>, ApiError> {
let info_hash =
InfoHash::from_str(&info_hash).map_err(|error| ApiError::bad_request(error.to_string()))?;
let record = tokio::task::spawn_blocking(move || state.repository.get(info_hash))
.await
.map_err(|error| ApiError::internal(error.to_string()))?
.map_err(|error| ApiError::internal(error.to_string()))?
.ok_or_else(|| ApiError::not_found("没有找到该 infohash"))?;
Ok(Json(record.into()))
}
pub(crate) struct ApiError {
status: StatusCode,
message: String,
}
impl ApiError {
fn bad_request(message: impl Into<String>) -> Self {
Self {
status: StatusCode::BAD_REQUEST,
message: message.into(),
}
}
fn not_found(message: impl Into<String>) -> Self {
Self {
status: StatusCode::NOT_FOUND,
message: message.into(),
}
}
fn internal(message: impl Into<String>) -> Self {
Self {
status: StatusCode::INTERNAL_SERVER_ERROR,
message: message.into(),
}
}
}
impl IntoResponse for ApiError {
fn into_response(self) -> Response {
(
self.status,
Json(ErrorResponse {
error: self.message,
}),
)
.into_response()
}
}
+36
View File
@@ -3,3 +3,39 @@
mod handlers;
mod request;
mod response;
use std::{net::SocketAddr, sync::Arc};
use axum::{Router, routing::get};
use dht_crawler::DhtRuntimeStats;
use dht_search::{search::SearchEngine, storage::TorrentRepository};
use tokio_util::sync::CancellationToken;
use crate::crawler::pipeline::PersistenceIngress;
#[derive(Clone)]
pub(crate) struct ApiState {
pub(crate) repository: Arc<dyn TorrentRepository>,
pub(crate) search: SearchEngine,
pub(crate) dht_stats: DhtRuntimeStats,
pub(crate) persistence: PersistenceIngress,
}
pub(crate) async fn serve(
listen: SocketAddr,
state: ApiState,
cancel: CancellationToken,
) -> std::io::Result<()> {
let router = Router::new()
.route("/health", get(handlers::health))
.route("/ready", get(handlers::ready))
.route("/stats", get(handlers::stats))
.route("/search", get(handlers::search))
.route("/torrents/{info_hash}", get(handlers::torrent))
.with_state(state);
let listener = tokio::net::TcpListener::bind(listen).await?;
tracing::info!(%listen, "HTTP 服务启动");
axum::serve(listener, router)
.with_graceful_shutdown(cancel.cancelled_owned())
.await
}
+19
View File
@@ -1 +1,20 @@
// 负责定义 HTTP 查询参数和输入校验模型
use serde::Deserialize;
fn default_limit() -> usize {
20
}
#[derive(Debug, Deserialize)]
pub(crate) struct SearchRequest {
#[serde(default)]
pub(crate) q: String,
#[serde(default)]
pub(crate) offset: usize,
#[serde(default = "default_limit")]
pub(crate) limit: usize,
pub(crate) min_size: Option<u64>,
pub(crate) max_size: Option<u64>,
pub(crate) extension: Option<String>,
}
+68
View File
@@ -1 +1,69 @@
// 负责定义稳定的 HTTP 响应模型和领域对象转换边界
use dht_search::domain::{TorrentFile, TorrentRecord};
use serde::Serialize;
#[derive(Debug, Serialize)]
pub(crate) struct StatusResponse {
pub(crate) status: &'static str,
}
#[derive(Debug, Serialize)]
pub(crate) struct ErrorResponse {
pub(crate) error: String,
}
#[derive(Debug, Serialize)]
pub(crate) struct StatsResponse {
pub(crate) nodes: usize,
pub(crate) udp_tx_packets: u64,
pub(crate) find_node_queries: u64,
pub(crate) peer_lookup_queries: u64,
pub(crate) peer_lookup_preferred_succeeded: u64,
pub(crate) peer_lookup_fallbacks: u64,
pub(crate) sample_queries: u64,
pub(crate) sampled_hashes: u64,
pub(crate) sampled_hashes_filtered: u64,
pub(crate) metadata_peer_attempts: u64,
pub(crate) metadata_in_flight: usize,
pub(crate) metadata_ok: u64,
pub(crate) metadata_failed: u64,
pub(crate) persistence_accepted: u64,
pub(crate) persistence_inserted: u64,
pub(crate) persistence_updated: u64,
pub(crate) persistence_rejected_full: u64,
pub(crate) persistence_queue: usize,
pub(crate) indexed_documents: u64,
}
#[derive(Debug, Serialize)]
pub(crate) struct TorrentResponse {
pub(crate) info_hash: String,
pub(crate) magnet_link: String,
pub(crate) name: String,
pub(crate) total_size: u64,
pub(crate) files: Vec<TorrentFile>,
pub(crate) piece_length: u64,
pub(crate) content_key: String,
pub(crate) first_seen: u64,
pub(crate) last_seen: u64,
pub(crate) seen_count: u64,
}
impl From<TorrentRecord> for TorrentResponse {
fn from(record: TorrentRecord) -> Self {
let info_hash = record.info_hash.to_string();
Self {
magnet_link: format!("magnet:?xt=urn:btih:{info_hash}"),
info_hash,
name: record.name,
total_size: record.total_size,
files: record.files,
piece_length: record.piece_length,
content_key: hex::encode(record.content_key),
first_seen: record.first_seen,
last_seen: record.last_seen,
seen_count: record.seen_count,
}
}
}
+311
View File
@@ -1 +1,312 @@
// 负责连接采集存储索引和接口层并定义应用级启动顺序
use std::{
str::FromStr,
sync::Arc,
time::{Duration, SystemTime, UNIX_EPOCH},
};
use dht_crawler::DHTServer;
use dht_search::{
domain::InfoHash,
search::SearchEngine,
storage::{RocksTorrentRepository, TorrentRepository},
};
use tokio_util::sync::CancellationToken;
use crate::{
api::{self, ApiState},
config::AppConfig,
crawler::pipeline::PersistencePipeline,
error::AppError,
shutdown,
};
pub(crate) async fn run(config: AppConfig) -> Result<(), AppError> {
std::fs::create_dir_all(&config.data_dir)?;
let database_path = config.data_dir.join("rocksdb");
let repository = Arc::new(RocksTorrentRepository::open(&database_path)?);
let (search, search_created) = SearchEngine::open_with_status(config.data_dir.join("tantivy"))?;
if search_created {
let records = repository.prepare_full_reindex()?;
tracing::info!(records, "检测到新搜索索引并准备全量重建");
}
let repository_api: Arc<dyn TorrentRepository> = repository.clone();
let mut persistence =
PersistencePipeline::start(repository_api, config.persistence_queue_capacity);
let ingress = persistence.ingress.clone();
let options = config.dht_options();
let server = DHTServer::new(options.clone()).await?;
server.on_error(|error| tracing::error!(%error, "DHT 运行时错误"));
let sampled_repository = repository.clone();
server.on_sampled_hashes(move |hashes| {
let repository = sampled_repository.clone();
async move {
let fallback = hashes.clone();
let info_hashes: Vec<_> = hashes.into_iter().map(InfoHash::from_bytes).collect();
let result = tokio::task::spawn_blocking(move || {
repository.filter_unknown_and_observe(&info_hashes, unix_timestamp())
})
.await;
match result {
Ok(Ok(unknown)) => unknown
.into_iter()
.map(|info_hash| *info_hash.as_bytes())
.collect(),
Ok(Err(error)) => {
tracing::error!(%error, "采样 infohash 批量持久化去重失败");
fallback
}
Err(error) => {
tracing::error!(%error, "采样 infohash 批量去重任务异常");
fallback
}
}
}
});
let gate_repository = repository.clone();
server.on_metadata_fetch(move |hash| {
let repository = gate_repository.clone();
async move {
let Ok(info_hash) = InfoHash::from_str(&hash) else {
tracing::warn!(%hash, "DHT 提供了无效 infohash");
return false;
};
let result = tokio::task::spawn_blocking(move || {
repository.observe_existing(info_hash, unix_timestamp())
})
.await;
match result {
Ok(Ok(already_exists)) => !already_exists,
Ok(Err(error)) => {
tracing::error!(%error, %hash, "持久化去重查询失败");
true
}
Err(error) => {
tracing::error!(%error, %hash, "持久化去重任务异常");
true
}
}
}
});
let callback_ingress = ingress.clone();
server.on_torrent_with_ack(move |torrent| callback_ingress.try_enqueue(torrent));
server.on_metadata_fetch_complete(|completion| {
tracing::debug!(
info_hash = %completion.info_hash,
status = ?completion.status,
attempts = completion.attempts,
"Metadata 任务完成"
)
});
tracing::info!(
dht_port = options.port,
data_dir = %config.data_dir.display(),
persistence_queue_capacity = config.persistence_queue_capacity,
"dht-search 启动"
);
let monitor_cancel = CancellationToken::new();
let monitor = tokio::spawn(monitor(
server.clone(),
ingress,
config.stats_interval_secs,
monitor_cancel.clone(),
));
let index_cancel = CancellationToken::new();
let (index_fatal_tx, mut index_fatal) = tokio::sync::oneshot::channel();
let index_task = tokio::spawn(run_indexer(
repository.clone(),
search.clone(),
config.index_batch_size,
Duration::from_millis(config.index_interval_millis),
index_cancel.clone(),
index_fatal_tx,
));
let api_cancel = CancellationToken::new();
let mut api_task = tokio::spawn(api::serve(
config.http.listen,
ApiState {
repository: repository.clone(),
search,
dht_stats: server.runtime_stats(),
persistence: persistence.ingress.clone(),
},
api_cancel.clone(),
));
let run_duration = async {
match config.run_duration_secs {
Some(seconds) => tokio::time::sleep(Duration::from_secs(seconds)).await,
None => std::future::pending().await,
}
};
tokio::pin!(run_duration);
let run_result = tokio::select! {
result = server.start() => result.map_err(AppError::from),
_ = shutdown::signal() => {
tracing::info!("收到退出信号");
Ok(())
}
_ = &mut run_duration => {
tracing::info!("达到配置的运行时长");
Ok(())
}
fatal = &mut persistence.fatal => {
let message = fatal.unwrap_or_else(|_| "持久化 worker 意外停止".to_owned());
Err(AppError::PersistenceWorker(message))
}
fatal = &mut index_fatal => {
let message = fatal.unwrap_or_else(|_| "索引 worker 意外停止".to_owned());
Err(AppError::IndexWorker(message))
}
result = &mut api_task => {
match result {
Ok(Ok(())) => Err(AppError::Config("HTTP 服务意外停止".to_owned())),
Ok(Err(error)) => Err(AppError::Io(error)),
Err(error) => Err(AppError::Config(format!("HTTP 服务任务异常: {error}"))),
}
}
};
server.shutdown();
monitor_cancel.cancel();
let _ = monitor.await;
persistence.close_and_join().await?;
index_cancel.cancel();
match index_task.await {
Ok(Ok(())) => {}
Ok(Err(error)) => return Err(AppError::IndexWorker(error)),
Err(error) => return Err(AppError::IndexWorker(error.to_string())),
}
api_cancel.cancel();
if !api_task.is_finished() {
match api_task.await {
Ok(Ok(())) => {}
Ok(Err(error)) => return Err(AppError::Io(error)),
Err(error) => return Err(AppError::Config(format!("HTTP 服务任务异常: {error}"))),
}
}
tracing::info!("dht-search 已安全停止");
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);
loop {
tokio::select! {
_ = cancel.cancelled() => break,
_ = ticker.tick() => {
let indexed = index_one_batch(repository.clone(), search.clone(), batch_size).await;
match indexed {
Ok(count) if count > 0 => tracing::debug!(count, "搜索索引已提交"),
Ok(_) => {}
Err(error) => {
let _ = fatal.send(error.clone());
return Err(error);
}
}
}
}
}
loop {
let count = index_one_batch(repository.clone(), search.clone(), batch_size).await?;
if count == 0 {
break;
}
}
Ok(())
}
async fn index_one_batch(
repository: Arc<RocksTorrentRepository>,
search: SearchEngine,
batch_size: usize,
) -> Result<usize, String> {
tokio::task::spawn_blocking(move || {
search
.index_pending(repository.as_ref(), batch_size, unix_timestamp())
.map_err(|error| error.to_string())
})
.await
.map_err(|error| error.to_string())?
}
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,
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)
.unwrap_or_default()
.as_secs()
}
+304
View File
@@ -1 +1,305 @@
// 负责加载校验和提供应用配置但不执行任何业务逻辑
use std::{fs, net::SocketAddr, path::PathBuf};
use clap::Parser;
use dht_crawler::{
BootstrapOptions, CrawlOptions, DHTOptions, MetadataOptions, NetMode, PeerLookupOptions,
PoolOptions, RateLimitOptions, SampleInfohashesOptions, SchedulerOptions, TargetOptions,
};
use serde::Deserialize;
use crate::error::AppError;
#[derive(Debug, Parser)]
#[command(name = "dht-search", version, about = "DHT 元数据采集和搜索服务")]
pub(crate) struct Cli {
#[arg(long, default_value = "dht-search.toml")]
config: PathBuf,
#[arg(long)]
data_dir: Option<PathBuf>,
#[arg(long)]
run_duration_secs: Option<u64>,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(default, deny_unknown_fields)]
pub(crate) struct AppConfig {
pub(crate) data_dir: PathBuf,
pub(crate) persistence_queue_capacity: usize,
pub(crate) stats_interval_secs: u64,
pub(crate) run_duration_secs: Option<u64>,
pub(crate) index_batch_size: usize,
pub(crate) index_interval_millis: u64,
pub(crate) dht: DhtConfig,
pub(crate) http: HttpConfig,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(default, deny_unknown_fields)]
pub(crate) struct DhtConfig {
pub(crate) port: u16,
pub(crate) netmode: NetworkMode,
pub(crate) hash_queue_capacity: usize,
pub(crate) max_outbound_queries_per_second: u32,
pub(crate) outbound_query_burst: u32,
pub(crate) metadata_timeout_secs: u64,
pub(crate) metadata_queue_capacity: usize,
pub(crate) metadata_workers: usize,
pub(crate) metadata_connects_per_second: u32,
pub(crate) sample_queries_per_second: u32,
pub(crate) peer_lookups_per_second: u32,
pub(crate) peer_lookup_max_active: usize,
pub(crate) find_node_queries_per_second: u32,
pub(crate) find_node_max_in_flight: usize,
pub(crate) new_destinations_per_minute: u32,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(default, deny_unknown_fields)]
pub(crate) struct HttpConfig {
pub(crate) listen: SocketAddr,
}
#[derive(Debug, Clone, Copy, Default, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub(crate) enum NetworkMode {
#[default]
Ipv4Only,
Ipv6Only,
DualStack,
}
impl Cli {
pub(crate) fn load(self) -> Result<AppConfig, AppError> {
let config_path = absolute_path(&self.config)?;
let mut config = if config_path.exists() {
let contents = fs::read_to_string(&config_path)?;
toml::from_str::<AppConfig>(&contents)?
} else {
AppConfig::default()
};
if let Some(data_dir) = self.data_dir {
config.data_dir = data_dir;
}
if self.run_duration_secs.is_some() {
config.run_duration_secs = self.run_duration_secs;
}
let base = config_path
.parent()
.ok_or_else(|| AppError::Config("配置文件没有父目录".to_owned()))?;
if config.data_dir.is_relative() {
config.data_dir = base.join(&config.data_dir);
}
config.data_dir = normalize_absolute(config.data_dir)?;
config.validate()?;
Ok(config)
}
}
impl AppConfig {
pub(crate) fn dht_options(&self) -> DHTOptions {
let defaults = DHTOptions::default();
DHTOptions {
port: self.dht.port,
netmode: self.dht.netmode.into(),
hash_queue_capacity: self.dht.hash_queue_capacity,
max_outbound_queries_per_second: self.dht.max_outbound_queries_per_second,
outbound_query_burst: self.dht.outbound_query_burst,
metadata: MetadataOptions {
timeout_secs: self.dht.metadata_timeout_secs,
max_queue_size: self.dht.metadata_queue_capacity,
max_worker_count: self.dht.metadata_workers,
max_connects_per_second: self.dht.metadata_connects_per_second,
..defaults.metadata
},
peer_lookup: PeerLookupOptions {
max_lookups_per_second: self.dht.peer_lookups_per_second,
burst: self.dht.peer_lookups_per_second.max(1),
max_active_lookups: self.dht.peer_lookup_max_active,
},
sample_infohashes: SampleInfohashesOptions {
max_queries_per_second: self.dht.sample_queries_per_second,
..defaults.sample_infohashes
},
crawl: CrawlOptions {
pool: PoolOptions {
..defaults.crawl.pool
},
rate_limit: RateLimitOptions {
max_find_node_rate_per_sec: self.dht.find_node_queries_per_second,
burst: self.dht.outbound_query_burst,
max_in_flight: self.dht.find_node_max_in_flight,
max_new_destinations_per_minute: self.dht.new_destinations_per_minute,
..defaults.crawl.rate_limit
},
bootstrap: BootstrapOptions {
..defaults.crawl.bootstrap
},
target: TargetOptions {
..defaults.crawl.target
},
scheduler: SchedulerOptions {
..defaults.crawl.scheduler
},
},
}
}
fn validate(&self) -> Result<(), AppError> {
if self.persistence_queue_capacity == 0 {
return Err(AppError::Config(
"persistence_queue_capacity 必须大于零".to_owned(),
));
}
if self.stats_interval_secs == 0 {
return Err(AppError::Config(
"stats_interval_secs 必须大于零".to_owned(),
));
}
if self.run_duration_secs == Some(0) {
return Err(AppError::Config(
"run_duration_secs 必须大于零或不设置".to_owned(),
));
}
if self.index_batch_size == 0 || self.index_interval_millis == 0 {
return Err(AppError::Config(
"索引批量大小和执行间隔必须大于零".to_owned(),
));
}
if self.dht.metadata_workers == 0 || self.dht.metadata_queue_capacity == 0 {
return Err(AppError::Config(
"Metadata worker 和队列容量必须大于零".to_owned(),
));
}
if self.dht.max_outbound_queries_per_second == 0
|| self.dht.outbound_query_burst == 0
|| self.dht.metadata_connects_per_second == 0
|| self.dht.peer_lookup_max_active == 0
|| self.dht.find_node_max_in_flight == 0
{
return Err(AppError::Config("网络速率和并发上限必须大于零".to_owned()));
}
Ok(())
}
}
impl Default for AppConfig {
fn default() -> Self {
Self {
data_dir: PathBuf::from("data"),
persistence_queue_capacity: 4_096,
stats_interval_secs: 10,
run_duration_secs: None,
index_batch_size: 512,
index_interval_millis: 5_000,
dht: DhtConfig::default(),
http: HttpConfig::default(),
}
}
}
impl Default for DhtConfig {
fn default() -> Self {
Self {
port: 12_313,
netmode: NetworkMode::Ipv4Only,
hash_queue_capacity: 10_000,
max_outbound_queries_per_second: 10,
outbound_query_burst: 2,
metadata_timeout_secs: 4,
metadata_queue_capacity: 10_000,
metadata_workers: 8,
metadata_connects_per_second: 2,
sample_queries_per_second: 1,
peer_lookups_per_second: 1,
peer_lookup_max_active: 4,
find_node_queries_per_second: 6,
find_node_max_in_flight: 12,
new_destinations_per_minute: 60,
}
}
}
impl Default for HttpConfig {
fn default() -> Self {
Self {
listen: SocketAddr::from(([127, 0, 0, 1], 8080)),
}
}
}
impl From<NetworkMode> for NetMode {
fn from(value: NetworkMode) -> Self {
match value {
NetworkMode::Ipv4Only => Self::Ipv4Only,
NetworkMode::Ipv6Only => Self::Ipv6Only,
NetworkMode::DualStack => Self::DualStack,
}
}
}
fn absolute_path(path: &PathBuf) -> Result<PathBuf, AppError> {
if path.is_absolute() {
Ok(path.clone())
} else {
Ok(std::env::current_dir()?.join(path))
}
}
fn normalize_absolute(path: PathBuf) -> Result<PathBuf, AppError> {
if path.is_absolute() {
Ok(path)
} else {
absolute_path(&path)
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
#[test]
fn defaults_produce_valid_dht_options() {
let config = AppConfig::default();
config.validate().unwrap();
let options = config.dht_options();
assert_eq!(options.port, 12_313);
assert_eq!(options.metadata.max_worker_count, 8);
assert_eq!(options.metadata.max_connects_per_second, 2);
assert_eq!(options.max_outbound_queries_per_second, 10);
assert_eq!(options.crawl.rate_limit.max_find_node_rate_per_sec, 6);
}
#[test]
fn relative_data_directory_is_resolved_from_config_file() {
let directory = TempDir::new().unwrap();
let config_path = directory.path().join("service.toml");
fs::write(&config_path, "data_dir = 'state'").unwrap();
let config = Cli {
config: config_path,
data_dir: None,
run_duration_secs: None,
}
.load()
.unwrap();
assert_eq!(config.data_dir, directory.path().join("state"));
}
#[test]
fn unknown_config_field_is_rejected() {
let directory = TempDir::new().unwrap();
let config_path = directory.path().join("service.toml");
fs::write(&config_path, "unknown = true").unwrap();
let error = Cli {
config: config_path,
data_dir: None,
run_duration_secs: None,
}
.load()
.unwrap_err();
assert!(matches!(error, AppError::Toml(_)));
}
}
+1 -1
View File
@@ -1,4 +1,4 @@
// 负责组合 DHT 发现 Metadata 下载和持久化提交管线
mod pipeline;
pub(crate) mod pipeline;
mod worker;
+240
View File
@@ -1 +1,241 @@
// 负责定义采集阶段之间的有界队列背压和任务流转规则
use std::{
sync::{
Arc, Mutex,
atomic::{AtomicU64, AtomicUsize, Ordering},
mpsc::{self, SyncSender, TrySendError},
},
thread::{self, JoinHandle},
};
use dht_crawler::TorrentInfo;
use dht_search::{
domain::TorrentRecord,
storage::{StorageError, TorrentRepository, UpsertOutcome},
};
use tokio::sync::oneshot;
use crate::error::AppError;
#[derive(Clone)]
pub(crate) struct PersistenceIngress {
sender: Arc<Mutex<Option<SyncSender<TorrentRecord>>>>,
stats: Arc<PersistenceStats>,
}
pub(crate) struct PersistencePipeline {
pub(crate) ingress: PersistenceIngress,
pub(crate) fatal: oneshot::Receiver<String>,
worker: JoinHandle<Result<(), StorageError>>,
}
#[derive(Default)]
pub(crate) struct PersistenceStats {
accepted: AtomicU64,
inserted: AtomicU64,
updated: AtomicU64,
rejected_full: AtomicU64,
invalid: AtomicU64,
failed: AtomicU64,
queue_depth: AtomicUsize,
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct PersistenceSnapshot {
pub(crate) accepted: u64,
pub(crate) inserted: u64,
pub(crate) updated: u64,
pub(crate) rejected_full: u64,
pub(crate) invalid: u64,
pub(crate) failed: u64,
pub(crate) queue_depth: usize,
}
impl PersistencePipeline {
pub(crate) fn start(repository: Arc<dyn TorrentRepository>, capacity: usize) -> Self {
let (sender, receiver) = mpsc::sync_channel::<TorrentRecord>(capacity);
let (fatal_tx, fatal) = oneshot::channel();
let stats = Arc::new(PersistenceStats::default());
let worker_stats = stats.clone();
let worker = thread::Builder::new()
.name("torrent-persistence".to_owned())
.spawn(move || {
while let Ok(record) = receiver.recv() {
worker_stats.queue_depth.fetch_sub(1, Ordering::Relaxed);
match repository.upsert(record) {
Ok(UpsertOutcome::Inserted) => {
worker_stats.inserted.fetch_add(1, Ordering::Relaxed);
}
Ok(UpsertOutcome::Updated { .. }) => {
worker_stats.updated.fetch_add(1, Ordering::Relaxed);
}
Err(error) => {
worker_stats.failed.fetch_add(1, Ordering::Relaxed);
let _ = fatal_tx.send(error.to_string());
return Err(error);
}
}
}
Ok(())
})
.expect("persistence worker thread must spawn");
Self {
ingress: PersistenceIngress {
sender: Arc::new(Mutex::new(Some(sender))),
stats,
},
fatal,
worker,
}
}
pub(crate) async fn close_and_join(self) -> Result<(), AppError> {
self.ingress.close();
let result = tokio::task::spawn_blocking(move || self.worker.join())
.await
.map_err(|_| AppError::PersistenceWorkerPanicked)?
.map_err(|_| AppError::PersistenceWorkerPanicked)?;
result.map_err(AppError::from)
}
}
impl PersistenceIngress {
pub(crate) fn try_enqueue(&self, torrent: TorrentInfo) -> bool {
let record = match TorrentRecord::try_from(torrent) {
Ok(record) => record,
Err(error) => {
self.stats.invalid.fetch_add(1, Ordering::Relaxed);
tracing::warn!(%error, "拒绝无效 Metadata");
return false;
}
};
let sender = self
.sender
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let Some(sender) = sender.as_ref() else {
return false;
};
match sender.try_send(record) {
Ok(()) => {
self.stats.accepted.fetch_add(1, Ordering::Relaxed);
self.stats.queue_depth.fetch_add(1, Ordering::Relaxed);
true
}
Err(TrySendError::Full(_)) => {
self.stats.rejected_full.fetch_add(1, Ordering::Relaxed);
false
}
Err(TrySendError::Disconnected(_)) => false,
}
}
pub(crate) fn close(&self) {
self.sender
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.take();
}
pub(crate) fn snapshot(&self) -> PersistenceSnapshot {
self.stats.snapshot()
}
}
impl PersistenceStats {
fn snapshot(&self) -> PersistenceSnapshot {
PersistenceSnapshot {
accepted: self.accepted.load(Ordering::Relaxed),
inserted: self.inserted.load(Ordering::Relaxed),
updated: self.updated.load(Ordering::Relaxed),
rejected_full: self.rejected_full.load(Ordering::Relaxed),
invalid: self.invalid.load(Ordering::Relaxed),
failed: self.failed.load(Ordering::Relaxed),
queue_depth: self.queue_depth.load(Ordering::Relaxed),
}
}
}
#[cfg(test)]
mod tests {
use std::sync::Mutex;
use dht_crawler::FileInfo;
use dht_search::{
domain::{IndexState, InfoHash},
storage::UpsertOutcome,
};
use super::*;
#[derive(Default)]
struct MemoryRepository {
records: Mutex<Vec<TorrentRecord>>,
}
impl TorrentRepository for MemoryRepository {
fn get(&self, info_hash: InfoHash) -> Result<Option<TorrentRecord>, StorageError> {
Ok(self
.records
.lock()
.unwrap()
.iter()
.find(|record| record.info_hash == info_hash)
.cloned())
}
fn upsert(&self, record: TorrentRecord) -> Result<UpsertOutcome, StorageError> {
self.records.lock().unwrap().push(record);
Ok(UpsertOutcome::Inserted)
}
fn observe_existing(&self, _: InfoHash, _: u64) -> Result<bool, StorageError> {
Ok(false)
}
fn by_content_key(&self, _: &[u8; 32], _: usize) -> Result<Vec<InfoHash>, StorageError> {
Ok(Vec::new())
}
fn pending_index(&self, _: usize) -> Result<Vec<InfoHash>, StorageError> {
Ok(Vec::new())
}
fn mark_indexed(&self, _: InfoHash, _: u64) -> Result<(), StorageError> {
Ok(())
}
fn prepare_full_reindex(&self) -> Result<u64, StorageError> {
Ok(0)
}
}
fn torrent() -> TorrentInfo {
TorrentInfo {
info_hash: "0101010101010101010101010101010101010101".into(),
magnet_link: String::new(),
name: "test".into(),
total_size: 1,
files: vec![FileInfo {
path: "test".into(),
size: 1,
}],
piece_length: 16_384,
peers: Vec::new(),
timestamp: 1,
}
}
#[tokio::test]
async fn accepted_record_is_drained_before_shutdown() {
let repository = Arc::new(MemoryRepository::default());
let pipeline = PersistencePipeline::start(repository.clone(), 1);
assert!(pipeline.ingress.try_enqueue(torrent()));
pipeline.close_and_join().await.unwrap();
let records = repository.records.lock().unwrap();
assert_eq!(records.len(), 1);
assert_eq!(records[0].index_state, IndexState::Pending);
}
}
+106
View File
@@ -1 +1,107 @@
// 负责规范化文件结构并生成用于内容聚合的稳定 BLAKE3 指纹
use unicode_normalization::UnicodeNormalization;
use super::torrent::{TorrentFile, TorrentRecordError};
const FINGERPRINT_VERSION: &[u8] = b"dht-search-content-v1\0";
pub fn content_key(files: &[TorrentFile]) -> Result<[u8; 32], TorrentRecordError> {
let mut normalized = normalize_files(files)?;
normalized.sort_unstable();
let mut hasher = blake3::Hasher::new();
hasher.update(FINGERPRINT_VERSION);
hasher.update(&(normalized.len() as u64).to_be_bytes());
for (path, size) in normalized {
hasher.update(&(path.len() as u64).to_be_bytes());
hasher.update(path.as_bytes());
hasher.update(&size.to_be_bytes());
}
Ok(*hasher.finalize().as_bytes())
}
fn normalize_files(files: &[TorrentFile]) -> Result<Vec<(String, u64)>, TorrentRecordError> {
let mut paths = Vec::with_capacity(files.len());
for file in files {
let parts: Vec<String> = file
.path
.replace('\\', "/")
.split('/')
.filter(|part| !part.is_empty() && *part != ".")
.map(|part| part.nfc().collect::<String>().to_lowercase())
.collect();
if parts.is_empty() {
return Err(TorrentRecordError::EmptyNormalizedPath);
}
paths.push((parts, file.size));
}
paths
.into_iter()
.map(|(parts, size)| {
let path = parts.join("/");
if path.is_empty() {
Err(TorrentRecordError::EmptyNormalizedPath)
} else {
Ok((path, size))
}
})
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn equivalent_layouts_have_the_same_content_key() {
let left = vec![
TorrentFile {
path: "Dir\\A.txt".into(),
size: 1,
},
TorrentFile {
path: "Dir/B.bin".into(),
size: 2,
},
];
let right = vec![
TorrentFile {
path: "dir/b.bin".into(),
size: 2,
},
TorrentFile {
path: "dir/a.TXT".into(),
size: 1,
},
];
assert_eq!(content_key(&left).unwrap(), content_key(&right).unwrap());
}
#[test]
fn real_subdirectory_is_part_of_the_content_key() {
let left = vec![TorrentFile {
path: "season1/a.mkv".into(),
size: 1,
}];
let right = vec![TorrentFile {
path: "season2/a.mkv".into(),
size: 1,
}];
assert_ne!(content_key(&left).unwrap(), content_key(&right).unwrap());
}
#[test]
fn file_size_is_part_of_the_content_key() {
let left = vec![TorrentFile {
path: "a.txt".into(),
size: 1,
}];
let right = vec![TorrentFile {
path: "a.txt".into(),
size: 2,
}];
assert_ne!(content_key(&left).unwrap(), content_key(&right).unwrap());
}
}
+7
View File
@@ -2,3 +2,10 @@
mod fingerprint;
mod torrent;
pub use fingerprint::content_key;
#[cfg(test)]
pub(crate) use torrent::test_record;
pub use torrent::{
IndexState, InfoHash, RECORD_SCHEMA_VERSION, TorrentFile, TorrentRecord, TorrentRecordError,
};
+197
View File
@@ -1 +1,198 @@
// 负责定义种子元数据文件条目发现状态和索引状态模型
use dht_crawler::TorrentInfo;
use serde::{Deserialize, Serialize};
use std::{fmt, str::FromStr};
use super::fingerprint::content_key;
pub const RECORD_SCHEMA_VERSION: u16 = 1;
const MAX_STORED_PEERS: usize = 32;
#[derive(Debug, Clone, Copy, PartialEq, Eq, 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 {
pub path: String,
pub size: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum IndexState {
Pending,
Indexed { indexed_at: u64 },
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct TorrentRecord {
pub schema_version: u16,
pub info_hash: InfoHash,
pub name: String,
pub total_size: u64,
pub files: Vec<TorrentFile>,
pub piece_length: u64,
pub source_peers: Vec<String>,
pub content_key: [u8; 32],
pub first_seen: u64,
pub last_seen: u64,
pub seen_count: u64,
pub index_state: IndexState,
}
impl TorrentRecord {
pub fn observe_again(&mut self, timestamp: u64, peers: &[String]) {
self.last_seen = self.last_seen.max(timestamp);
self.seen_count = self.seen_count.saturating_add(1);
for peer in peers {
if self.source_peers.len() >= MAX_STORED_PEERS {
break;
}
if !self.source_peers.contains(peer) {
self.source_peers.push(peer.clone());
}
}
}
}
impl TryFrom<TorrentInfo> for TorrentRecord {
type Error = TorrentRecordError;
fn try_from(info: TorrentInfo) -> Result<Self, Self::Error> {
let info_hash = InfoHash::from_str(&info.info_hash)?;
let files: Vec<_> = info
.files
.into_iter()
.map(|file| TorrentFile {
path: file.path,
size: file.size,
})
.collect();
if info.name.trim().is_empty() {
return Err(TorrentRecordError::EmptyName);
}
if files.is_empty() {
return Err(TorrentRecordError::EmptyFileList);
}
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);
Ok(Self {
schema_version: RECORD_SCHEMA_VERSION,
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,
index_state: IndexState::Pending,
})
}
}
#[derive(Debug, thiserror::Error, PartialEq, Eq)]
pub enum TorrentRecordError {
#[error("infohash 必须是二十字节的十六进制字符串")]
InvalidInfoHash,
#[error("种子名称不能为空")]
EmptyName,
#[error("文件列表不能为空")]
EmptyFileList,
#[error("文件总大小溢出")]
SizeOverflow,
#[error("声明大小 {declared} 与文件计算大小 {calculated} 不一致")]
TotalSizeMismatch { declared: u64, calculated: u64 },
#[error("文件路径规范化后为空")]
EmptyNormalizedPath,
}
#[cfg(test)]
pub(crate) fn test_record(hash_byte: u8, timestamp: u64) -> TorrentRecord {
let files = vec![TorrentFile {
path: "Example/file.txt".to_owned(),
size: 42,
}];
TorrentRecord {
schema_version: RECORD_SCHEMA_VERSION,
info_hash: InfoHash::from_bytes([hash_byte; 20]),
name: "Example".to_owned(),
total_size: 42,
content_key: content_key(&files).expect("test file path is valid"),
files,
piece_length: 16_384,
source_peers: vec!["127.0.0.1:6881".to_owned()],
first_seen: timestamp,
last_seen: timestamp,
seen_count: 1,
index_state: IndexState::Pending,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn infohash_round_trips_as_lowercase_hex() {
let hash = InfoHash::from_str("ABABABABABABABABABABABABABABABABABABABAB").unwrap();
assert_eq!(hash.to_string(), "abababababababababababababababababababab");
}
#[test]
fn repeated_observation_updates_time_count_and_unique_peers() {
let mut record = test_record(1, 10);
record.observe_again(20, &["127.0.0.1:6881".into(), "127.0.0.2:6881".into()]);
assert_eq!(record.first_seen, 10);
assert_eq!(record.last_seen, 20);
assert_eq!(record.seen_count, 2);
assert_eq!(record.source_peers.len(), 2);
}
}
+22
View File
@@ -1 +1,23 @@
// 负责定义应用层统一错误类型和跨模块错误转换边界
#[derive(Debug, thiserror::Error)]
pub(crate) enum AppError {
#[error("I/O 操作失败: {0}")]
Io(#[from] std::io::Error),
#[error("配置解析失败: {0}")]
Toml(#[from] toml::de::Error),
#[error("配置无效: {0}")]
Config(String),
#[error("DHT 服务失败: {0}")]
Dht(#[from] dht_crawler::DHTError),
#[error("存储失败: {0}")]
Storage(#[from] dht_search::storage::StorageError),
#[error("搜索失败: {0}")]
Search(#[from] dht_search::search::SearchError),
#[error("持久化 worker 异常退出")]
PersistenceWorkerPanicked,
#[error("持久化 worker 失败: {0}")]
PersistenceWorker(String),
#[error("索引 worker 失败: {0}")]
IndexWorker(String),
}
+5
View File
@@ -0,0 +1,5 @@
// 负责导出可测试可组合的领域模型和持久化能力
pub mod domain;
pub mod search;
pub mod storage;
+14 -4
View File
@@ -1,14 +1,24 @@
// 负责组装应用依赖启动运行时并协调服务生命周期
use clap::Parser;
mod api;
mod app;
mod config;
mod crawler;
mod domain;
mod error;
mod search;
mod shutdown;
mod storage;
mod telemetry;
fn main() {}
#[tokio::main]
async fn main() {
telemetry::init();
let result = match config::Cli::parse().load() {
Ok(config) => app::run(config).await,
Err(error) => Err(error),
};
if let Err(error) = result {
tracing::error!(%error, "dht-search 退出");
std::process::exit(1);
}
}
+351
View File
@@ -1 +1,352 @@
// 负责批量写入删除提交和从权威存储重建 Tantivy 索引
use std::{
collections::BTreeSet,
ops::Bound,
path::Path,
sync::{Arc, Mutex},
};
use tantivy::{
Index, IndexReader, IndexWriter, ReloadPolicy, TantivyDocument, Term,
collector::{Count, TopDocs},
directory::MmapDirectory,
query::{AllQuery, BooleanQuery, Query, QueryParser, RangeQuery, TermQuery},
schema::{IndexRecordOption, Value},
};
use crate::domain::TorrentRecord;
use crate::storage::TorrentRepository;
use super::{
IndexingError, SearchError,
query::{SearchHit, SearchOptions, SearchPage},
schema::{SearchFields, build_schema},
};
const INDEX_WRITER_MEMORY_BYTES: usize = 64 * 1024 * 1024;
const MAX_PAGE_SIZE: usize = 100;
const MAX_OFFSET: usize = 10_000;
#[derive(Clone)]
pub struct SearchEngine {
inner: Arc<SearchInner>,
}
struct SearchInner {
index: Index,
reader: IndexReader,
writer: Mutex<IndexWriter>,
fields: SearchFields,
}
impl SearchEngine {
pub fn open(path: impl AsRef<Path>) -> Result<Self, SearchError> {
Self::open_with_status(path).map(|(engine, _)| engine)
}
pub fn open_with_status(path: impl AsRef<Path>) -> Result<(Self, bool), SearchError> {
std::fs::create_dir_all(path.as_ref())
.map_err(|error| SearchError::Directory(error.to_string()))?;
let directory = MmapDirectory::open(path.as_ref())
.map_err(|error| SearchError::Directory(error.to_string()))?;
let (expected_schema, fields) = build_schema();
let exists =
Index::exists(&directory).map_err(|error| SearchError::Directory(error.to_string()))?;
let index = if exists {
let index = Index::open(directory)?;
if index.schema() != expected_schema {
return Err(SearchError::IncompatibleSchema);
}
index
} else {
Index::open_or_create(directory, expected_schema)?
};
let reader = index
.reader_builder()
.reload_policy(ReloadPolicy::OnCommitWithDelay)
.try_into()?;
let writer = index.writer_with_num_threads(1, INDEX_WRITER_MEMORY_BYTES)?;
Ok((
Self {
inner: Arc::new(SearchInner {
index,
reader,
writer: Mutex::new(writer),
fields,
}),
},
!exists,
))
}
pub fn index_records(&self, records: &[TorrentRecord]) -> Result<(), SearchError> {
if records.is_empty() {
return Ok(());
}
let fields = self.inner.fields;
let mut writer = self
.inner
.writer
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
for record in records {
writer.delete_term(Term::from_field_text(
fields.info_hash,
&record.info_hash.to_string(),
));
writer.add_document(document(record, fields))?;
}
writer.commit()?;
self.inner.reader.reload()?;
Ok(())
}
pub fn num_docs(&self) -> u64 {
self.inner.reader.searcher().num_docs()
}
pub fn index_pending(
&self,
repository: &dyn TorrentRepository,
limit: usize,
indexed_at: u64,
) -> Result<usize, IndexingError> {
let hashes = repository.pending_index(limit)?;
let mut records = Vec::with_capacity(hashes.len());
for info_hash in hashes {
if let Some(record) = repository.get(info_hash)? {
records.push(record);
}
}
self.index_records(&records)?;
for record in &records {
repository.mark_indexed(record.info_hash, indexed_at)?;
}
Ok(records.len())
}
pub fn search(
&self,
query_text: &str,
offset: usize,
limit: usize,
) -> Result<SearchPage, SearchError> {
self.search_with(SearchOptions {
query: query_text.to_owned(),
offset,
limit,
..SearchOptions::default()
})
}
pub fn search_with(&self, options: SearchOptions) -> Result<SearchPage, SearchError> {
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 {
let parser = QueryParser::for_index(
&self.inner.index,
vec![fields.name, fields.files_text, fields.info_hash],
);
clauses.push(parser.parse_query(query_text)?);
}
if options.min_size.is_some() || options.max_size.is_some() {
let lower = options
.min_size
.map(|value| Bound::Included(Term::from_field_u64(fields.total_size, value)))
.unwrap_or(Bound::Unbounded);
let upper = options
.max_size
.map(|value| Bound::Included(Term::from_field_u64(fields.total_size, value)))
.unwrap_or(Bound::Unbounded);
clauses.push(Box::new(RangeQuery::new(lower, upper)));
}
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,
)));
}
}
let query: Box<dyn Query> = if clauses.len() == 1 {
clauses.pop().expect("one query clause exists")
} else {
Box::new(BooleanQuery::intersection(clauses))
};
let searcher = self.inner.reader.searcher();
let (total, documents) = searcher.search(
query.as_ref(),
&(
Count,
TopDocs::with_limit(limit)
.and_offset(offset)
.order_by_score(),
),
)?;
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.name, "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")?,
score,
});
}
Ok(SearchPage {
total,
offset,
limit,
hits,
})
}
}
fn document(record: &TorrentRecord, fields: SearchFields) -> TantivyDocument {
let mut document = TantivyDocument::default();
document.add_text(fields.info_hash, record.info_hash.to_string());
document.add_text(fields.name, &record.name);
document.add_text(
fields.files_text,
record
.files
.iter()
.map(|file| file.path.as_str())
.collect::<Vec<_>>()
.join(" "),
);
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, record.first_seen);
document.add_u64(fields.last_seen, record.last_seen);
document.add_u64(fields.seen_count, record.seen_count);
document.add_text(fields.content_key, hex::encode(record.content_key));
document
}
fn extensions(record: &TorrentRecord) -> BTreeSet<String> {
record
.files
.iter()
.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 tempfile::TempDir;
use crate::domain::{IndexState, InfoHash, TorrentFile, TorrentRecord};
use super::*;
fn record() -> TorrentRecord {
TorrentRecord {
schema_version: 1,
info_hash: InfoHash::from_bytes([1; 20]),
name: "Ubuntu Linux 24.04".into(),
total_size: 42,
files: vec![TorrentFile {
path: "ubuntu.iso".into(),
size: 42,
}],
piece_length: 16_384,
source_peers: Vec::new(),
content_key: [2; 32],
first_seen: 10,
last_seen: 20,
seen_count: 3,
index_state: IndexState::Pending,
}
}
#[test]
fn record_is_searchable_and_update_is_idempotent() {
let directory = TempDir::new().unwrap();
let engine = SearchEngine::open(directory.path()).unwrap();
let mut record = record();
engine.index_records(&[record.clone()]).unwrap();
record.seen_count = 4;
engine.index_records(&[record]).unwrap();
let page = engine.search("ubuntu", 0, 10).unwrap();
assert_eq!(page.total, 1);
assert_eq!(page.hits[0].name, "Ubuntu Linux 24.04");
assert_eq!(page.hits[0].seen_count, 4);
}
#[test]
fn file_path_is_searchable() {
let directory = TempDir::new().unwrap();
let engine = SearchEngine::open(directory.path()).unwrap();
engine.index_records(&[record()]).unwrap();
assert_eq!(engine.search("ubuntu.iso", 0, 10).unwrap().total, 1);
}
#[test]
fn size_and_extension_filters_use_the_index() {
let directory = TempDir::new().unwrap();
let engine = SearchEngine::open(directory.path()).unwrap();
engine.index_records(&[record()]).unwrap();
let matching = engine
.search_with(SearchOptions {
query: String::new(),
min_size: Some(40),
max_size: Some(50),
extension: Some("ISO".into()),
limit: 10,
..SearchOptions::default()
})
.unwrap();
assert_eq!(matching.total, 1);
let excluded = engine
.search_with(SearchOptions {
query: String::new(),
min_size: Some(100),
limit: 10,
..SearchOptions::default()
})
.unwrap();
assert_eq!(excluded.total, 0);
}
}
+25
View File
@@ -3,3 +3,28 @@
mod indexer;
mod query;
mod schema;
pub use indexer::SearchEngine;
pub use query::{SearchHit, SearchOptions, SearchPage};
#[derive(Debug, thiserror::Error)]
pub enum SearchError {
#[error("搜索索引操作失败: {0}")]
Tantivy(#[from] tantivy::TantivyError),
#[error("搜索查询无效: {0}")]
Query(#[from] tantivy::query::QueryParserError),
#[error("无法打开搜索索引目录: {0}")]
Directory(String),
#[error("搜索索引 schema 与当前版本不兼容")]
IncompatibleSchema,
#[error("搜索文档缺少字段 {0}")]
MissingField(&'static str),
}
#[derive(Debug, thiserror::Error)]
pub enum IndexingError {
#[error(transparent)]
Search(#[from] SearchError),
#[error(transparent)]
Storage(#[from] crate::storage::StorageError),
}
+33
View File
@@ -1 +1,34 @@
// 负责构建全文查询过滤排序分页和内容聚合逻辑
use serde::Serialize;
#[derive(Debug, Clone, Default)]
pub struct SearchOptions {
pub query: String,
pub offset: usize,
pub limit: usize,
pub min_size: Option<u64>,
pub max_size: Option<u64>,
pub extension: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Serialize)]
pub struct SearchHit {
pub info_hash: String,
pub name: String,
pub total_size: u64,
pub file_count: u64,
pub first_seen: u64,
pub last_seen: u64,
pub seen_count: u64,
pub content_key: String,
pub score: f32,
}
#[derive(Debug, Clone, PartialEq, Serialize)]
pub struct SearchPage {
pub total: usize,
pub offset: usize,
pub limit: usize,
pub hits: Vec<SearchHit>,
}
+46
View File
@@ -1 +1,47 @@
// 负责定义 Tantivy 字段分词索引存储和快速字段策略
use tantivy::schema::{FAST, Field, STORED, STRING, Schema, TEXT};
#[derive(Debug, Clone, Copy)]
pub(crate) struct SearchFields {
pub(crate) info_hash: Field,
pub(crate) name: Field,
pub(crate) files_text: Field,
pub(crate) extensions: Field,
pub(crate) total_size: Field,
pub(crate) file_count: Field,
pub(crate) first_seen: Field,
pub(crate) last_seen: Field,
pub(crate) seen_count: Field,
pub(crate) content_key: Field,
}
pub(crate) fn build_schema() -> (Schema, SearchFields) {
let mut builder = Schema::builder();
let info_hash = builder.add_text_field("info_hash", STRING | STORED);
let name = builder.add_text_field("name", TEXT | STORED);
let files_text = builder.add_text_field("files_text", TEXT);
let extensions = builder.add_text_field("extensions", STRING);
let total_size = builder.add_u64_field("total_size", FAST | STORED);
let file_count = builder.add_u64_field("file_count", FAST | STORED);
let first_seen = builder.add_u64_field("first_seen", FAST | STORED);
let last_seen = builder.add_u64_field("last_seen", FAST | STORED);
let seen_count = builder.add_u64_field("seen_count", FAST | STORED);
let content_key = builder.add_text_field("content_key", STRING | STORED);
let schema = builder.build();
(
schema,
SearchFields {
info_hash,
name,
files_text,
extensions,
total_size,
file_count,
first_seen,
last_seen,
seen_count,
content_key,
},
)
}
+22
View File
@@ -1 +1,23 @@
// 负责监听退出信号并协调有界队列排空提交和资源关闭
pub(crate) async fn signal() {
#[cfg(unix)]
{
use tokio::signal::unix::{SignalKind, signal};
let mut terminate = signal(SignalKind::terminate()).expect("SIGTERM handler must install");
tokio::select! {
result = tokio::signal::ctrl_c() => {
if let Err(error) = result {
tracing::error!(%error, "无法监听 Ctrl+C")
}
}
_ = terminate.recv() => {}
}
}
#[cfg(not(unix))]
if let Err(error) = tokio::signal::ctrl_c().await {
tracing::error!(%error, "无法监听 Ctrl+C")
}
}
+83
View File
@@ -1 +1,84 @@
// 负责定义稳定的 RocksDB 键空间编码和版本边界
use crate::domain::InfoHash;
pub(crate) const DATABASE_SCHEMA_VERSION: u32 = 1;
pub(crate) const SCHEMA_VERSION_KEY: &[u8] = b"\x00schema-version";
const TORRENT_PREFIX: u8 = b't';
const CONTENT_PREFIX: u8 = b'c';
const PENDING_INDEX_PREFIX: u8 = b'p';
pub(crate) fn torrent_key(info_hash: InfoHash) -> [u8; 1 + InfoHash::BYTE_LEN] {
prefixed_info_hash(TORRENT_PREFIX, info_hash)
}
pub(crate) fn torrent_prefix() -> [u8; 1] {
[TORRENT_PREFIX]
}
pub(crate) fn pending_index_key(info_hash: InfoHash) -> [u8; 1 + InfoHash::BYTE_LEN] {
prefixed_info_hash(PENDING_INDEX_PREFIX, info_hash)
}
pub(crate) fn pending_index_prefix() -> [u8; 1] {
[PENDING_INDEX_PREFIX]
}
pub(crate) fn content_member_key(
content_key: &[u8; 32],
info_hash: InfoHash,
) -> [u8; 1 + 32 + InfoHash::BYTE_LEN] {
let mut key = [0_u8; 1 + 32 + InfoHash::BYTE_LEN];
key[0] = CONTENT_PREFIX;
key[1..33].copy_from_slice(content_key);
key[33..].copy_from_slice(info_hash.as_bytes());
key
}
pub(crate) fn content_members_prefix(content_key: &[u8; 32]) -> [u8; 1 + 32] {
let mut key = [0_u8; 1 + 32];
key[0] = CONTENT_PREFIX;
key[1..].copy_from_slice(content_key);
key
}
pub(crate) fn decode_content_member_info_hash(
key: &[u8],
content_key: &[u8; 32],
) -> Option<InfoHash> {
let expected_prefix = content_members_prefix(content_key);
if key.len() != 1 + 32 + InfoHash::BYTE_LEN || !key.starts_with(&expected_prefix) {
return None;
}
let bytes: [u8; InfoHash::BYTE_LEN] = key[33..].try_into().ok()?;
Some(InfoHash::from_bytes(bytes))
}
pub(crate) fn decode_pending_info_hash(key: &[u8]) -> Option<InfoHash> {
if key.len() != 1 + InfoHash::BYTE_LEN || key.first().copied() != Some(PENDING_INDEX_PREFIX) {
return None;
}
let bytes: [u8; InfoHash::BYTE_LEN] = key[1..].try_into().ok()?;
Some(InfoHash::from_bytes(bytes))
}
fn prefixed_info_hash(prefix: u8, info_hash: InfoHash) -> [u8; 1 + InfoHash::BYTE_LEN] {
let mut key = [0_u8; 1 + InfoHash::BYTE_LEN];
key[0] = prefix;
key[1..].copy_from_slice(info_hash.as_bytes());
key
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn pending_key_round_trips_infohash() {
let hash = InfoHash::from_bytes([7; 20]);
assert_eq!(
decode_pending_info_hash(&pending_index_key(hash)),
Some(hash)
);
}
}
+5
View File
@@ -2,4 +2,9 @@
mod keys;
mod repository;
#[cfg(feature = "rocksdb-storage")]
mod rocks;
pub use repository::{StorageError, TorrentRepository, UpsertOutcome};
#[cfg(feature = "rocksdb-storage")]
pub use rocks::RocksTorrentRepository;
+64
View File
@@ -1 +1,65 @@
// 负责定义元数据去重状态恢复和索引任务所需的存储接口
use crate::domain::{InfoHash, TorrentRecord};
pub trait TorrentRepository: Send + Sync {
fn get(&self, info_hash: InfoHash) -> Result<Option<TorrentRecord>, StorageError>;
fn contains(&self, info_hash: InfoHash) -> Result<bool, StorageError> {
self.get(info_hash).map(|record| record.is_some())
}
fn upsert(&self, observation: TorrentRecord) -> Result<UpsertOutcome, StorageError>;
fn observe_existing(&self, info_hash: InfoHash, observed_at: u64)
-> Result<bool, StorageError>;
fn filter_unknown_and_observe(
&self,
info_hashes: &[InfoHash],
observed_at: u64,
) -> Result<Vec<InfoHash>, StorageError> {
let mut unknown = Vec::with_capacity(info_hashes.len());
for info_hash in info_hashes {
if !self.observe_existing(*info_hash, observed_at)? {
unknown.push(*info_hash);
}
}
Ok(unknown)
}
fn by_content_key(
&self,
content_key: &[u8; 32],
limit: usize,
) -> Result<Vec<InfoHash>, StorageError>;
fn pending_index(&self, limit: usize) -> Result<Vec<InfoHash>, StorageError>;
fn mark_indexed(&self, info_hash: InfoHash, indexed_at: u64) -> Result<(), StorageError>;
fn prepare_full_reindex(&self) -> Result<u64, StorageError>;
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum UpsertOutcome {
Inserted,
Updated { seen_count: u64 },
}
#[derive(Debug, thiserror::Error)]
pub enum StorageError {
#[cfg(feature = "rocksdb-storage")]
#[error("RocksDB 操作失败: {0}")]
RocksDb(#[from] rocksdb::Error),
#[error("记录编码失败: {0}")]
Encode(#[from] rmp_serde::encode::Error),
#[error("记录解码失败: {0}")]
Decode(#[from] rmp_serde::decode::Error),
#[error("数据库 schema 版本不受支持 expected={expected} actual={actual}")]
SchemaVersion { expected: u32, actual: u32 },
#[error("数据库 schema 版本数据损坏")]
CorruptSchemaVersion,
#[error("待索引记录不存在 infohash={0}")]
MissingRecord(InfoHash),
}
+420
View File
@@ -1 +1,421 @@
// 负责实现 RocksDB 打开配置批量写入精确查询和关闭流程
use std::{path::Path, sync::Mutex};
use rocksdb::{
BlockBasedOptions, Cache, DB, DBCompressionType, Direction, IteratorMode, Options,
SliceTransform, WriteBatch,
};
use crate::domain::{IndexState, InfoHash, TorrentRecord};
use super::{
keys::{
DATABASE_SCHEMA_VERSION, SCHEMA_VERSION_KEY, content_member_key, content_members_prefix,
decode_content_member_info_hash, decode_pending_info_hash, pending_index_key,
pending_index_prefix, torrent_key, torrent_prefix,
},
repository::{StorageError, TorrentRepository, UpsertOutcome},
};
const DEFAULT_BLOCK_CACHE_BYTES: usize = 64 * 1024 * 1024;
pub struct RocksTorrentRepository {
db: DB,
write_lock: Mutex<()>,
}
impl RocksTorrentRepository {
pub fn open(path: impl AsRef<Path>) -> Result<Self, StorageError> {
let mut block_options = BlockBasedOptions::default();
block_options.set_bloom_filter(10.0, false);
let block_cache = Cache::new_lru_cache(DEFAULT_BLOCK_CACHE_BYTES);
block_options.set_block_cache(&block_cache);
let mut options = Options::default();
options.create_if_missing(true);
options.set_compression_type(DBCompressionType::Lz4);
options.set_block_based_table_factory(&block_options);
options.set_prefix_extractor(SliceTransform::create_fixed_prefix(1));
options.set_max_open_files(256);
let repository = Self {
db: DB::open(&options, path)?,
write_lock: Mutex::new(()),
};
repository.ensure_schema_version()?;
Ok(repository)
}
fn ensure_schema_version(&self) -> Result<(), StorageError> {
let expected = DATABASE_SCHEMA_VERSION.to_be_bytes();
match self.db.get(SCHEMA_VERSION_KEY)? {
None => self
.db
.put(SCHEMA_VERSION_KEY, expected)
.map_err(Into::into),
Some(value) if value.as_ref() == expected => Ok(()),
Some(value) => {
let actual = value
.as_slice()
.try_into()
.map(u32::from_be_bytes)
.map_err(|_| StorageError::CorruptSchemaVersion)?;
Err(StorageError::SchemaVersion {
expected: DATABASE_SCHEMA_VERSION,
actual,
})
}
}
}
fn encode(record: &TorrentRecord) -> Result<Vec<u8>, StorageError> {
rmp_serde::to_vec_named(record).map_err(Into::into)
}
fn decode(bytes: &[u8]) -> Result<TorrentRecord, StorageError> {
rmp_serde::from_slice(bytes).map_err(Into::into)
}
}
impl TorrentRepository for RocksTorrentRepository {
fn get(&self, info_hash: InfoHash) -> Result<Option<TorrentRecord>, StorageError> {
self.db
.get(torrent_key(info_hash))?
.map(|bytes| Self::decode(&bytes))
.transpose()
}
fn upsert(&self, observation: TorrentRecord) -> Result<UpsertOutcome, StorageError> {
let _guard = self
.write_lock
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if let Some(mut current) = self.get(observation.info_hash)? {
current.observe_again(observation.last_seen, &observation.source_peers);
current.index_state = IndexState::Pending;
let mut batch = WriteBatch::default();
batch.put(torrent_key(current.info_hash), Self::encode(&current)?);
batch.put(pending_index_key(current.info_hash), []);
self.db.write(batch)?;
return Ok(UpsertOutcome::Updated {
seen_count: current.seen_count,
});
}
let mut batch = WriteBatch::default();
batch.put(
torrent_key(observation.info_hash),
Self::encode(&observation)?,
);
batch.put(
content_member_key(&observation.content_key, observation.info_hash),
[],
);
batch.put(pending_index_key(observation.info_hash), []);
self.db.write(batch)?;
Ok(UpsertOutcome::Inserted)
}
fn observe_existing(
&self,
info_hash: InfoHash,
observed_at: u64,
) -> Result<bool, StorageError> {
let _guard = self
.write_lock
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let Some(mut record) = self.get(info_hash)? else {
return Ok(false);
};
record.observe_again(observed_at, &[]);
record.index_state = IndexState::Pending;
let mut batch = WriteBatch::default();
batch.put(torrent_key(info_hash), Self::encode(&record)?);
batch.put(pending_index_key(info_hash), []);
self.db.write(batch)?;
Ok(true)
}
fn filter_unknown_and_observe(
&self,
info_hashes: &[InfoHash],
observed_at: u64,
) -> Result<Vec<InfoHash>, StorageError> {
if info_hashes.is_empty() {
return Ok(Vec::new());
}
let _guard = self
.write_lock
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let keys: Vec<_> = info_hashes.iter().copied().map(torrent_key).collect();
let records = self.db.multi_get(keys.iter());
let mut unknown = Vec::with_capacity(info_hashes.len());
let mut batch = WriteBatch::default();
let mut updated = 0_usize;
for ((info_hash, key), record) in info_hashes.iter().zip(&keys).zip(records) {
let Some(bytes) = record? else {
unknown.push(*info_hash);
continue;
};
let mut record = Self::decode(&bytes)?;
record.observe_again(observed_at, &[]);
record.index_state = IndexState::Pending;
batch.put(key, Self::encode(&record)?);
batch.put(pending_index_key(*info_hash), []);
updated += 1;
}
if updated > 0 {
self.db.write(batch)?;
}
Ok(unknown)
}
fn by_content_key(
&self,
content_key: &[u8; 32],
limit: usize,
) -> Result<Vec<InfoHash>, StorageError> {
if limit == 0 {
return Ok(Vec::new());
}
let prefix = content_members_prefix(content_key);
let iterator = self
.db
.iterator(IteratorMode::From(&prefix, Direction::Forward));
let mut hashes = Vec::with_capacity(limit.min(1024));
for entry in iterator {
let (key, _) = entry?;
let Some(info_hash) = decode_content_member_info_hash(&key, content_key) else {
break;
};
hashes.push(info_hash);
if hashes.len() == limit {
break;
}
}
Ok(hashes)
}
fn pending_index(&self, limit: usize) -> Result<Vec<InfoHash>, StorageError> {
if limit == 0 {
return Ok(Vec::new());
}
let prefix = pending_index_prefix();
let iterator = self
.db
.iterator(IteratorMode::From(&prefix, Direction::Forward));
let mut hashes = Vec::with_capacity(limit.min(1024));
for entry in iterator {
let (key, _) = entry?;
let Some(info_hash) = decode_pending_info_hash(&key) else {
break;
};
hashes.push(info_hash);
if hashes.len() == limit {
break;
}
}
Ok(hashes)
}
fn mark_indexed(&self, info_hash: InfoHash, indexed_at: u64) -> Result<(), StorageError> {
let _guard = self
.write_lock
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let Some(mut record) = self.get(info_hash)? else {
return Err(StorageError::MissingRecord(info_hash));
};
record.index_state = IndexState::Indexed { indexed_at };
let mut batch = WriteBatch::default();
batch.put(torrent_key(info_hash), Self::encode(&record)?);
batch.delete(pending_index_key(info_hash));
self.db.write(batch)?;
Ok(())
}
fn prepare_full_reindex(&self) -> Result<u64, StorageError> {
const BATCH_SIZE: usize = 1_000;
let _guard = self
.write_lock
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let prefix = torrent_prefix();
let iterator = self
.db
.iterator(IteratorMode::From(&prefix, Direction::Forward));
let mut batch = WriteBatch::default();
let mut batch_len = 0_usize;
let mut total = 0_u64;
for entry in iterator {
let (key, value) = entry?;
if !key.starts_with(&prefix) {
break;
}
let mut record = Self::decode(&value)?;
record.index_state = IndexState::Pending;
batch.put(&key, Self::encode(&record)?);
batch.put(pending_index_key(record.info_hash), []);
batch_len += 1;
total += 1;
if batch_len == BATCH_SIZE {
self.db.write(batch)?;
batch = WriteBatch::default();
batch_len = 0;
}
}
if batch_len > 0 {
self.db.write(batch)?;
}
Ok(total)
}
}
#[cfg(test)]
mod tests {
use tempfile::TempDir;
use crate::domain::{IndexState, test_record};
use super::*;
#[test]
fn record_survives_close_and_reopen() {
let directory = TempDir::new().unwrap();
let expected = test_record(1, 10);
{
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
assert_eq!(
repository.upsert(expected.clone()).unwrap(),
UpsertOutcome::Inserted
);
}
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
assert_eq!(repository.get(expected.info_hash).unwrap(), Some(expected));
}
#[test]
fn duplicate_updates_observation_without_creating_another_pending_item() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let first = test_record(2, 10);
let mut second = first.clone();
second.last_seen = 20;
second.source_peers = vec!["127.0.0.2:6881".into()];
assert_eq!(
repository.upsert(first.clone()).unwrap(),
UpsertOutcome::Inserted
);
assert!(repository.contains(first.info_hash).unwrap());
assert_eq!(
repository.upsert(second).unwrap(),
UpsertOutcome::Updated { seen_count: 2 }
);
let stored = repository.get(first.info_hash).unwrap().unwrap();
assert_eq!(stored.first_seen, 10);
assert_eq!(stored.last_seen, 20);
assert_eq!(stored.seen_count, 2);
assert_eq!(repository.pending_index(10).unwrap(), vec![first.info_hash]);
}
#[test]
fn batch_triage_updates_existing_and_returns_only_unknown_hashes() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let existing = test_record(1, 10);
let unknown = test_record(2, 10).info_hash;
repository.upsert(existing.clone()).unwrap();
let admitted = repository
.filter_unknown_and_observe(&[existing.info_hash, unknown], 20)
.unwrap();
assert_eq!(admitted, vec![unknown]);
let updated = repository.get(existing.info_hash).unwrap().unwrap();
assert_eq!(updated.last_seen, 20);
assert_eq!(updated.seen_count, existing.seen_count + 1);
assert_eq!(updated.index_state, IndexState::Pending);
}
#[test]
fn existing_hash_can_be_observed_without_downloading_metadata_again() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let record = test_record(6, 10);
assert!(!repository.observe_existing(record.info_hash, 5).unwrap());
repository.upsert(record.clone()).unwrap();
assert!(repository.observe_existing(record.info_hash, 30).unwrap());
let stored = repository.get(record.info_hash).unwrap().unwrap();
assert_eq!(stored.first_seen, 10);
assert_eq!(stored.last_seen, 30);
assert_eq!(stored.seen_count, 2);
}
#[test]
fn equal_content_maps_multiple_infohashes_without_merging_records() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let first = test_record(4, 10);
let mut second = test_record(5, 20);
second.content_key = first.content_key;
repository.upsert(first.clone()).unwrap();
repository.upsert(second.clone()).unwrap();
let mut hashes = repository.by_content_key(&first.content_key, 10).unwrap();
hashes.sort_unstable_by_key(ToString::to_string);
assert_eq!(hashes, vec![first.info_hash, second.info_hash]);
assert!(repository.get(first.info_hash).unwrap().is_some());
assert!(repository.get(second.info_hash).unwrap().is_some());
}
#[test]
fn marking_indexed_is_atomic_with_removing_pending_marker() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let record = test_record(3, 10);
repository.upsert(record.clone()).unwrap();
repository.mark_indexed(record.info_hash, 30).unwrap();
assert!(repository.pending_index(10).unwrap().is_empty());
assert_eq!(
repository
.get(record.info_hash)
.unwrap()
.unwrap()
.index_state,
IndexState::Indexed { indexed_at: 30 }
);
}
#[test]
fn full_reindex_restores_pending_markers_for_every_record() {
let directory = TempDir::new().unwrap();
let repository = RocksTorrentRepository::open(directory.path()).unwrap();
let first = test_record(7, 10);
let second = test_record(8, 10);
repository.upsert(first.clone()).unwrap();
repository.upsert(second.clone()).unwrap();
repository.mark_indexed(first.info_hash, 20).unwrap();
repository.mark_indexed(second.info_hash, 20).unwrap();
assert!(repository.pending_index(10).unwrap().is_empty());
assert_eq!(repository.prepare_full_reindex().unwrap(), 2);
assert_eq!(repository.pending_index(10).unwrap().len(), 2);
assert_eq!(
repository
.get(first.info_hash)
.unwrap()
.unwrap()
.index_state,
IndexState::Pending
);
}
}
+15
View File
@@ -1 +1,16 @@
// 负责初始化结构化日志指标和运行状态观测
use tracing_subscriber::EnvFilter;
pub(crate) fn init() {
let filter = EnvFilter::try_from_default_env()
.unwrap_or_else(|_| EnvFilter::new("warn,dht_search=info,dht_crawler=info"));
if let Err(error) = tracing_subscriber::fmt()
.with_env_filter(filter)
.with_target(true)
.with_ansi(std::io::IsTerminal::is_terminal(&std::io::stderr()))
.try_init()
{
eprintln!("无法初始化日志订阅器: {error}");
}
}