diff --git a/.gitignore b/.gitignore index fa1c213..8e85455 100644 --- a/.gitignore +++ b/.gitignore @@ -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/ \ No newline at end of file +opencodes/ diff --git a/AGENTS.md b/AGENTS.md index 2c4f47e..4f97205 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -90,6 +90,10 @@ Bloom Filter 只能作为前置加速结构不得作为最终去重依据 新业务代码写入 `dht-search` +项目阶段任务完成状态和验收标准统一维护在根目录 `TODOS.md` + +需求实现或技术决策发生变化时必须同步更新 `TODOS.md` + `dht-crawler` 只接受可复用的 DHT 基础能力不得包含数据库搜索接口或部署逻辑 每个 Rust 文件顶部必须使用中文行注释描述该文件的功能边界且注释行尾不添加标点 diff --git a/README.md b/README.md index 1a51ee4..67ea47c 100644 --- a/README.md +++ b/README.md @@ -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) diff --git a/TODOS.md b/TODOS.md new file mode 100644 index 0000000..8a4eb5c --- /dev/null +++ b/TODOS.md @@ -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 端到端测试 diff --git a/dht-crawler/README.md b/dht-crawler/README.md index f426180..6dbc62a 100644 --- a/dht-crawler/README.md +++ b/dht-crawler/README.md @@ -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 端口； diff --git a/dht-crawler/src/budget.rs b/dht-crawler/src/budget.rs index f76bbe8..acfb00e 100644 --- a/dht-crawler/src/budget.rs +++ b/dht-crawler/src/budget.rs @@ -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>, +} + +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)); + } } diff --git a/dht-crawler/src/crawl_engine.rs b/dht-crawler/src/crawl_engine.rs index e850dbf..24f68e4 100644 --- a/dht-crawler/src/crawl_engine.rs +++ b/dht-crawler/src/crawl_engine.rs @@ -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>, pub(crate) node_count: Arc, 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, 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")); diff --git a/dht-crawler/src/metadata.rs b/dht-crawler/src/metadata.rs index 517ddc6..d062a3f 100644 --- a/dht-crawler/src/metadata.rs +++ b/dht-crawler/src/metadata.rs @@ -41,6 +41,36 @@ enum PeerFailureReason { ConnectFailed, } +struct ConnectRateLimiter { + interval: Duration, + next_start: Mutex, +} + +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, + connect_rate_limiter: Arc, } 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, diff --git a/dht-crawler/src/peer_lookup.rs b/dht-crawler/src/peer_lookup.rs index 9ec8780..06ca1ae 100644 --- a/dht-crawler/src/peer_lookup.rs +++ b/dht-crawler/src/peer_lookup.rs @@ -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(); } } diff --git a/dht-crawler/src/runtime_stats.rs b/dht-crawler/src/runtime_stats.rs index 771df77..8d2bd44 100644 --- a/dht-crawler/src/runtime_stats.rs +++ b/dht-crawler/src/runtime_stats.rs @@ -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, diff --git a/dht-crawler/src/sample_infohashes.rs b/dht-crawler/src/sample_infohashes.rs index ab5ece8..c51f70f 100644 --- a/dht-crawler/src/sample_infohashes.rs +++ b/dht-crawler/src/sample_infohashes.rs @@ -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> + Send>> + + Send + + Sync + + 'static, +>; + #[derive(Clone)] pub(crate) struct SampleInfohashesHandle { response_tx: mpsc::Sender, @@ -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>, 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>, snapshot: Arc>, crawl_engine: Arc, - peer_lookup: PeerLookupHandle, + admission_tx: mpsc::Sender, response_rx: mpsc::Receiver, 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, + hash_admission: Arc>, + peer_lookup_tx: mpsc::Sender, + 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 = 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(); + } } diff --git a/dht-crawler/src/server.rs b/dht-crawler/src/server.rs index a6783d6..8bd5c06 100644 --- a/dht-crawler/src/server.rs +++ b/dht-crawler/src/server.rs @@ -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>, hash_filter: Arc>, on_metadata_fetch: Arc>, + sample_hash_admission: Arc>, metadata_completion_callback: Arc>, error_callback: Arc>, crawl_engine: Arc, @@ -353,18 +354,26 @@ impl DHTServer { mpsc::channel::(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(&self, callback: F) + where + F: Fn(Vec<[u8; 20]>) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { + let callback: Arc = + 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. diff --git a/dht-crawler/src/types.rs b/dht-crawler/src/types.rs index 18b81e9..890b799 100644 --- a/dht-crawler/src/types.rs +++ b/dht-crawler/src/types.rs @@ -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, diff --git a/dht-search.example.toml b/dht-search.example.toml new file mode 100644 index 0000000..557918d --- /dev/null +++ b/dht-search.example.toml @@ -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" diff --git a/dht-search/Cargo.toml b/dht-search/Cargo.toml index 51f4ded..9ff7fc6 100644 --- a/dht-search/Cargo.toml +++ b/dht-search/Cargo.toml @@ -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" diff --git a/dht-search/README.md b/dht-search/README.md new file mode 100644 index 0000000..d4226e5 --- /dev/null +++ b/dht-search/README.md @@ -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 服务 diff --git a/dht-search/src/api/handlers.rs b/dht-search/src/api/handlers.rs index 202e1b6..4816ee3 100644 --- a/dht-search/src/api/handlers.rs +++ b/dht-search/src/api/handlers.rs @@ -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 { + Json(StatusResponse { status: "ok" }) +} + +pub(crate) async fn ready() -> Json { + Json(StatusResponse { status: "ready" }) +} + +pub(crate) async fn stats(State(state): State) -> Json { + 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, + Query(request): Query, +) -> Result, 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, + Path(info_hash): Path, +) -> Result, 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) -> Self { + Self { + status: StatusCode::BAD_REQUEST, + message: message.into(), + } + } + + fn not_found(message: impl Into) -> Self { + Self { + status: StatusCode::NOT_FOUND, + message: message.into(), + } + } + + fn internal(message: impl Into) -> 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() + } +} diff --git a/dht-search/src/api/mod.rs b/dht-search/src/api/mod.rs index 43a6961..5bb779d 100644 --- a/dht-search/src/api/mod.rs +++ b/dht-search/src/api/mod.rs @@ -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, + 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 +} diff --git a/dht-search/src/api/request.rs b/dht-search/src/api/request.rs index 06c215a..aa78dd9 100644 --- a/dht-search/src/api/request.rs +++ b/dht-search/src/api/request.rs @@ -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, + pub(crate) max_size: Option, + pub(crate) extension: Option, +} diff --git a/dht-search/src/api/response.rs b/dht-search/src/api/response.rs index 2bc0779..ae9b311 100644 --- a/dht-search/src/api/response.rs +++ b/dht-search/src/api/response.rs @@ -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, + 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 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, + } + } +} diff --git a/dht-search/src/app.rs b/dht-search/src/app.rs index 7960a46..358e6de 100644 --- a/dht-search/src/app.rs +++ b/dht-search/src/app.rs @@ -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 = 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, + search: SearchEngine, + batch_size: usize, + interval: Duration, + cancel: CancellationToken, + fatal: tokio::sync::oneshot::Sender, +) -> 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, + search: SearchEngine, + batch_size: usize, +) -> Result { + 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() +} diff --git a/dht-search/src/config.rs b/dht-search/src/config.rs index 161389b..a0354fd 100644 --- a/dht-search/src/config.rs +++ b/dht-search/src/config.rs @@ -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, + #[arg(long)] + run_duration_secs: Option, +} + +#[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, + 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 { + 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::(&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 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 { + if path.is_absolute() { + Ok(path.clone()) + } else { + Ok(std::env::current_dir()?.join(path)) + } +} + +fn normalize_absolute(path: PathBuf) -> Result { + 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(_))); + } +} diff --git a/dht-search/src/crawler/mod.rs b/dht-search/src/crawler/mod.rs index 057d0fe..5611e51 100644 --- a/dht-search/src/crawler/mod.rs +++ b/dht-search/src/crawler/mod.rs @@ -1,4 +1,4 @@ // 负责组合 DHT 发现 Metadata 下载和持久化提交管线 -mod pipeline; +pub(crate) mod pipeline; mod worker; diff --git a/dht-search/src/crawler/pipeline.rs b/dht-search/src/crawler/pipeline.rs index a7f52ce..3c73257 100644 --- a/dht-search/src/crawler/pipeline.rs +++ b/dht-search/src/crawler/pipeline.rs @@ -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>>>, + stats: Arc, +} + +pub(crate) struct PersistencePipeline { + pub(crate) ingress: PersistenceIngress, + pub(crate) fatal: oneshot::Receiver, + worker: JoinHandle>, +} + +#[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, capacity: usize) -> Self { + let (sender, receiver) = mpsc::sync_channel::(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>, + } + + impl TorrentRepository for MemoryRepository { + fn get(&self, info_hash: InfoHash) -> Result, StorageError> { + Ok(self + .records + .lock() + .unwrap() + .iter() + .find(|record| record.info_hash == info_hash) + .cloned()) + } + + fn upsert(&self, record: TorrentRecord) -> Result { + self.records.lock().unwrap().push(record); + Ok(UpsertOutcome::Inserted) + } + + fn observe_existing(&self, _: InfoHash, _: u64) -> Result { + Ok(false) + } + + fn by_content_key(&self, _: &[u8; 32], _: usize) -> Result, StorageError> { + Ok(Vec::new()) + } + + fn pending_index(&self, _: usize) -> Result, StorageError> { + Ok(Vec::new()) + } + + fn mark_indexed(&self, _: InfoHash, _: u64) -> Result<(), StorageError> { + Ok(()) + } + + fn prepare_full_reindex(&self) -> Result { + 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); + } +} diff --git a/dht-search/src/domain/fingerprint.rs b/dht-search/src/domain/fingerprint.rs index 89850bc..f66d015 100644 --- a/dht-search/src/domain/fingerprint.rs +++ b/dht-search/src/domain/fingerprint.rs @@ -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, TorrentRecordError> { + let mut paths = Vec::with_capacity(files.len()); + for file in files { + let parts: Vec = file + .path + .replace('\\', "/") + .split('/') + .filter(|part| !part.is_empty() && *part != ".") + .map(|part| part.nfc().collect::().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()); + } +} diff --git a/dht-search/src/domain/mod.rs b/dht-search/src/domain/mod.rs index b217dbc..1d692bd 100644 --- a/dht-search/src/domain/mod.rs +++ b/dht-search/src/domain/mod.rs @@ -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, +}; diff --git a/dht-search/src/domain/torrent.rs b/dht-search/src/domain/torrent.rs index 9933fdc..ab8eb4e 100644 --- a/dht-search/src/domain/torrent.rs +++ b/dht-search/src/domain/torrent.rs @@ -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 { + 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, + pub piece_length: u64, + pub source_peers: Vec, + 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 for TorrentRecord { + type Error = TorrentRecordError; + + fn try_from(info: TorrentInfo) -> Result { + 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); + } +} diff --git a/dht-search/src/error.rs b/dht-search/src/error.rs index 464004c..6a78c3e 100644 --- a/dht-search/src/error.rs +++ b/dht-search/src/error.rs @@ -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), +} diff --git a/dht-search/src/lib.rs b/dht-search/src/lib.rs new file mode 100644 index 0000000..fd4913c --- /dev/null +++ b/dht-search/src/lib.rs @@ -0,0 +1,5 @@ +// 负责导出可测试可组合的领域模型和持久化能力 + +pub mod domain; +pub mod search; +pub mod storage; diff --git a/dht-search/src/main.rs b/dht-search/src/main.rs index 5c224ad..0fa921c 100644 --- a/dht-search/src/main.rs +++ b/dht-search/src/main.rs @@ -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); + } +} diff --git a/dht-search/src/search/indexer.rs b/dht-search/src/search/indexer.rs index 17fb063..e8ee0dd 100644 --- a/dht-search/src/search/indexer.rs +++ b/dht-search/src/search/indexer.rs @@ -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, +} + +struct SearchInner { + index: Index, + reader: IndexReader, + writer: Mutex, + fields: SearchFields, +} + +impl SearchEngine { + pub fn open(path: impl AsRef) -> Result { + Self::open_with_status(path).map(|(engine, _)| engine) + } + + pub fn open_with_status(path: impl AsRef) -> 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 { + 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 { + self.search_with(SearchOptions { + query: query_text.to_owned(), + offset, + limit, + ..SearchOptions::default() + }) + } + + pub fn search_with(&self, options: SearchOptions) -> Result { + 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> = 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 = 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::>() + .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 { + 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 { + 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 { + 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); + } +} diff --git a/dht-search/src/search/mod.rs b/dht-search/src/search/mod.rs index c821804..c17d65c 100644 --- a/dht-search/src/search/mod.rs +++ b/dht-search/src/search/mod.rs @@ -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), +} diff --git a/dht-search/src/search/query.rs b/dht-search/src/search/query.rs index 92b7af0..1a2d4bb 100644 --- a/dht-search/src/search/query.rs +++ b/dht-search/src/search/query.rs @@ -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, + pub max_size: Option, + pub extension: Option, +} + +#[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, +} diff --git a/dht-search/src/search/schema.rs b/dht-search/src/search/schema.rs index 9272546..13b143b 100644 --- a/dht-search/src/search/schema.rs +++ b/dht-search/src/search/schema.rs @@ -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, + }, + ) +} diff --git a/dht-search/src/shutdown.rs b/dht-search/src/shutdown.rs index d3fd465..6f79fad 100644 --- a/dht-search/src/shutdown.rs +++ b/dht-search/src/shutdown.rs @@ -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") + } +} diff --git a/dht-search/src/storage/keys.rs b/dht-search/src/storage/keys.rs index fd9192f..ece03d2 100644 --- a/dht-search/src/storage/keys.rs +++ b/dht-search/src/storage/keys.rs @@ -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 { + 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 { + 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) + ); + } +} diff --git a/dht-search/src/storage/mod.rs b/dht-search/src/storage/mod.rs index ba41e10..b5a78d9 100644 --- a/dht-search/src/storage/mod.rs +++ b/dht-search/src/storage/mod.rs @@ -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; diff --git a/dht-search/src/storage/repository.rs b/dht-search/src/storage/repository.rs index dadd914..9a50141 100644 --- a/dht-search/src/storage/repository.rs +++ b/dht-search/src/storage/repository.rs @@ -1 +1,65 @@ // 负责定义元数据去重状态恢复和索引任务所需的存储接口 + +use crate::domain::{InfoHash, TorrentRecord}; + +pub trait TorrentRepository: Send + Sync { + fn get(&self, info_hash: InfoHash) -> Result, StorageError>; + + fn contains(&self, info_hash: InfoHash) -> Result { + self.get(info_hash).map(|record| record.is_some()) + } + + fn upsert(&self, observation: TorrentRecord) -> Result; + + fn observe_existing(&self, info_hash: InfoHash, observed_at: u64) + -> Result; + + fn filter_unknown_and_observe( + &self, + info_hashes: &[InfoHash], + observed_at: u64, + ) -> Result, 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, StorageError>; + + fn pending_index(&self, limit: usize) -> Result, StorageError>; + + fn mark_indexed(&self, info_hash: InfoHash, indexed_at: u64) -> Result<(), StorageError>; + + fn prepare_full_reindex(&self) -> Result; +} + +#[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), +} diff --git a/dht-search/src/storage/rocks.rs b/dht-search/src/storage/rocks.rs index 0899d6d..045136d 100644 --- a/dht-search/src/storage/rocks.rs +++ b/dht-search/src/storage/rocks.rs @@ -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) -> Result { + 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, StorageError> { + rmp_serde::to_vec_named(record).map_err(Into::into) + } + + fn decode(bytes: &[u8]) -> Result { + rmp_serde::from_slice(bytes).map_err(Into::into) + } +} + +impl TorrentRepository for RocksTorrentRepository { + fn get(&self, info_hash: InfoHash) -> Result, StorageError> { + self.db + .get(torrent_key(info_hash))? + .map(|bytes| Self::decode(&bytes)) + .transpose() + } + + fn upsert(&self, observation: TorrentRecord) -> Result { + 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(¤t)?); + 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 { + 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, 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, 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, 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 { + 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 + ); + } +} diff --git a/dht-search/src/telemetry.rs b/dht-search/src/telemetry.rs index 3ad4ce5..81573d6 100644 --- a/dht-search/src/telemetry.rs +++ b/dht-search/src/telemetry.rs @@ -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}"); + } +}