commit d366c088a378a09bd355f352f5095a01faf8154b Author: chuan Date: Thu Aug 6 23:02:18 2026 +0800 feat: init diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..8e85455 --- /dev/null +++ b/.gitignore @@ -0,0 +1,27 @@ +# Rust +/target/ +/.tools/ +/.run-data/ +/.remote-data/ +/data/ +/dht-search.toml +Cargo.lock +**/*.rs.bk + +# IDE +.idea/ +.vscode/ +*.swp + +# 输出文件 +torrents/ +*.json +*.db + +# 日志 +*.log + +# 个人脚本（不提交到仓库） +scripts/ + +opencodes/ diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..4f97205 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,115 @@ +# DHT 元数据搜索服务开发约定 + +## 项目目标 + +本项目用于持续或间歇地从 BitTorrent DHT 网络发现 infohash 获取元数据并提供本地全文搜索和高性能过滤能力 + +基础 DHT 协议和抓取能力保留在 `dht-crawler` 中 + +面向最终用户运行的服务代码统一放在 `dht-search` 中 + +## 技术方案 + +- Tokio 负责异步任务调度网络任务和有界队列 +- RocksDB 负责权威数据持久化精确去重抓取状态和索引状态 +- Tantivy 负责可重建的全文搜索过滤排序和结果聚合 +- Axum 负责 HTTP 搜索接口详情接口和运行状态接口 +- BLAKE3 负责计算规范化内容结构指纹 +- Serde 负责配置领域对象和接口数据的序列化 +- Tracing 负责结构化日志和故障定位 + +RocksDB 是唯一权威数据源 + +Tantivy 索引必须能够从 RocksDB 完整重建 + +## 数据处理流程 + +1. DHT 采集器发现 infohash +2. 内存近期缓存过滤高频重复 +3. RocksDB 精确判断 infohash 是否已处理 +4. 未处理的 infohash 进入有界 Metadata 下载队列 +5. Metadata 完成校验和规范化后计算内容指纹 +6. 使用 RocksDB WriteBatch 原子保存元数据去重映射和待索引状态 +7. 后台索引任务批量写入 Tantivy +8. Tantivy 提交成功后将记录状态更新为已索引 +9. Axum 只通过搜索和存储抽象读取数据 + +所有任务队列必须有明确容量并在队列满时产生背压 + +禁止通过无限队列维持表面吞吐 + +## 去重规则 + +第一层以 infohash 做精确去重并阻止相同 Metadata 被重复下载 + +第二层根据规范化文件路径和文件大小计算 content key 将不同 infohash 的相同内容聚合展示 + +模糊名称相似度只用于搜索结果聚合不得直接删除数据 + +Bloom Filter 只能作为前置加速结构不得作为最终去重依据 + +再次发现已有 infohash 时只更新最后发现时间和发现次数 + +## 持久化和恢复 + +程序必须支持间歇运行和跨重启恢复 + +正常关闭时先停止接收新任务再排空或持久化队列最后提交搜索索引 + +异常退出后依靠 RocksDB WAL 恢复已提交数据 + +每条搜索文档必须具有待索引和已索引状态以便启动后补建索引 + +内存队列不得成为任何权威状态的唯一保存位置 + +数据目录必须通过配置指定且不得依赖当前工作目录 + +## 代码边界 + +- `crawler` 只负责协调 DHT 事件和 Metadata 下载 +- `domain` 只定义领域模型规范化规则和内容指纹 +- `storage` 只负责 RocksDB 数据布局原子写入查询和恢复 +- `search` 只负责 Tantivy schema 文档转换索引和查询 +- `api` 只负责 HTTP 协议参数校验和响应转换 +- `config` 只负责读取校验和暴露配置 +- `telemetry` 只负责日志指标和运行观测 +- `shutdown` 只负责关闭信号和优雅退出协调 + +模块之间通过明确的数据结构和 trait 通信不得跨层直接访问内部实现 + +## 资源约束 + +- Metadata 下载并发必须可配置 +- 单个 Metadata 大小文件数量和路径长度必须有限制 +- RocksDB 写入使用批处理并明确控制 block cache +- Tantivy 写入使用批量提交并明确控制 IndexWriter 内存预算 +- 日志不得输出完整 Metadata 或大文件列表 +- 长期运行的集合必须有容量上限过期规则或磁盘持久化方案 + +## 开发规则 + +新业务代码写入 `dht-search` + +项目阶段任务完成状态和验收标准统一维护在根目录 `TODOS.md` + +需求实现或技术决策发生变化时必须同步更新 `TODOS.md` + +`dht-crawler` 只接受可复用的 DHT 基础能力不得包含数据库搜索接口或部署逻辑 + +每个 Rust 文件顶部必须使用中文行注释描述该文件的功能边界且注释行尾不添加标点 + +新增行为必须包含与风险相称的测试 + +修改持久化 key 或编码格式时必须提供兼容或迁移方案 + +不得把 `opencodes` 下的参考项目纳入 workspace 或修改其内容 + +## 构建准备 + +应用骨架默认不启用 RocksDB 原生构建 + +实现存储层时通过 `rocksdb-storage` feature 启用 RocksDB + +Windows 构建 RocksDB 前需要安装 LLVM 并确保 `LIBCLANG_PATH` 指向包含 `libclang.dll` 的目录 + +启用后的检查命令为 `cargo check -p dht-search --features rocksdb-storage` diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..7db45ad --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,29 @@ +[workspace] +members = ["dht-crawler", "dht-search"] +exclude = ["opencodes"] +resolver = "3" + +[workspace.package] +edition = "2024" +authors = ["chuan "] +license = "MIT" +repository = "https://git.pchuan.top/tools/dht" + +[workspace.dependencies] +serde = { version = "1.0", features = ["derive"] } +thiserror = "2.0.19" +tokio = "1.53" +tokio-util = "0.7" +tracing = "0.1" +tracing-subscriber = "0.3" + +[profile.release] +opt-level = 3 +lto = "fat" +codegen-units = 1 +panic = "abort" +strip = true + +[profile.dev] +opt-level = 0 +debug = true diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..ecb2e4a --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024 dht-crawler contributors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..67ea47c --- /dev/null +++ b/README.md @@ -0,0 +1,21 @@ +# DHT 元数据搜索服务 + +这是一个用于发现持久化索引和搜索 BitTorrent DHT 元数据的 Rust workspace + +## 目录 + +```text +dht-search/ 最终运行的采集存储搜索和接口应用 +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..067c99f --- /dev/null +++ b/TODOS.md @@ -0,0 +1,294 @@ +# 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] 建立带时间衰减的 DHT 活跃度分数和用户可读等级 +- [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 +- [ ] 增加请求延迟错误率和并发指标 +- [x] 增加搜索详情字段和按需验证入队 API 端到端测试 +- [ ] 增加其余 API 单元测试和端到端测试 + +### 验收标准 + +- [x] API 能搜索真实采集数据 +- [x] 非法参数返回稳定的客户端错误 +- [x] 搜索查询在独立阻塞任务执行不会阻塞异步 worker +- [x] 健康检查能区分进程存活和服务可用 + +## 阶段五 质量过滤和重复内容控制 + +### 目标 + +减少垃圾数据和重复展示同时避免不可恢复的误删 + +### 任务 + +- [x] 将种子可用性定义为最近通过 DHT 找到并完成 BitTorrent 握手 +- [x] 区分 Metadata 结构有效和 swarm 当前可用性 +- [x] 实现未验证活跃和可能失效三态模型 +- [x] 实现详情高优先级和搜索普通优先级的仅按需验证 +- [x] 使用持久化有界验证队列租约恢复去重和失败退避 +- [x] Metadata 与验证握手共享 TCP 建连总预算 +- [x] RocksDB schema v1 到 v2 可恢复迁移并触发安全重建索引 +- [x] 搜索和详情接口返回热度与可用性数据 +- [x] `/stats` 返回验证队列发现握手成功失败和拒绝指标 + +- [ ] 统计真实数据的 infohash 重复率和内容重复率 +- [ ] 定义可配置的名称路径扩展名和大小过滤规则 +- [ ] 定义 Metadata 最大大小文件数和路径长度限制 +- [ ] 识别空名称异常路径大小溢出和文件数量攻击 +- [ ] 设计可解释的名称标准化规则 +- [ ] 为模糊相似结果生成聚合候选但不自动删除 +- [ ] 支持黑名单规则版本和命中原因 +- [ ] 保留被过滤记录的计数指标但避免保存大内容 +- [ ] 增加误判测试和边界数据集 + +### 验收标准 + +- [x] 搜索和详情响应不等待 DHT 或 Peer 网络验证 +- [x] 进程重启后已接受的验证任务能够通过租约恢复 +- [x] 一次验证失败不会删除记录或标记为绝对失效 +- [x] 旧数据库记录能够迁移并重新建立搜索文档 + +- [ ] 精确重复不会重复下载和重复展示 +- [ ] 内容重复可以折叠并保留全部 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 成功率 +- [ ] 统计按需验证的 Peer 发现率握手成功率和平均验证耗时 +- [x] 完成首轮真实按需验证并确认旧记录两次握手均成功更新为活跃 +- [x] 验证新 Metadata 记录直接继承成功来源 Peer 的活跃状态 +- [x] 验证启用按需可用性功能后保守预算运行三分钟并安全停止 +- [ ] 根据真实验证数据校准热度权重等级边界和失败退避时间 +- [ ] 根据公网设备长期实测设计超时率自动降速 +- [ ] 建立采集存储索引和查询基准测试 +- [ ] 记录每条元数据和每个索引文档的平均磁盘占用 +- [ ] 记录 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 旁路只影响爬虫进程 +- [ ] 更新失败时可以恢复上一版本二进制和数据 + +## 当前下一步 + +验证新可用性管线的真实运行数据并继续完善阶段三和阶段四 + +下一步运行更长时间的按需验证采样并实现时间文件数过滤排序策略和内容聚合展示 diff --git a/dht-crawler/CHANGELOG.md b/dht-crawler/CHANGELOG.md new file mode 100644 index 0000000..fcd1299 --- /dev/null +++ b/dht-crawler/CHANGELOG.md @@ -0,0 +1,48 @@ +# Changelog + +本项目遵循语义化版本。0.2.1 是包含公开 API 变更的 breaking release。 + +## Unreleased + +### Added + +- BEP-51 `sample_infohashes` 主动发现、按节点 interval/退避的采样 actor，以及有界 Hash + 去重。 +- 采样 Hash 优先查询来源节点，再通过迭代式 `get_peers` 补充 Peer。 + +### Changed + +- 主动 `get_peers` ingress 改为有界排队，避免突发采样在速率预算耗尽时直接丢弃。 +- 默认主动 Peer lookup 提升到每秒 128 个、最多 256 个并发 lookup。 +- 空节点池的 Bootstrap 默认改为 30 秒重试、每轮最多 16 个端点，降低坏 DNS + 地址导致冷启动停滞的概率。 + +## 0.2.1 - 2026-07-30 + +### Breaking changes + +- 将 `DHTOptions` 的 Metadata 和 crawl 参数改为嵌套结构：`MetadataOptions`、 + `CrawlOptions`、`RateLimitOptions`、`PoolOptions`、`BootstrapOptions`、 + `TargetOptions`、`SchedulerOptions`。 +- 删除旧的 `metadata_timeout`、`max_metadata_queue_size`、 + `max_metadata_worker_count`、`node_queue_capacity` 等扁平字段。 +- 删除旧 active/candidate frontier 和 sharded queue 实现，改用单所有者严格 FIFO + 节点池、recent-probe set 与 responsive-node ring。 +- Metadata 调度改为有界、按 InfoHash 去重、最多三个 Peer、60 秒 freshness TTL。 + +### Added + +- 独立的主动爬取 QPS、新目标、节点替换、回复包/字节、单来源回复、总在途和子网在途限制。 +- 根据 Metadata 队列压力自动降低实际 `find_node` QPS。 +- Bootstrap 来源退避、低水位触发和响应节点快照。 +- `on_torrent_with_ack`、`on_metadata_fetch_complete`、 + `MetadataFetchCompletionStatus` 和真实 Peer `attempts`。 +- 按 `SocketAddr` 缓存 Metadata Peer 的 timeout/connect failure。 +- `DhtRuntimeStats::snapshot()`、`observability_snapshot()` 和三组固定桶直方图。 +- DHT、UDP、节点池、Metadata scheduler/fetcher 的低基数 Prometheus 指标。 + +### Changed + +- Metadata timeout 现在覆盖连接、握手、传输、SHA1 和解析的完整 Peer 尝试。 +- UDP ingress、crawl events 和 Metadata queues 全部有界，并暴露 drop/depth 指标。 +- DHT 回复增加总包、总字节、单来源限流，以及 `ping`/`get_peers` 10% 保底预算。 diff --git a/dht-crawler/Cargo.toml b/dht-crawler/Cargo.toml new file mode 100644 index 0000000..c3caa8b --- /dev/null +++ b/dht-crawler/Cargo.toml @@ -0,0 +1,54 @@ +# 定义可复用 DHT 爬虫基础库的包元数据依赖和功能开关 + +[package] +name = "dht-crawler" +version = "0.2.1" +edition.workspace = true +authors.workspace = true +description = "高性能的 Rust DHT 爬虫基础库" +license.workspace = true +documentation = "https://docs.rs/dht-crawler" +repository.workspace = true +keywords = ["dht", "bittorrent", "crawler", "p2p", "torrent"] +categories = ["network-programming", "asynchronous"] +readme = "README.md" + +[lib] +name = "dht_crawler" +crate-type = ["rlib"] + +[dependencies] +ahash = "0.8" +arc-swap = "1.7" +async-channel = "2.5.0" +bytes = "1.0" +crossbeam-queue = "0.3" +hex = "0.4" +log = "0.4" +metrics = { version = "0.24", optional = true } +rand = "0.10.2" +rbit = "0.2" +serde.workspace = true +serde_bencode = "0.2" +serde_bytes = "0.11.19" +sha1 = "0.11.0" +socket2 = { version = "0.6.5", features = ["all"] } +thiserror.workspace = true +tokio = { workspace = true, features = ["rt", "rt-multi-thread", "net", "sync", "time", "macros"] } +tokio-util.workspace = true + +[dev-dependencies] +metrics-exporter-prometheus = { version = "0.18", default-features = false, features = ["http-listener"] } +mimalloc = "0.1" +tokio = { workspace = true, features = ["signal"] } +tracing.workspace = true +tracing-subscriber = { workspace = true, features = ["env-filter", "fmt", "tracing-log"] } + +[features] +default = [] +metrics = ["dep:metrics"] +mimalloc = [] + +[[example]] +name = "dht_crawler_example" +path = "examples/main.rs" diff --git a/dht-crawler/README.md b/dht-crawler/README.md new file mode 100644 index 0000000..6dbc62a --- /dev/null +++ b/dht-crawler/README.md @@ -0,0 +1,269 @@ +# dht-crawler + +[![Crates.io](https://img.shields.io/crates/v/dht-crawler.svg)](https://crates.io/crates/dht-crawler) +[![Documentation](https://docs.rs/dht-crawler/badge.svg)](https://docs.rs/dht-crawler) +[![License](https://img.shields.io/crates/l/dht-crawler.svg)](../LICENSE) + +基于 Rust 和 Tokio 的 BitTorrent DHT 爬虫库。它参与 BEP-5 DHT 网络，通过 BEP-51 +`sample_infohashes` 主动发现 InfoHash，也接收 `announce_peer`，并通过 BEP-9 +`ut_metadata` 获取、校验和解析 torrent 元数据。 + +`dht-crawler` 提供： + +- IPv4、IPv6 和双栈 DHT； +- 主动节点发现、BEP-51 InfoHash 采样与 `get_peers` 查询； +- 有界、去重的 Metadata 下载队列； +- InfoHash 过滤、异步准入、结果交付和完成通知； +- 默认可用的运行时统计，以及可选的 `metrics` 集成。 + +## 安装 + +```bash +cargo add dht-crawler +cargo add tokio --features rt-multi-thread,macros,signal +``` + +或在 `Cargo.toml` 中添加： + +```toml +[dependencies] +dht-crawler = "0.2" +tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal"] } +``` + +## 快速开始 + +```rust +use dht_crawler::prelude::*; + +#[tokio::main] +async fn main() -> Result<()> { + let server = DHTServer::new(DHTOptions { + port: 6881, + netmode: NetMode::Ipv4Only, + ..Default::default() + }) + .await?; + + server.on_torrent(|torrent| { + println!( + "{} {} {}", + torrent.info_hash, + torrent.name, + torrent.format_size() + ); + }); + + server.on_error(|error| { + eprintln!("DHT runtime error: {error}"); + }); + + let shutdown = server.clone(); + tokio::spawn(async move { + if tokio::signal::ctrl_c().await.is_ok() { + shutdown.shutdown(); + } + }); + + // 一直运行，直到 shutdown() 被调用。 + server.start().await +} +``` + +运行仓库中的完整示例： + +```bash +cargo run --release --example dht_crawler_example +``` + +## 核心 API + +`DHTServer` 是主要入口： + +| API | 用途 | +|---|---| +| `DHTServer::new(options)` | 校验配置、绑定 UDP Socket 并创建内部管道 | +| `start().await` | 启动 DHT 与爬取任务，等待 `shutdown()` | +| `shutdown()` | 停止 UDP、爬取和 Metadata 任务；可重复调用 | +| `filter(callback)` | 在 InfoHash 进入队列前执行同步过滤 | +| `on_metadata_fetch(callback)` | 在第一次 Peer 下载前执行异步准入 | +| `on_torrent(callback)` | 接收已校验的 `TorrentInfo` | +| `on_torrent_with_ack(callback)` | 接收结果并显式确认是否接受交付 | +| `on_metadata_fetch_complete(callback)` | 接收已准入任务的最终状态 | +| `on_error(callback)` | 接收运行期错误 | +| `runtime_stats()` | 获取可复制的运行时统计句柄 | + +同类回调重复注册时，新回调会替换旧回调。 + +### 过滤与准入 + +`filter` 是同步的早期过滤器，适合拦截已处理过的 InfoHash： + +```rust +server.filter(|info_hash| !already_exists(info_hash)); +``` + +`on_metadata_fetch` 是异步准入回调，在实际连接 Peer 前调用： + +```rust +server.on_metadata_fetch(|info_hash| async move { + should_download(&info_hash).await +}); +``` + +返回 `false` 会终止任务，不下载 Metadata，也不会触发 torrent 或 completion 回调。 +未注册准入回调时默认允许下载。 + +### 交付确认 + +不需要确认下游是否接收时使用 `on_torrent`。需要确认下游是否成功接收时使用 +`on_torrent_with_ack`： + +```rust +server.on_torrent_with_ack(|torrent| { + output.try_send(torrent).is_ok() +}); + +server.on_metadata_fetch_complete(|completion| { + println!( + "{}: {:?}, attempts={}", + completion.info_hash, + completion.status, + completion.attempts + ); +}); +``` + +完成状态： + +| 状态 | 含义 | +|---|---| +| `Accepted` | Metadata 下载成功，结果已被回调接受 | +| `FetchFailed` | 所有可用 Peer 尝试均失败 | +| `DeliveryRejected` | Metadata 下载成功，但结果未被回调接受 | + +`attempts` 只统计实际发起的 Peer 网络请求。异步准入拒绝不会产生 completion 事件。 + +## 配置 + +大多数调用方可以从 `DHTOptions::default()` 开始，只覆盖监听方式和容量限制： + +```rust +let options = DHTOptions { + port: 6881, + netmode: NetMode::DualStack, + hash_queue_capacity: 20_000, + metadata: MetadataOptions { + timeout_secs: 5, + max_queue_size: 20_000, + max_worker_count: 8, + max_connects_per_second: 2, + ..Default::default() + }, + crawl: CrawlOptions { + rate_limit: RateLimitOptions { + max_find_node_rate_per_sec: 6, + max_in_flight: 12, + ..Default::default() + }, + ..Default::default() + }, + ..Default::default() +}; +``` + +配置分组： + +| 类型 | 控制内容 | +|---|---| +| `DHTOptions` | 监听端口、网络模式、顶层队列和主动 UDP 查询总预算 | +| `MetadataOptions` | 下载超时、队列、并发、每秒 TCP 建连和失败 Peer 缓存 | +| `PeerLookupOptions` | 主动 `get_peers` 的速率与并发 | +| `SampleInfohashesOptions` | BEP-51 采样速率、并发、超时、退避和 Hash 去重容量 | +| `RateLimitOptions` | `find_node`、在途请求和 UDP 回复预算 | +| `PoolOptions` | 节点池、最近探测记录和响应节点缓存 | +| `BootstrapOptions` | Bootstrap 节点与失败退避 | +| `TargetOptions` | 主动爬取目标生成策略 | +| `SchedulerOptions` | 内部事件队列、批处理与快照限制 | + +完整字段和默认值以 [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 端口； +- 空节点池默认每 30 秒重新尝试 Bootstrap，每轮最多使用 16 个已解析端点； +- `MetadataOptions::timeout_secs` 是单个 Peer 尝试的端到端期限； +- `PeerLookupOptions::max_lookups_per_second = 0` 会关闭主动 `get_peers`； +- `SampleInfohashesOptions::max_queries_per_second = 0` 会关闭 BEP-51 主动采样； +- Metadata 和爬取队列都是有界的，容量应与下游处理能力一起调整。 + +## 数据与运行语义 + +`TorrentInfo` 包含 `info_hash`、`magnet_link`、`name`、`total_size`、`files`、 +`piece_length`、`peers` 和 `timestamp`。只有通过 SHA1 校验并成功解析的 Metadata +才会交付给 torrent 回调。 + +库使用有界队列控制内存占用。队列满或速率预算耗尽时，新事件可能被拒绝、淘汰或计入 +drop 指标。Metadata 队列按 InfoHash 去重，每个任务可尝试多个候选 Peer；连接超时和 +连接失败的 Peer 会被短期缓存，避免反复占用 worker。 + +`start()` 返回后，当前实例不能再次启动。如需重新运行，请创建新的 `DHTServer`。 + +## 可观测性 + +运行时快照无需启用 Cargo feature： + +```rust +let stats = server.runtime_stats(); +let snapshot = stats.snapshot(); + +println!( + "nodes={} metadata={}/{} workers={}", + snapshot.node_pool_size, + snapshot.metadata_queue_depth, + snapshot.metadata_queue_max, + snapshot.metadata_in_flight, +); +``` + +`observability_snapshot()` 提供 UDP、查询、队列、Metadata 失败原因和固定桶直方图。 +这些快照面向监控，读取时不是跨字段事务视图。 + +启用 `metrics` 后，库通过 [`metrics`](https://crates.io/crates/metrics) facade 记录 +指标，但不会安装 recorder 或启动 HTTP 服务： + +```toml +[dependencies] +dht-crawler = { version = "0.2", features = ["metrics"] } +``` + +指标名称、类型、标签和单位见 [docs/metrics.md](docs/metrics.md)。 + +## Cargo features + +默认不启用任何 feature。 + +| Feature | 用途 | +|---|---| +| `metrics` | 通过 `metrics` facade 记录指标 | +| `mimalloc` | 将 mimalloc 注册为全局分配器 | + +启用 `mimalloc` 前，请确认最终二进制没有注册其他全局分配器。 + +## 开发 + +```bash +cargo fmt --all --check +cargo check --all-targets --all-features +cargo clippy --all-targets --all-features -- -D warnings +cargo test --all-features +cargo doc --no-deps --all-features +``` + +## 许可证 + +[MIT](LICENSE) diff --git a/dht-crawler/docs/metrics.md b/dht-crawler/docs/metrics.md new file mode 100644 index 0000000..fdcea15 --- /dev/null +++ b/dht-crawler/docs/metrics.md @@ -0,0 +1,110 @@ +# dht-crawler 指标参考 + +启用 Cargo feature `metrics` 后，库通过 `metrics` facade 记录以下指标。库不安装 +recorder、不监听端口，也不依赖任何特定导出协议；Prometheus exporter 应由最终应用安装。 + +所有 `*_total` 都是进程生命周期累计 counter。Gauge 是当前值。Histogram 的记录值 +使用下表标出的单位。 + +## UDP 与 KRPC + +| 指标 | 类型 | 标签 | 单位/含义 | +|---|---|---|---| +| `dht_udp_bytes_received_total` | counter | — | Socket 接收字节 | +| `dht_udp_packets_received_total` | counter | `status=ok|dropped_size|dropped_magic|queue_full` | UDP ingress 结果 | +| `dht_udp_bytes_sent_total` | counter | — | 成功发送字节 | +| `dht_udp_packets_sent_total` | counter | `type=query|response` | 成功发送包 | +| `dht_udp_query_size_bytes` | histogram | — | find_node query 编码长度，bytes | +| `dht_messages_processed_total` | counter | `type=q|r|e|unknown` | 成功解析的 KRPC 消息类型 | +| `dht_messages_parse_error_total` | counter | — | bencode/KRPC 解析失败 | +| `dht_queries_total` | counter | `q=ping|find_node|get_peers|announce_peer|vote|other_or_invalid` | 入站查询类型 | +| `dht_udp_responses_dropped_total` | counter | `reason=rate_limit` | 最终未发送的限流回复 | +| `dht_udp_responses_priority_reserved_total` | counter | `query=ping|get_peers` | 使用 10% 保底预算的回复 | + +`dht_udp_bytes_received_total` 包含后续被判定为 invalid/queue-full 的 Datagram；发送侧只在 +`send_to` 成功后累计。 + +## 主动爬取与节点池 + +| 指标 | 类型 | 标签 | 含义 | +|---|---|---|---| +| `dht_node_pool_size` | gauge | — | 当前 FIFO 节点数 | +| `dht_node_pool_oldest_age_seconds` | gauge | — | FIFO 最老节点年龄 | +| `dht_node_pool_admissions_total` | counter | — | 新准入节点 | +| `dht_node_pool_replacements_total` | counter | — | 满池替换 | +| `dht_node_pool_dropped_total` | counter | `reason=duplicate|rate_limit|invalid` | 节点拒绝原因 | +| `dht_find_node_in_flight` | gauge | — | 当前在途 find_node | +| `dht_find_node_effective_rate_per_second` | gauge | — | Metadata 压力调整后的实际预算 | +| `dht_crawl_queries_sent_total` | counter | `kind=new|revisit|bootstrap` | 已交给 egress 的查询用途；发送失败另计 | +| `dht_find_node_responses_total` | counter | — | 与 pending transaction 匹配的回复 | +| `dht_find_node_response_unmatched_total` | counter | — | 无匹配 pending 的回复 | +| `dht_find_node_timeouts_total` | counter | — | pending 超时 | +| `dht_find_node_send_failures_total` | counter | — | UDP query 发送失败 | +| `dht_crawl_events_dropped_total` | counter | `kind=discovered|response` | 有界 actor channel 丢弃 | +| `dht_metadata_queue_pressure_ratio` | gauge | — | Metadata depth/capacity，范围 0..1 | + +actor 每秒把内部增量 flush 到 counter，因此 exporter 看到的 counter 可能最多延迟约一秒。 + +## BEP-51 InfoHash 采样 + +| 指标 | 类型 | 标签 | 含义 | +|---|---|---|---| +| `dht_sample_infohashes_queries_total` | counter | — | 已发送的 BEP-51 查询 | +| `dht_sample_infohashes_responses_total` | counter | — | 匹配的 BEP-51 响应 | +| `dht_sample_infohashes_timeouts_total` | counter | — | 超时的 BEP-51 请求 | +| `dht_sample_infohashes_hashes_total` | counter | — | 已接受并送往 Peer lookup 的新 Hash | +| `dht_sample_infohashes_in_flight` | gauge | — | 当前在途采样请求 | +| `dht_sample_infohashes_dropped_total` | counter | `reason=response_queue_full|peer_lookup_queue_full` | 有界队列丢弃 | + +## announce 与 Metadata ingress + +| 指标 | 类型 | 标签 | 含义 | +|---|---|---|---| +| `dht_announce_peer_blocked_total` | counter | `reason=invalid_token|filtered` | announce 拒绝原因 | +| `dht_info_hashes_discovered_total` | counter | — | token/hash/filter 校验通过的 InfoHash | +| `dht_metadata_ingress_dropped_total` | counter | `reason=queue_full` | Hash ingress 满导致的丢弃 | + +## Metadata scheduler + +| 指标 | 类型 | 标签 | 单位/含义 | +|---|---|---|---| +| `dht_metadata_queue_depth` | gauge | — | Pending Hash 数 | +| `dht_metadata_in_flight` | gauge | — | 当前 job 数 | +| `dht_metadata_queue_events_total` | counter | `result=inserted|deduplicated|evicted_oldest|stale|expired` | 队列事件 | +| `dht_metadata_queue_wait_seconds` | histogram | — | Hash 从最近发现到首次分派的秒数 | +| `dht_metadata_jobs_dispatched_total` | counter | — | 分派 job 数 | +| `dht_metadata_jobs_completed_total` | counter | `result=accepted|fetch_failed|delivery_rejected|gate_rejected` | job 终态 | +| `dht_metadata_worker_join_error_total` | counter | — | worker task join 失败 | +| `dht_metadata_completion_callback_panics_total` | counter | — | 完成回调 panic | + +## Metadata Peer 下载 + +| 指标 | 类型 | 标签 | 单位/含义 | +|---|---|---|---| +| `dht_metadata_fetch_attempts_total` | counter | — | 实际 Peer 尝试 | +| `dht_metadata_peer_attempts_total` | counter | — | 与 fetch attempts 相同的 Peer 尝试计数 | +| `dht_metadata_fetch_success_total` | counter | — | 成功下载和解析 | +| `dht_metadata_fetch_result_total` | counter | `result=success|failed|timeout` | Peer 尝试结果 | +| `dht_metadata_fetch_fail_total` | counter | `reason=timeout|send_error|size_limit|sha1_mismatch|parse_error` | 详细失败原因 | +| `dht_metadata_connection_result_total` | counter | `result=success|failed` | TCP/BitTorrent connect 结果 | +| `dht_metadata_handshake_result_total` | counter | `result=success|no_extension_support` | 扩展能力/最终校验结果 | +| `dht_metadata_fetch_duration_seconds` | histogram | — | 端到端 Peer 尝试秒数 | +| `dht_metadata_size_bytes` | histogram | — | 完整 bencoded info payload 字节数 | +| `dht_metadata_bytes_downloaded_total` | counter | — | 收到的 Metadata piece 数据字节数 | +| `dht_metadata_peer_failure_cache_hits_total` | counter | `reason=timeout|connect_failed` | 坏 Peer 缓存命中 | +| `dht_metadata_peer_failure_cache_inserts_total` | counter | `reason=timeout|connect_failed` | 坏 Peer 缓存写入 | +| `dht_metadata_peer_failure_cache_entries` | gauge | — | 当前缓存条目数 | + +`dht_metadata_fetch_result_total{result="failed"}` 汇总非 timeout 的失败，不适合单独用于 +分析具体原因；详细原因应结合 `fetch_fail`、connection 和 handshake 指标。 + +## 与原子快照的关系 + +`DhtRuntimeStats` 始终可用，与 `metrics` feature 无关： + +- `snapshot()` 提供队列、节点池、crawl、Peer 和 UDP 运行状态。 +- `observability_snapshot()` 提供 UDP 字节/包、入站查询分类、announce、节点准入、 + Metadata 失败分类、failure cache 分类和固定桶。 + +两套出口在同一事件点更新，但读取时都不是跨字段事务快照；短时间内可能相差一个并发 +事件，Prometheus 的 crawl actor counter 还可能有最多约一秒 flush 延迟。 diff --git a/dht-crawler/examples/main.rs b/dht-crawler/examples/main.rs new file mode 100644 index 0000000..280e12a --- /dev/null +++ b/dht-crawler/examples/main.rs @@ -0,0 +1,149 @@ +#[cfg(feature = "mimalloc")] +#[global_allocator] +static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; + +use dht_crawler::prelude::*; +#[cfg(feature = "metrics")] +use metrics_exporter_prometheus::PrometheusBuilder; +#[cfg(feature = "metrics")] +use std::net::SocketAddr; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use tracing_subscriber::EnvFilter; +#[tokio::main] +async fn main() -> Result<()> { + let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")); + + tracing_subscriber::fmt() + .with_env_filter(filter) + .with_ansi(true) + .init(); + + // 初始化 Prometheus metrics 导出器 + #[cfg(feature = "metrics")] + { + let addr: SocketAddr = "0.0.0.0:9000" + .parse() + .map_err(|e| DHTError::Init(format!("无效的 metrics 监听地址: {e}")))?; + PrometheusBuilder::new() + .with_http_listener(addr) + .install() + .map_err(|e| DHTError::Init(format!("无法安装 Prometheus metrics 导出器: {e}")))?; + log::info!("📊 Prometheus metrics 导出器已启动，访问 http://localhost:9000/metrics"); + } + + let options = DHTOptions { + port: 12313, + netmode: NetMode::Ipv4Only, + metadata: MetadataOptions { + timeout_secs: 4, + max_queue_size: 10_000, + max_worker_count: 256, + ..MetadataOptions::default() + }, + ..Default::default() + }; + + // 统计计数器 + let torrent_count = Arc::new(AtomicUsize::new(0)); + let torrent_count_clone = torrent_count.clone(); + + // 🚀 初始化 DHT Server + log::info!("🔧 正在初始化 DHT Server..."); + let server = DHTServer::new(options.clone()).await?; + + log::info!("🚀 DHT Server 启动，监听端口: {}", options.port); + + // 注册错误回调，将运行时错误输出而不是 panic + server.on_error(|err| { + log::error!("DHT 运行时错误: {}", err); + }); + + // 设置 torrent 回调 + server.on_torrent(move |_torrent| { + let _count = torrent_count_clone.fetch_add(1, Ordering::Relaxed) + 1; + + // 🔇 取消打印 torrent 信息，减少日志输出 + // let total_size: u64 = torrent.files.iter().map(|f| f.size).sum(); + // let files_display = if torrent.files.len() <= 3 { + // torrent.files.iter() + // .map(|f| format!("{} ({})", f.path, format_size(f.size))) + // .collect::>() + // .join(", ") + // } else { + // format!("{}个文件", torrent.files.len()) + // }; + // + // log::info!( + // "🎉 [{}] {} ({}, {})", + // count, + // torrent.name, + // format_size(total_size), + // files_display + // ); + }); + + // 设置元数据获取前的检查回调 + server.on_metadata_fetch(|_hash| async move { true }); + + // 一个通过 gate 的 Hash 最终只会收到一次完成状态。 + server.on_metadata_fetch_complete(|completion| { + log::debug!( + "Metadata complete: hash={}, status={:?}, peer_attempts={}", + completion.info_hash, + completion.status, + completion.attempts + ); + }); + + // 启动监控任务 + let count_monitor = torrent_count.clone(); + let runtime_stats = server.runtime_stats(); + tokio::spawn(async move { + let mut interval = tokio::time::interval(std::time::Duration::from_secs(5)); + let start_time = std::time::Instant::now(); + + loop { + interval.tick().await; + let success_fetch = count_monitor.load(Ordering::Relaxed); + let uptime = start_time.elapsed().as_secs(); + let runtime = runtime_stats.snapshot(); + + // ✅ 监控：爬虫运行状态 + log::info!( + "📊 [监控] 时长: {}s | 成功抓取: ✨ {} | 节点: {} | BEP51: hash={}, resp={}, timeout={} | Lookup: peer={} | Fetch: ok={}, fail={}, connect={}, timeout={}, noext={} | Metadata: {}/{} | worker: {}", + uptime, + success_fetch, + runtime.node_pool_size, + runtime.sample_infohashes_hashes_discovered, + runtime.sample_infohashes_responses, + runtime.sample_infohashes_timeouts, + runtime.peer_lookup_peers_found, + runtime.metadata_peer_succeeded, + runtime.metadata_peer_failed, + runtime.metadata_connect_failed, + runtime.metadata_peer_timeouts, + runtime.metadata_no_extension, + runtime.metadata_queue_depth, + runtime.metadata_queue_max, + runtime.metadata_in_flight, + ); + + if uptime > 0 && success_fetch > 0 { + let speed = (success_fetch as f64) / (uptime as f64 / 60.0); + log::info!("📈 平均抓取速度: {:.2} 种子/分钟", speed); + } + } + }); + + let shutdown_server = server.clone(); + tokio::spawn(async move { + if tokio::signal::ctrl_c().await.is_ok() { + log::info!("收到 Ctrl-C，正在停止 DHT Server"); + shutdown_server.shutdown(); + } + }); + + server.start().await?; + Ok(()) +} diff --git a/dht-crawler/src/addr.rs b/dht-crawler/src/addr.rs new file mode 100644 index 0000000..42a79d1 --- /dev/null +++ b/dht-crawler/src/addr.rs @@ -0,0 +1,71 @@ +use crate::types::NetMode; +use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; + +pub(crate) fn addr_allowed_by_netmode(addr: &SocketAddr, netmode: NetMode) -> bool { + match netmode { + NetMode::Ipv4Only => addr.is_ipv4(), + NetMode::Ipv6Only => addr.is_ipv6(), + NetMode::DualStack => true, + } +} + +pub(crate) fn is_valid_node_addr(addr: &SocketAddr) -> bool { + if addr.port() == 0 { + return false; + } + + match addr.ip() { + IpAddr::V4(ip) => is_valid_ipv4_node_addr(ip), + IpAddr::V6(ip) => is_valid_ipv6_node_addr(ip), + } +} + +fn is_valid_ipv4_node_addr(ip: Ipv4Addr) -> bool { + let octets = ip.octets(); + let is_cgnat = octets[0] == 100 && (octets[1] & 0b1100_0000) == 64; + let is_benchmark = octets[0] == 198 && (octets[1] == 18 || octets[1] == 19); + let is_reserved = octets[0] >= 240; + + !ip.is_unspecified() + && !ip.is_loopback() + && !ip.is_private() + && !ip.is_link_local() + && !ip.is_multicast() + && !ip.is_broadcast() + && !ip.is_documentation() + && !is_cgnat + && !is_benchmark + && !is_reserved +} + +fn is_valid_ipv6_node_addr(ip: Ipv6Addr) -> bool { + let octets = ip.octets(); + let is_unique_local = (octets[0] & 0xfe) == 0xfc; + let is_unicast_link_local = octets[0] == 0xfe && (octets[1] & 0xc0) == 0x80; + let is_documentation = + octets[0] == 0x20 && octets[1] == 0x01 && octets[2] == 0x0d && octets[3] == 0xb8; + + !ip.is_unspecified() + && !ip.is_loopback() + && !ip.is_multicast() + && !is_unique_local + && !is_unicast_link_local + && !is_documentation +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn invalid_node_addresses_are_filtered() { + assert!(is_valid_node_addr(&"8.8.8.8:6881".parse().unwrap())); + assert!(is_valid_node_addr( + &"[2001:4860:4860::8888]:6881".parse().unwrap() + )); + assert!(!is_valid_node_addr(&"8.8.8.8:0".parse().unwrap())); + assert!(!is_valid_node_addr(&"10.0.0.1:6881".parse().unwrap())); + assert!(!is_valid_node_addr(&"127.0.0.1:6881".parse().unwrap())); + assert!(!is_valid_node_addr(&"[fc00::1]:6881".parse().unwrap())); + } +} diff --git a/dht-crawler/src/bootstrap.rs b/dht-crawler/src/bootstrap.rs new file mode 100644 index 0000000..e9547cf --- /dev/null +++ b/dht-crawler/src/bootstrap.rs @@ -0,0 +1,221 @@ +use crate::addr::{addr_allowed_by_netmode, is_valid_node_addr}; +use crate::crawl_config::ResolvedCrawlConfig; +use crate::types::NetMode; +use std::collections::{HashMap, HashSet}; +use std::net::SocketAddr; +use std::time::{Duration, Instant}; + +pub(crate) struct BootstrapGate { + last_bootstrap: Option, +} + +impl BootstrapGate { + pub(crate) fn new() -> Self { + Self { + last_bootstrap: None, + } + } + + pub(crate) fn should_bootstrap( + &mut self, + pool_len: usize, + config: &ResolvedCrawlConfig, + now: Instant, + ) -> bool { + if config.bootstrap_max_nodes_per_round == 0 { + return false; + } + if pool_len >= config.low_watermark { + return false; + } + if let Some(last) = self.last_bootstrap + && now.checked_duration_since(last).unwrap_or_default() < config.bootstrap_interval + { + return false; + } + self.last_bootstrap = Some(now); + true + } +} + +#[derive(Default)] +struct BootstrapSourceState { + last_attempt: Option, + last_success: Option, + fail_count: u32, + backoff_until: Option, +} + +pub(crate) struct BootstrapSourcePool { + pub(crate) hosts: Vec, + states: HashMap, + backoff_base: Duration, + backoff_max: Duration, +} + +impl BootstrapSourcePool { + pub(crate) fn new(hosts: Vec, backoff_base: Duration, backoff_max: Duration) -> Self { + Self { + hosts, + states: HashMap::new(), + backoff_base, + backoff_max, + } + } + + pub(crate) fn select( + &mut self, + candidates: Vec, + max_nodes: usize, + now: Instant, + ) -> Vec { + let mut selected = Vec::with_capacity(max_nodes); + let mut seen = HashSet::with_capacity(candidates.len()); + let mut earliest_backoff: Option<(SocketAddr, Instant)> = None; + + for addr in candidates { + if !seen.insert(addr) { + continue; + } + + let state = self.states.entry(addr).or_default(); + if let Some(backoff_until) = state.backoff_until + && backoff_until > now + { + if earliest_backoff.is_none_or(|(_, current)| backoff_until < current) { + earliest_backoff = Some((addr, backoff_until)); + } + continue; + } + + selected.push(addr); + if selected.len() >= max_nodes { + return selected; + } + } + + if selected.is_empty() + && max_nodes > 0 + && let Some((addr, _)) = earliest_backoff + { + selected.push(addr); + } + selected + } + + pub(crate) fn mark_attempt(&mut self, addr: SocketAddr, now: Instant) { + self.states.entry(addr).or_default().last_attempt = Some(now); + } + + pub(crate) fn mark_success(&mut self, addr: SocketAddr, now: Instant) { + let state = self.states.entry(addr).or_default(); + state.last_success = Some(now); + state.fail_count = 0; + state.backoff_until = None; + } + + pub(crate) fn mark_timeout(&mut self, addr: SocketAddr, now: Instant) { + let state = self.states.entry(addr).or_default(); + state.fail_count = state.fail_count.saturating_add(1); + let multiplier = 1u32 + .checked_shl(state.fail_count.saturating_sub(1).min(16)) + .unwrap_or(u32::MAX); + let backoff = self + .backoff_base + .saturating_mul(multiplier) + .min(self.backoff_max); + state.backoff_until = Some(now + backoff); + } +} + +pub(crate) async fn resolve_bootstrap_nodes(hosts: &[String], netmode: NetMode) -> Vec { + let mut resolved = Vec::new(); + for host in hosts { + if let Ok(addrs) = tokio::net::lookup_host(host).await { + for addr in addrs { + if !addr_allowed_by_netmode(&addr, netmode) || !is_valid_node_addr(&addr) { + continue; + } + resolved.push(addr); + } + } + } + resolved +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::types::CrawlOptions; + + fn test_config() -> ResolvedCrawlConfig { + ResolvedCrawlConfig::from_options(&CrawlOptions::default()) + } + + #[test] + fn bootstrap_pool_backs_off_dead_sources_without_spending_quota() { + let start = Instant::now(); + let addr1: SocketAddr = "8.8.8.8:6881".parse().unwrap(); + let addr2: SocketAddr = "1.1.1.1:6881".parse().unwrap(); + let mut pool = BootstrapSourcePool::new( + vec!["example.invalid:6881".to_string()], + Duration::from_secs(300), + Duration::from_secs(3600), + ); + + pool.mark_timeout(addr1, start); + let selected = pool.select(vec![addr1, addr2], 1, start + Duration::from_secs(1)); + assert_eq!(selected, vec![addr2]); + + pool.mark_success(addr1, start + Duration::from_secs(2)); + let selected = pool.select(vec![addr1], 1, start + Duration::from_secs(3)); + assert_eq!(selected, vec![addr1]); + } + + #[test] + fn bootstrap_pool_forces_one_retry_when_all_sources_backed_off() { + let start = Instant::now(); + let addr1: SocketAddr = "8.8.8.8:6881".parse().unwrap(); + let addr2: SocketAddr = "1.1.1.1:6881".parse().unwrap(); + let mut pool = BootstrapSourcePool::new( + vec!["example.invalid:6881".to_string()], + Duration::from_secs(300), + Duration::from_secs(3600), + ); + pool.mark_timeout(addr1, start); + pool.mark_timeout(addr2, start); + + let selected = pool.select(vec![addr1, addr2], 3, start + Duration::from_secs(1)); + assert_eq!(selected.len(), 1); + } + + #[test] + fn bootstrap_pool_deduplicates_resolved_addresses() { + let start = Instant::now(); + let addr: SocketAddr = "8.8.8.8:6881".parse().unwrap(); + let mut pool = BootstrapSourcePool::new( + vec!["example.invalid:6881".to_string()], + Duration::from_secs(300), + Duration::from_secs(3600), + ); + + let selected = pool.select(vec![addr, addr], 10, start); + assert_eq!(selected, vec![addr]); + } + + #[test] + fn bootstrap_gate_uses_pool_low_water_mark() { + let start = Instant::now(); + let config = test_config(); + let mut gate = BootstrapGate::new(); + + assert!(!gate.should_bootstrap(config.low_watermark, &config, start)); + assert!(gate.should_bootstrap(0, &config, start)); + assert!(!gate.should_bootstrap( + 999, + &config, + start + config.bootstrap_interval - Duration::from_secs(1) + )); + assert!(gate.should_bootstrap(999, &config, start + config.bootstrap_interval)); + } +} diff --git a/dht-crawler/src/budget.rs b/dht-crawler/src/budget.rs new file mode 100644 index 0000000..acfb00e --- /dev/null +++ b/dht-crawler/src/budget.rs @@ -0,0 +1,180 @@ +use std::{ + sync::{Arc, Mutex}, + time::{Duration, Instant}, +}; + +/// Single-owner token bucket. It deliberately contains no atomics or locks. +pub(crate) struct RateBucket { + rate_per_sec: f64, + capacity: f64, + tokens: f64, + 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, + burst: u32, + initially_full: bool, + now: Instant, + ) -> Self { + Self::new( + f64::from(rate_per_sec), + f64::from(burst), + initially_full, + now, + ) + } + + pub(crate) fn per_minute( + rate_per_minute: u32, + burst: u32, + initially_full: bool, + now: Instant, + ) -> Self { + Self::new( + f64::from(rate_per_minute) / 60.0, + f64::from(burst), + initially_full, + now, + ) + } + + fn new(rate_per_sec: f64, capacity: f64, initially_full: bool, now: Instant) -> Self { + let capacity = if rate_per_sec <= 0.0 { + 0.0 + } else { + capacity.max(1.0) + }; + Self { + rate_per_sec, + capacity, + tokens: if initially_full { capacity } else { 0.0 }, + last_refill: now, + } + } + + pub(crate) fn set_per_second_rate(&mut self, rate_per_sec: u32, now: Instant) { + self.refill(now); + self.rate_per_sec = f64::from(rate_per_sec); + if self.rate_per_sec <= 0.0 { + self.tokens = 0.0; + } else { + self.tokens = self.tokens.min(self.capacity); + } + } + + pub(crate) fn try_take_one(&mut self, now: Instant) -> bool { + self.try_take_exact(1, now) + } + + pub(crate) fn try_take_exact(&mut self, count: usize, now: Instant) -> bool { + if count == 0 { + return true; + } + if self.rate_per_sec <= 0.0 || self.capacity <= 0.0 { + return false; + } + self.refill(now); + if self.tokens < count as f64 { + return false; + } + self.tokens -= count as f64; + true + } + + pub(crate) fn try_take(&mut self, max: usize, now: Instant) -> usize { + if max == 0 || self.rate_per_sec <= 0.0 || self.capacity <= 0.0 { + return 0; + } + self.refill(now); + let taken = (self.tokens.floor() as usize).min(max); + self.tokens -= taken as f64; + taken + } + + pub(crate) fn refund_one(&mut self) { + self.refund(1); + } + + pub(crate) fn refund(&mut self, count: usize) { + self.tokens = (self.tokens + count as f64).min(self.capacity); + } + + fn refill(&mut self, now: Instant) { + let elapsed = now + .checked_duration_since(self.last_refill) + .unwrap_or(Duration::ZERO) + .as_secs_f64(); + if elapsed > 0.0 { + self.tokens = (self.tokens + elapsed * self.rate_per_sec).min(self.capacity); + self.last_refill = now; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn second_bucket_is_smooth_and_bounded() { + let start = Instant::now(); + let mut bucket = RateBucket::per_second(100, 20, false, start); + assert_eq!(bucket.try_take(100, start), 0); + assert_eq!(bucket.try_take(100, start + Duration::from_millis(100)), 10); + assert_eq!(bucket.try_take(100, start + Duration::from_secs(10)), 20); + } + + #[test] + fn minute_bucket_refills_fractionally() { + let start = Instant::now(); + 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_config.rs b/dht-crawler/src/crawl_config.rs new file mode 100644 index 0000000..e622ed4 --- /dev/null +++ b/dht-crawler/src/crawl_config.rs @@ -0,0 +1,137 @@ +use crate::types::CrawlOptions; +use std::time::Duration; + +#[derive(Debug, Clone)] +pub(crate) struct ResolvedCrawlConfig { + pub(crate) max_find_node_rate_per_sec: u32, + pub(crate) burst: u32, + pub(crate) max_in_flight: usize, + pub(crate) request_timeout: Duration, + pub(crate) max_new_destinations_per_minute: u32, + pub(crate) max_response_rate_per_sec: u32, + pub(crate) max_response_bytes_per_sec: u64, + pub(crate) max_response_rate_per_source: u32, + pub(crate) metadata_pressure_floor_percent: u8, + + pub(crate) pool_capacity: usize, + pub(crate) max_replacements_per_minute: u32, + pub(crate) recent_probe_ttl: Duration, + pub(crate) responsive_capacity: usize, + pub(crate) responsive_ttl: Duration, + pub(crate) low_watermark: usize, + pub(crate) max_in_flight_per_subnet: usize, + + pub(crate) bootstrap_nodes: Vec, + pub(crate) bootstrap_interval: Duration, + pub(crate) bootstrap_max_nodes_per_round: usize, + pub(crate) bootstrap_backoff_base: Duration, + pub(crate) bootstrap_backoff_max: Duration, + + pub(crate) random_walk_percent: u8, + pub(crate) sparse_bucket_percent: u8, + pub(crate) neighbor_sender_id: bool, + + pub(crate) priority_event_channel_capacity: usize, + pub(crate) discovery_event_channel_capacity: usize, + pub(crate) event_batch_limit: usize, + pub(crate) node_batch_limit: usize, + pub(crate) routing_snapshot_size: usize, + pub(crate) snapshot_refresh: Duration, +} + +impl ResolvedCrawlConfig { + pub(crate) fn from_options(options: &CrawlOptions) -> Self { + let capacity = options.pool.capacity.max(1); + Self { + max_find_node_rate_per_sec: options.rate_limit.max_find_node_rate_per_sec, + burst: if options.rate_limit.max_find_node_rate_per_sec == 0 { + 0 + } else { + options.rate_limit.burst.max(1) + }, + max_in_flight: options.rate_limit.max_in_flight.max(1), + request_timeout: Duration::from_secs(options.rate_limit.request_timeout_secs.max(1)), + max_new_destinations_per_minute: options.rate_limit.max_new_destinations_per_minute, + max_response_rate_per_sec: options.rate_limit.max_response_rate_per_sec, + max_response_bytes_per_sec: options.rate_limit.max_response_bytes_per_sec, + max_response_rate_per_source: options.rate_limit.max_response_rate_per_source, + metadata_pressure_floor_percent: options + .rate_limit + .metadata_pressure_floor_percent + .min(100), + + pool_capacity: capacity, + max_replacements_per_minute: options.rate_limit.max_replacements_per_minute, + recent_probe_ttl: Duration::from_secs(options.pool.recent_probe_ttl_secs.max(1)), + responsive_capacity: options.pool.responsive_capacity.max(1), + responsive_ttl: Duration::from_secs(options.pool.responsive_ttl_secs.max(1)), + low_watermark: options.pool.low_watermark.min(capacity), + max_in_flight_per_subnet: options.rate_limit.max_in_flight_per_subnet.max(1), + + bootstrap_nodes: if options.bootstrap.nodes.is_empty() { + crate::types::BootstrapOptions::default().nodes + } else { + options.bootstrap.nodes.clone() + }, + bootstrap_interval: Duration::from_secs(options.bootstrap.interval_secs), + bootstrap_max_nodes_per_round: options.bootstrap.max_nodes_per_round, + bootstrap_backoff_base: Duration::from_secs( + options.bootstrap.source_backoff_base_secs.max(1), + ), + bootstrap_backoff_max: Duration::from_secs( + options.bootstrap.source_backoff_max_secs.max(1), + ), + + random_walk_percent: options.target.random_walk_percent.min(100), + sparse_bucket_percent: options.target.sparse_bucket_percent.min(100), + neighbor_sender_id: options.target.neighbor_sender_id, + + priority_event_channel_capacity: options + .scheduler + .priority_event_channel_capacity + .max(1), + discovery_event_channel_capacity: options + .scheduler + .discovery_event_channel_capacity + .max(1), + event_batch_limit: options.scheduler.event_batch_limit.max(1), + node_batch_limit: options.scheduler.node_batch_limit.max(1), + routing_snapshot_size: options.scheduler.routing_snapshot_size.max(1), + snapshot_refresh: Duration::from_millis( + options.scheduler.snapshot_refresh_millis.max(100), + ), + } + } + + pub(crate) fn rate_for_metadata_pressure(&self, pressure: f64) -> u32 { + let max = f64::from(self.max_find_node_rate_per_sec); + if pressure < 0.80 { + return self.max_find_node_rate_per_sec; + } + let floor = max * f64::from(self.metadata_pressure_floor_percent) / 100.0; + if pressure >= 0.95 { + return floor.round() as u32; + } + let progress = ((pressure - 0.80) / 0.15).clamp(0.0, 1.0); + (max - (max - floor) * progress).round() as u32 + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn resolves_pool_and_rate_limit() { + let mut options = CrawlOptions::default(); + options.rate_limit.max_find_node_rate_per_sec = 200; + options.rate_limit.metadata_pressure_floor_percent = 25; + options.pool.capacity = 42; + let resolved = ResolvedCrawlConfig::from_options(&options); + + assert_eq!(resolved.pool_capacity, 42); + assert_eq!(resolved.rate_for_metadata_pressure(0.79), 200); + assert_eq!(resolved.rate_for_metadata_pressure(0.95), 50); + assert_eq!(resolved.rate_for_metadata_pressure(1.0), 50); + } +} diff --git a/dht-crawler/src/crawl_engine.rs b/dht-crawler/src/crawl_engine.rs new file mode 100644 index 0000000..24f68e4 --- /dev/null +++ b/dht-crawler/src/crawl_engine.rs @@ -0,0 +1,960 @@ +use crate::bootstrap::{BootstrapGate, BootstrapSourcePool, resolve_bootstrap_nodes}; +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::{ + TransactionId, bucket_index, neighbor_node_id, random_node_id, target_for_bucket, +}; +use crate::node_pool::{AdmissionOutcome, NodePool, ResponsiveReservoir, SubnetKey}; +use crate::protocol::DhtResponse; +use crate::routing_snapshot::RoutingSnapshot; +#[cfg(test)] +use crate::runtime_stats::DhtRuntimeLimits; +use crate::runtime_stats::DhtRuntimeStats; +use crate::types::{NetMode, NodeTuple}; +use ahash::AHashMap; +use arc_swap::ArcSwap; +use bytes::BytesMut; +#[cfg(feature = "metrics")] +use metrics::{counter, gauge}; +use rand::RngExt; +use std::collections::{HashMap, VecDeque}; +use std::net::SocketAddr; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; +use tokio::net::UdpSocket; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; + +const SCHEDULE_INTERVAL: Duration = Duration::from_millis(5); +const MAX_POOL_SCAN_PER_SCHEDULE: usize = 256; + +#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)] +struct PendingKey { + addr: SocketAddr, + tid: TransactionId, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ProbePurpose { + New, + Revisit, + Bootstrap, +} + +#[derive(Debug, Clone, Copy)] +struct PendingRequest { + node: NodeTuple, + purpose: ProbePurpose, + deadline: Instant, + subnet: SubnetKey, +} + +enum PriorityEvent { + Response { + remote_addr: SocketAddr, + tid: TransactionId, + response: DhtResponse, + }, + SendFailed(PendingKey), + BootstrapResolved(Vec), +} + +struct OutboundRequest { + key: PendingKey, + node: NodeTuple, + target: [u8; 20], + sender_id: [u8; 20], +} + +pub(crate) struct CrawlEngine { + config: ResolvedCrawlConfig, + priority_tx: mpsc::Sender, + discovery_tx: mpsc::Sender, + /// One-shot handoff used only by `DHTServer::start`; never touched by the crawl hot path. + receivers: Mutex, mpsc::Receiver)>>, + 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, + 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 { + config, + priority_tx, + discovery_tx, + receivers: Mutex::new(Some((priority_rx, discovery_rx))), + snapshot: Arc::new(ArcSwap::from_pointee(RoutingSnapshot::default())), + node_count: Arc::new(AtomicUsize::new(0)), + runtime_stats, + outbound_query_budget, + } + } + + pub(crate) fn route_discovered(&self, node: NodeTuple) { + let enqueue_result = self.discovery_tx.try_send(node); + self.runtime_stats.set_crawl_discovery_queue_depth( + self.discovery_tx + .max_capacity() + .saturating_sub(self.discovery_tx.capacity()), + ); + if enqueue_result.is_err() { + self.runtime_stats.crawl_event_dropped_discovered(); + #[cfg(feature = "metrics")] + counter!("dht_crawl_events_dropped_total", "kind" => "discovered").increment(1); + } + } + + pub(crate) fn route_response( + &self, + remote_addr: SocketAddr, + tid: TransactionId, + response: DhtResponse, + ) { + let enqueue_result = self.priority_tx.try_send(PriorityEvent::Response { + remote_addr, + tid, + response, + }); + self.runtime_stats.set_crawl_priority_queue_depth( + self.priority_tx + .max_capacity() + .saturating_sub(self.priority_tx.capacity()), + ); + if enqueue_result.is_err() { + self.runtime_stats.crawl_event_dropped_response(); + #[cfg(feature = "metrics")] + counter!("dht_crawl_events_dropped_total", "kind" => "response").increment(1); + } + } + + pub(crate) fn spawn( + &self, + netmode: NetMode, + local_id: [u8; 20], + sockets: &HashMap>, + metadata_queue_len: Arc, + max_metadata_queue_size: usize, + shutdown: CancellationToken, + ) { + let Some((priority_rx, discovery_rx)) = self + .receivers + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .take() + else { + return; + }; + + let mut egress_v4 = None; + let mut egress_v6 = None; + for (bind_addr, socket) in sockets.iter() { + let (tx, rx) = mpsc::channel(self.config.max_in_flight.max(1)); + spawn_egress( + socket.clone(), + rx, + self.priority_tx.clone(), + self.runtime_stats.clone(), + shutdown.clone(), + ); + if bind_addr.is_ipv4() { + egress_v4 = Some(tx); + } else { + egress_v6 = Some(tx); + } + } + + let actor = CrawlActor::new(CrawlActorInit { + config: self.config.clone(), + netmode, + local_id, + priority_rx, + discovery_rx, + priority_tx: self.priority_tx.clone(), + egress_v4, + egress_v6, + snapshot: self.snapshot.clone(), + node_count: self.node_count.clone(), + 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()); + } +} + +fn spawn_egress( + socket: Arc, + mut rx: mpsc::Receiver, + priority_tx: mpsc::Sender, + runtime_stats: DhtRuntimeStats, + shutdown: CancellationToken, +) { + tokio::spawn(async move { + let mut buffer = BytesMut::with_capacity(128); + loop { + tokio::select! { + _ = shutdown.cancelled() => break, + request = rx.recv() => { + let Some(request) = request else { break }; + if !send_find_node_query( + &request.node.addr, + &request.key.tid, + &request.target, + &request.sender_id, + &socket, + &mut buffer, + ).await { + let _ = priority_tx.try_send(PriorityEvent::SendFailed(request.key)); + runtime_stats.set_crawl_priority_queue_depth( + priority_tx + .max_capacity() + .saturating_sub(priority_tx.capacity()), + ); + } else { + runtime_stats.udp_sent(buffer.len()); + } + } + } + } + }); +} + +#[derive(Default)] +struct ActorMetrics { + admitted: u64, + replaced: u64, + duplicate: u64, + admission_limited: u64, + invalid: u64, + queries_new: u64, + queries_revisit: u64, + queries_bootstrap: u64, + responses: u64, + timeouts: u64, + send_failures: u64, + unmatched_responses: u64, +} + +impl ActorMetrics { + fn record_admission(&mut self, outcome: AdmissionOutcome) { + match outcome { + AdmissionOutcome::Admitted => self.admitted += 1, + AdmissionOutcome::Replaced => self.replaced += 1, + AdmissionOutcome::Duplicate => self.duplicate += 1, + AdmissionOutcome::RateLimited => self.admission_limited += 1, + AdmissionOutcome::Invalid => self.invalid += 1, + } + } +} + +fn record_runtime_admission(stats: &DhtRuntimeStats, outcome: AdmissionOutcome) { + match outcome { + AdmissionOutcome::Admitted => stats.node_admitted(), + AdmissionOutcome::Replaced => stats.node_replaced(), + AdmissionOutcome::Duplicate => stats.node_dropped_duplicate(), + AdmissionOutcome::RateLimited => stats.node_dropped_rate_limited(), + AdmissionOutcome::Invalid => stats.node_dropped_invalid(), + } +} + +struct CrawlActorInit { + config: ResolvedCrawlConfig, + netmode: NetMode, + local_id: [u8; 20], + priority_rx: mpsc::Receiver, + discovery_rx: mpsc::Receiver, + priority_tx: mpsc::Sender, + egress_v4: Option>, + egress_v6: Option>, + snapshot: Arc>, + node_count: Arc, + metadata_queue_len: Arc, + max_metadata_queue_size: usize, + runtime_stats: DhtRuntimeStats, + outbound_query_budget: SharedRateBudget, + shutdown: CancellationToken, +} + +struct CrawlActor { + config: ResolvedCrawlConfig, + netmode: NetMode, + local_id: [u8; 20], + priority_rx: mpsc::Receiver, + discovery_rx: mpsc::Receiver, + priority_tx: mpsc::Sender, + egress_v4: Option>, + egress_v6: Option>, + pool: NodePool, + responsive: ResponsiveReservoir, + pending: AHashMap, + pending_expiry: VecDeque<(Instant, PendingKey)>, + subnet_in_flight: AHashMap, + query_budget: RateBucket, + destination_budget: RateBucket, + bootstrap_gate: BootstrapGate, + bootstrap_pool: BootstrapSourcePool, + bootstrap_queue: VecDeque, + snapshot: Arc>, + node_count: Arc, + metadata_queue_len: Arc, + max_metadata_queue_size: usize, + next_tid: u64, + metrics: ActorMetrics, + runtime_stats: DhtRuntimeStats, + outbound_query_budget: SharedRateBudget, + shutdown: CancellationToken, +} + +impl CrawlActor { + fn new(init: CrawlActorInit) -> Self { + let CrawlActorInit { + config, + netmode, + local_id, + priority_rx, + discovery_rx, + priority_tx, + egress_v4, + egress_v6, + snapshot, + node_count, + metadata_queue_len, + max_metadata_queue_size, + runtime_stats, + outbound_query_budget, + shutdown, + } = init; + let now = Instant::now(); + let destination_burst = config.max_new_destinations_per_minute.div_ceil(60).max(1); + Self { + pool: NodePool::new( + config.pool_capacity, + config.max_replacements_per_minute, + config.recent_probe_ttl, + now, + ), + responsive: ResponsiveReservoir::new(config.responsive_capacity, config.responsive_ttl), + pending: AHashMap::with_capacity(config.max_in_flight), + pending_expiry: VecDeque::with_capacity(config.max_in_flight), + subnet_in_flight: AHashMap::new(), + query_budget: RateBucket::per_second( + config.max_find_node_rate_per_sec, + config.burst, + false, + now, + ), + destination_budget: RateBucket::per_minute( + config.max_new_destinations_per_minute, + destination_burst, + false, + now, + ), + bootstrap_pool: BootstrapSourcePool::new( + config.bootstrap_nodes.clone(), + config.bootstrap_backoff_base, + config.bootstrap_backoff_max, + ), + config, + netmode, + local_id, + priority_rx, + discovery_rx, + priority_tx, + egress_v4, + egress_v6, + bootstrap_gate: BootstrapGate::new(), + bootstrap_queue: VecDeque::new(), + snapshot, + node_count, + metadata_queue_len, + max_metadata_queue_size, + next_tid: 1, + metrics: ActorMetrics::default(), + runtime_stats, + outbound_query_budget, + shutdown, + } + } + + async fn run(mut self) { + let mut schedule_tick = tokio::time::interval(SCHEDULE_INTERVAL); + schedule_tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + let mut snapshot_tick = tokio::time::interval(self.config.snapshot_refresh); + snapshot_tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + let mut metrics_tick = tokio::time::interval(Duration::from_secs(1)); + metrics_tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + + loop { + tokio::select! { + biased; + _ = self.shutdown.cancelled() => break, + _ = schedule_tick.tick() => self.on_schedule_tick(Instant::now()), + _ = snapshot_tick.tick() => self.publish_snapshot(Instant::now()), + _ = metrics_tick.tick() => self.flush_metrics(Instant::now()), + event = self.priority_rx.recv() => { + self.runtime_stats + .set_crawl_priority_queue_depth(self.priority_rx.len()); + let Some(event) = event else { break }; + self.handle_priority(event, Instant::now()); + self.drain_events(); + } + node = self.discovery_rx.recv() => { + self.runtime_stats + .set_crawl_discovery_queue_depth(self.discovery_rx.len()); + let Some(node) = node else { break }; + self.admit(node, Instant::now()); + self.drain_events(); + } + } + } + self.runtime_stats.set_crawl_priority_queue_depth(0); + self.runtime_stats.set_crawl_discovery_queue_depth(0); + } + + fn drain_events(&mut self) { + let mut events = 1; + let mut nodes = 0; + while events < self.config.event_batch_limit && nodes < self.config.node_batch_limit { + if events % 8 == 0 + && let Ok(node) = self.discovery_rx.try_recv() + { + self.runtime_stats + .set_crawl_discovery_queue_depth(self.discovery_rx.len()); + self.admit(node, Instant::now()); + events += 1; + nodes += 1; + continue; + } + if let Ok(event) = self.priority_rx.try_recv() { + self.runtime_stats + .set_crawl_priority_queue_depth(self.priority_rx.len()); + self.handle_priority(event, Instant::now()); + events += 1; + continue; + } + if let Ok(node) = self.discovery_rx.try_recv() { + self.runtime_stats + .set_crawl_discovery_queue_depth(self.discovery_rx.len()); + self.admit(node, Instant::now()); + events += 1; + nodes += 1; + continue; + } + break; + } + } + + fn handle_priority(&mut self, event: PriorityEvent, now: Instant) { + match event { + PriorityEvent::Response { + remote_addr, + tid, + response, + } => self.handle_response(remote_addr, tid, response, now), + PriorityEvent::SendFailed(key) => { + if let Some(pending) = self.pending.remove(&key) { + self.release_pending(pending); + self.metrics.send_failures += 1; + self.runtime_stats.send_failure(); + if pending.purpose == ProbePurpose::Bootstrap { + self.bootstrap_pool.mark_timeout(pending.node.addr, now); + } + } + } + PriorityEvent::BootstrapResolved(candidates) => { + let selected = self.bootstrap_pool.select( + candidates, + self.config.bootstrap_max_nodes_per_round, + now, + ); + self.bootstrap_queue.extend(selected); + } + } + } + + fn handle_response( + &mut self, + remote_addr: SocketAddr, + tid: TransactionId, + response: DhtResponse, + now: Instant, + ) { + let key = PendingKey { + addr: remote_addr, + tid, + }; + let Some(pending) = self.pending.remove(&key) else { + self.metrics.unmatched_responses += 1; + self.runtime_stats.unmatched_response(); + return; + }; + self.release_pending(pending); + self.metrics.responses += 1; + self.runtime_stats.response(); + if pending.purpose == ProbePurpose::Bootstrap { + self.bootstrap_pool.mark_success(remote_addr, now); + } + + let mut responsive_node = pending.node; + if let Some(id) = response.id.as_ref() + && let Ok(id) = <[u8; 20]>::try_from(id.as_slice()) + { + responsive_node.id = id; + } + self.responsive.record(responsive_node, now); + + let pool = &mut self.pool; + let metrics = &mut self.metrics; + let runtime_stats = &self.runtime_stats; + for_each_response_node(&response, self.netmode, |node| { + let outcome = pool.admit(node, now); + metrics.record_admission(outcome); + record_runtime_admission(runtime_stats, outcome); + }); + self.sync_node_count(); + } + + fn admit(&mut self, node: NodeTuple, now: Instant) { + let outcome = self.pool.admit(node, now); + self.metrics.record_admission(outcome); + record_runtime_admission(&self.runtime_stats, outcome); + self.sync_node_count(); + } + + fn on_schedule_tick(&mut self, now: Instant) { + self.expire_pending(now); + self.maybe_resolve_bootstrap(now); + + let rate = self + .config + .rate_for_metadata_pressure(self.metadata_pressure()); + self.runtime_stats.set_find_node_effective_rate(rate); + self.query_budget.set_per_second_rate(rate, now); + let budget = self.query_budget.try_take(self.config.burst as usize, now); + for _ in 0..budget { + if self.pending.len() >= self.config.max_in_flight { + self.query_budget.refund_one(); + break; + } + if !self.schedule_one(now) { + self.query_budget.refund_one(); + break; + } + } + self.sync_node_count(); + } + + 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() + { + let is_new = !self.pool.contains_recent(&addr, now); + if is_new && !self.destination_budget.try_take_one(now) { + return self.schedule_revisit(now); + } + let node = NodeTuple { + id: self.local_id, + addr, + }; + if self.try_dispatch(node, ProbePurpose::Bootstrap, now) { + self.bootstrap_queue.pop_front(); + self.pool.record_probe(addr, now); + self.bootstrap_pool.mark_attempt(addr, now); + self.metrics.queries_bootstrap += 1; + self.runtime_stats.query_bootstrap(); + return true; + } + if is_new { + self.destination_budget.refund_one(); + } + } + + let scan_limit = self.pool.len().min(MAX_POOL_SCAN_PER_SCHEDULE); + for _ in 0..scan_limit { + let Some(node) = self.pool.front() else { + break; + }; + let subnet = SubnetKey::from_addr(&node.addr); + if self.subnet_count(&subnet) >= self.config.max_in_flight_per_subnet { + self.pool.rotate_front_to_back(); + continue; + } + if !self.destination_budget.try_take_one(now) { + return self.schedule_revisit(now); + } + let node = self + .pool + .take_front_for_probe(now) + .expect("front node exists"); + if self.try_dispatch(node, ProbePurpose::New, now) { + self.metrics.queries_new += 1; + self.runtime_stats.query_new(); + return true; + } + self.pool.restore_front(node, now); + self.destination_budget.refund_one(); + break; + } + self.schedule_revisit(now) + } + + fn schedule_revisit(&mut self, now: Instant) -> bool { + let Some(node) = self.responsive.next_revisit(now) else { + return false; + }; + let was_recent = self.pool.contains_recent(&node.addr, now); + if !was_recent && !self.destination_budget.try_take_one(now) { + return false; + } + if self.try_dispatch(node, ProbePurpose::Revisit, now) { + self.pool.record_probe(node.addr, now); + self.metrics.queries_revisit += 1; + self.runtime_stats.query_revisit(); + return true; + } + if !was_recent { + self.destination_budget.refund_one(); + } + false + } + + fn try_dispatch(&mut self, node: NodeTuple, purpose: ProbePurpose, now: Instant) -> bool { + if self.pending.len() >= self.config.max_in_flight { + return false; + } + let subnet = SubnetKey::from_addr(&node.addr); + if self.subnet_count(&subnet) >= self.config.max_in_flight_per_subnet { + return false; + } + let tx = if node.addr.is_ipv4() { + self.egress_v4.clone() + } else { + self.egress_v6.clone() + }; + let Some(tx) = tx else { + return false; + }; + let Ok(permit) = tx.try_reserve() else { + return false; + }; + + let tid = self.next_tid.to_be_bytes(); + self.next_tid = self.next_tid.wrapping_add(1).max(1); + let key = PendingKey { + addr: node.addr, + tid, + }; + let deadline = now + self.config.request_timeout; + self.pending.insert( + key, + PendingRequest { + node, + purpose, + deadline, + subnet, + }, + ); + self.pending_expiry.push_back((deadline, key)); + *self.subnet_in_flight.entry(subnet).or_insert(0) += 1; + self.runtime_stats + .set_find_node_in_flight(self.pending.len()); + + let sender_id = if self.config.neighbor_sender_id { + let generated = neighbor_node_id(&node.id, &self.local_id); + generated + .as_slice() + .try_into() + .expect("neighbor id is always 20 bytes") + } else { + self.local_id + }; + permit.send(OutboundRequest { + key, + node, + target: self.choose_target(&node), + sender_id, + }); + true + } + + fn choose_target(&self, node: &NodeTuple) -> [u8; 20] { + let total = self + .config + .random_walk_percent + .saturating_add(self.config.sparse_bucket_percent) + .max(1); + if rand::rng().random_range(0..total) < self.config.sparse_bucket_percent { + target_for_bucket(&self.local_id, bucket_index(&node.id, &self.local_id)) + } else { + random_node_id() + } + } + + fn expire_pending(&mut self, now: Instant) { + while let Some((deadline, key)) = self.pending_expiry.front().copied() { + if deadline > now { + break; + } + self.pending_expiry.pop_front(); + let should_remove = self + .pending + .get(&key) + .is_some_and(|pending| pending.deadline == deadline); + if !should_remove { + continue; + } + let pending = self.pending.remove(&key).expect("pending entry exists"); + self.release_pending(pending); + self.metrics.timeouts += 1; + self.runtime_stats.timeout(); + if pending.purpose == ProbePurpose::Bootstrap { + self.bootstrap_pool.mark_timeout(pending.node.addr, now); + } + } + } + + fn release_pending(&mut self, pending: PendingRequest) { + if let Some(count) = self.subnet_in_flight.get_mut(&pending.subnet) { + *count = count.saturating_sub(1); + if *count == 0 { + self.subnet_in_flight.remove(&pending.subnet); + } + } + self.runtime_stats + .set_find_node_in_flight(self.pending.len()); + } + + fn subnet_count(&self, subnet: &SubnetKey) -> usize { + self.subnet_in_flight.get(subnet).copied().unwrap_or(0) + } + + fn maybe_resolve_bootstrap(&mut self, now: Instant) { + if !self.bootstrap_queue.is_empty() + || !self + .bootstrap_gate + .should_bootstrap(self.pool.len(), &self.config, now) + { + return; + } + let hosts = self.bootstrap_pool.hosts.clone(); + let netmode = self.netmode; + let priority_tx = self.priority_tx.clone(); + let runtime_stats = self.runtime_stats.clone(); + tokio::spawn(async move { + let resolved = resolve_bootstrap_nodes(&hosts, netmode).await; + let _ = priority_tx.try_send(PriorityEvent::BootstrapResolved(resolved)); + runtime_stats.set_crawl_priority_queue_depth( + priority_tx + .max_capacity() + .saturating_sub(priority_tx.capacity()), + ); + }); + } + + fn publish_snapshot(&self, now: Instant) { + let nodes = self + .responsive + .snapshot(self.config.routing_snapshot_size, now); + self.snapshot.store(Arc::new(RoutingSnapshot::from_nodes( + nodes, + self.config.routing_snapshot_size, + ))); + } + + fn metadata_pressure(&self) -> f64 { + if self.max_metadata_queue_size == 0 { + 1.0 + } else { + (self.metadata_queue_len.load(Ordering::Relaxed) as f64 + / self.max_metadata_queue_size as f64) + .min(1.0) + } + } + + fn sync_node_count(&self) { + self.node_count.store(self.pool.len(), Ordering::Relaxed); + self.runtime_stats.set_node_pool_size(self.pool.len()); + } + + fn flush_metrics(&mut self, now: Instant) { + #[cfg(feature = "metrics")] + { + let metadata_pressure = self.metadata_pressure(); + gauge!("dht_node_pool_size").set(self.pool.len() as f64); + gauge!("dht_node_pool_oldest_age_seconds").set(self.pool.oldest_age(now).as_secs_f64()); + gauge!("dht_find_node_in_flight").set(self.pending.len() as f64); + gauge!("dht_metadata_queue_pressure_ratio").set(metadata_pressure); + gauge!("dht_find_node_effective_rate_per_second") + .set(self.config.rate_for_metadata_pressure(metadata_pressure) as f64); + counter!("dht_node_pool_admissions_total").increment(self.metrics.admitted); + counter!("dht_node_pool_replacements_total").increment(self.metrics.replaced); + counter!("dht_node_pool_dropped_total", "reason" => "duplicate") + .increment(self.metrics.duplicate); + counter!("dht_node_pool_dropped_total", "reason" => "rate_limit") + .increment(self.metrics.admission_limited); + counter!("dht_node_pool_dropped_total", "reason" => "invalid") + .increment(self.metrics.invalid); + counter!("dht_crawl_queries_sent_total", "kind" => "new") + .increment(self.metrics.queries_new); + counter!("dht_crawl_queries_sent_total", "kind" => "revisit") + .increment(self.metrics.queries_revisit); + counter!("dht_crawl_queries_sent_total", "kind" => "bootstrap") + .increment(self.metrics.queries_bootstrap); + counter!("dht_find_node_responses_total").increment(self.metrics.responses); + counter!("dht_find_node_timeouts_total").increment(self.metrics.timeouts); + counter!("dht_find_node_send_failures_total").increment(self.metrics.send_failures); + counter!("dht_find_node_response_unmatched_total") + .increment(self.metrics.unmatched_responses); + } + #[cfg(not(feature = "metrics"))] + let _ = now; + self.metrics = ActorMetrics::default(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::types::CrawlOptions; + + fn node(id: u8, addr: &str) -> NodeTuple { + NodeTuple { + id: [id; 20], + addr: addr.parse().unwrap(), + } + } + + fn test_actor( + config: ResolvedCrawlConfig, + ) -> (CrawlActor, mpsc::Receiver, DhtRuntimeStats) { + let (priority_tx, priority_rx) = mpsc::channel(16); + let (_discovery_tx, discovery_rx) = mpsc::channel(16); + let (egress_tx, egress_rx) = mpsc::channel(config.max_in_flight); + let runtime_stats = DhtRuntimeStats::with_limits(DhtRuntimeLimits { + metadata_queue: 100_000, + node_pool: config.pool_capacity, + node_pool_low_watermark: config.low_watermark, + find_node_in_flight: config.max_in_flight, + initial_find_node_rate: config.max_find_node_rate_per_sec, + hash_ingress_queue: 0, + crawl_priority_queue: config.priority_event_channel_capacity, + crawl_discovery_queue: config.discovery_event_channel_capacity, + }); + let actor = CrawlActor::new(CrawlActorInit { + config, + netmode: NetMode::Ipv4Only, + local_id: [7; 20], + priority_rx, + discovery_rx, + priority_tx, + egress_v4: Some(egress_tx), + egress_v6: None, + snapshot: Arc::new(ArcSwap::from_pointee(RoutingSnapshot::default())), + node_count: Arc::new(AtomicUsize::new(0)), + 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) + } + + #[test] + 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"); + let eligible = node(2, "1.1.1.1:6881"); + + assert_eq!(actor.pool.admit(blocked, now), AdmissionOutcome::Admitted); + assert_eq!(actor.pool.admit(eligible, now), AdmissionOutcome::Admitted); + actor + .subnet_in_flight + .insert(SubnetKey::from_addr(&blocked.addr), max_per_subnet); + + assert!(actor.schedule_one(now)); + let request = egress_rx.try_recv().expect("eligible node was dispatched"); + + assert_eq!(request.node, eligible); + assert_eq!(actor.pool.front(), Some(blocked)); + assert_eq!(actor.pool.admit(blocked, now), AdmissionOutcome::Duplicate); + assert!(!actor.pool.contains_recent(&blocked.addr, now)); + assert!(actor.pool.contains_recent(&eligible.addr, now)); + 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, max_in_flight); + } + + #[test] + fn runtime_stats_count_dropped_crawl_events() { + let mut options = CrawlOptions::default(); + options.scheduler.priority_event_channel_capacity = 1; + options.scheduler.discovery_event_channel_capacity = 1; + let config = ResolvedCrawlConfig::from_options(&options); + let stats = DhtRuntimeStats::with_limits(DhtRuntimeLimits { + metadata_queue: 100, + node_pool: config.pool_capacity, + node_pool_low_watermark: config.low_watermark, + find_node_in_flight: config.max_in_flight, + initial_find_node_rate: config.max_find_node_rate_per_sec, + hash_ingress_queue: 0, + crawl_priority_queue: config.priority_event_channel_capacity, + crawl_discovery_queue: config.discovery_event_channel_capacity, + }); + 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")); + let response = || DhtResponse { + id: None, + nodes: None, + nodes6: None, + values: None, + samples: None, + num: None, + interval: None, + }; + engine.route_response("8.8.8.8:1".parse().unwrap(), [1; 8], response()); + engine.route_response("1.1.1.1:2".parse().unwrap(), [2; 8], response()); + + let snapshot = stats.snapshot(); + assert_eq!(snapshot.crawl_events_dropped_discovered, 1); + assert_eq!(snapshot.crawl_events_dropped_response, 1); + assert_eq!(snapshot.crawl_discovery_queue_depth, 1); + assert_eq!(snapshot.crawl_discovery_queue_capacity, 1); + assert_eq!(snapshot.crawl_priority_queue_depth, 1); + assert_eq!(snapshot.crawl_priority_queue_capacity, 1); + } +} diff --git a/dht-crawler/src/error.rs b/dht-crawler/src/error.rs new file mode 100644 index 0000000..26b5dde --- /dev/null +++ b/dht-crawler/src/error.rs @@ -0,0 +1,28 @@ +use thiserror::Error; + +#[derive(Error, Debug)] +/// Error returned during DHT initialization or execution. +pub enum DHTError { + /// Socket or other network I/O failed. + #[error("网络错误: {0}")] + Network(#[from] std::io::Error), + + /// A shared lock was poisoned. + #[error("锁中毒: {0}")] + LockPoisoned(String), + + /// Server initialization failed, for example because a socket could not bind. + #[error("初始化错误: {0}")] + Init(String), + + /// An internal invariant or worker operation failed. + #[error("内部错误: {0}")] + Internal(String), + + /// Another error represented by a human-readable message. + #[error("{0}")] + Other(String), +} + +/// Result type used by the crate's public APIs. +pub type Result = std::result::Result; diff --git a/dht-crawler/src/krpc.rs b/dht-crawler/src/krpc.rs new file mode 100644 index 0000000..805f7cb --- /dev/null +++ b/dht-crawler/src/krpc.rs @@ -0,0 +1,308 @@ +use crate::addr::{addr_allowed_by_netmode, is_valid_node_addr}; +use crate::node_id::TransactionId; +use crate::protocol::DhtResponse; +use crate::types::{NetMode, NodeTuple}; +use bytes::BytesMut; +#[cfg(feature = "metrics")] +use metrics::{counter, histogram}; +use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; +use std::sync::Arc; +use tokio::net::UdpSocket; + +pub(crate) fn for_each_response_node( + response: &DhtResponse, + netmode: NetMode, + mut visit: impl FnMut(NodeTuple), +) -> usize { + let mut count = 0; + if netmode != NetMode::Ipv6Only + && let Some(nodes) = response.nodes.as_deref() + && nodes.len() % 26 == 0 + { + for chunk in nodes.chunks_exact(26) { + let id: [u8; 20] = chunk[..20].try_into().expect("compact v4 id is 20 bytes"); + let ip = Ipv4Addr::new(chunk[20], chunk[21], chunk[22], chunk[23]); + let port = u16::from_be_bytes([chunk[24], chunk[25]]); + let addr = SocketAddr::new(IpAddr::V4(ip), port); + if is_valid_node_addr(&addr) { + visit(NodeTuple { id, addr }); + count += 1; + } + } + } + if netmode != NetMode::Ipv4Only + && let Some(nodes) = response.nodes6.as_deref() + && nodes.len() % 38 == 0 + { + for chunk in nodes.chunks_exact(38) { + let id: [u8; 20] = chunk[..20].try_into().expect("compact v6 id is 20 bytes"); + let ip_bytes: [u8; 16] = chunk[20..36] + .try_into() + .expect("compact v6 address is 16 bytes"); + let port = u16::from_be_bytes([chunk[36], chunk[37]]); + let addr = SocketAddr::new(IpAddr::V6(Ipv6Addr::from(ip_bytes)), port); + if is_valid_node_addr(&addr) { + visit(NodeTuple { id, addr }); + count += 1; + } + } + } + count +} + +pub(crate) fn for_each_response_peer( + response: &DhtResponse, + netmode: NetMode, + mut visit: impl FnMut(SocketAddr), +) -> usize { + let mut count = 0; + let Some(values) = response.values.as_ref() else { + return count; + }; + for value in values { + let bytes = value.as_ref(); + let addr = match bytes.len() { + 6 => { + let ip = Ipv4Addr::new(bytes[0], bytes[1], bytes[2], bytes[3]); + let port = u16::from_be_bytes([bytes[4], bytes[5]]); + SocketAddr::new(IpAddr::V4(ip), port) + } + 18 => { + let ip_bytes: [u8; 16] = bytes[..16] + .try_into() + .expect("compact IPv6 Peer address is 16 bytes"); + let port = u16::from_be_bytes([bytes[16], bytes[17]]); + SocketAddr::new(IpAddr::V6(Ipv6Addr::from(ip_bytes)), port) + } + _ => continue, + }; + if addr_allowed_by_netmode(&addr, netmode) && is_valid_node_addr(&addr) { + visit(addr); + count += 1; + } + } + count +} + +pub(crate) fn encode_find_node_query( + buffer: &mut BytesMut, + tid: &TransactionId, + target: &[u8; 20], + sender_id: &[u8; 20], +) { + buffer.clear(); + buffer.reserve(112); + buffer.extend_from_slice(b"d1:ad2:id20:"); + buffer.extend_from_slice(sender_id); + buffer.extend_from_slice(b"6:target20:"); + buffer.extend_from_slice(target); + buffer.extend_from_slice(b"e1:q9:find_node1:t8:"); + buffer.extend_from_slice(tid); + buffer.extend_from_slice(b"1:y1:qe"); +} + +pub(crate) fn encode_get_peers_query( + buffer: &mut BytesMut, + tid: &TransactionId, + info_hash: &[u8; 20], + sender_id: &[u8; 20], +) { + buffer.clear(); + buffer.reserve(111); + buffer.extend_from_slice(b"d1:ad2:id20:"); + buffer.extend_from_slice(sender_id); + buffer.extend_from_slice(b"9:info_hash20:"); + buffer.extend_from_slice(info_hash); + buffer.extend_from_slice(b"e1:q9:get_peers1:t8:"); + buffer.extend_from_slice(tid); + buffer.extend_from_slice(b"1:y1:qe"); +} + +pub(crate) fn encode_sample_infohashes_query( + buffer: &mut BytesMut, + tid: &TransactionId, + target: &[u8; 20], + sender_id: &[u8; 20], +) { + buffer.clear(); + buffer.reserve(128); + buffer.extend_from_slice(b"d1:ad2:id20:"); + buffer.extend_from_slice(sender_id); + buffer.extend_from_slice(b"6:target20:"); + buffer.extend_from_slice(target); + buffer.extend_from_slice(b"e1:q17:sample_infohashes1:t8:"); + buffer.extend_from_slice(tid); + buffer.extend_from_slice(b"1:y1:qe"); +} + +pub(crate) fn encode_response( + buffer: &mut BytesMut, + tid: &[u8], + node_id: &[u8; 20], + token: &[u8; 8], + nodes: &[NodeTuple], + ipv6: bool, +) { + buffer.clear(); + buffer.reserve(384); + buffer.extend_from_slice(b"d1:rd2:id20:"); + buffer.extend_from_slice(node_id); + + let compact_len = if ipv6 { + nodes.iter().filter(|node| node.addr.is_ipv6()).count() * 38 + } else { + nodes.iter().filter(|node| node.addr.is_ipv4()).count() * 26 + }; + if compact_len > 0 { + if ipv6 { + buffer.extend_from_slice(b"6:nodes6"); + } else { + buffer.extend_from_slice(b"5:nodes"); + } + push_usize(buffer, compact_len); + buffer.extend_from_slice(b":"); + for node in nodes { + match node.addr.ip() { + IpAddr::V4(ip) if !ipv6 => { + buffer.extend_from_slice(&node.id); + buffer.extend_from_slice(&ip.octets()); + buffer.extend_from_slice(&node.addr.port().to_be_bytes()); + } + IpAddr::V6(ip) if ipv6 => { + buffer.extend_from_slice(&node.id); + buffer.extend_from_slice(&ip.octets()); + buffer.extend_from_slice(&node.addr.port().to_be_bytes()); + } + _ => {} + } + } + } + + buffer.extend_from_slice(b"5:token8:"); + buffer.extend_from_slice(token); + buffer.extend_from_slice(b"e1:t"); + push_usize(buffer, tid.len()); + buffer.extend_from_slice(b":"); + buffer.extend_from_slice(tid); + buffer.extend_from_slice(b"1:y1:re"); +} + +fn push_usize(buffer: &mut BytesMut, mut value: usize) { + let mut digits = [0u8; 20]; + let mut cursor = digits.len(); + loop { + cursor -= 1; + digits[cursor] = b'0' + (value % 10) as u8; + value /= 10; + if value == 0 { + break; + } + } + buffer.extend_from_slice(&digits[cursor..]); +} + +pub(crate) async fn send_find_node_query( + addr: &SocketAddr, + tid: &TransactionId, + target: &[u8; 20], + sender_id: &[u8; 20], + socket: &Arc, + buffer: &mut BytesMut, +) -> bool { + encode_find_node_query(buffer, tid, target, sender_id); + match socket.send_to(buffer, addr).await { + Ok(len) => { + #[cfg(feature = "metrics")] + { + counter!("dht_udp_bytes_sent_total").increment(len as u64); + counter!("dht_udp_packets_sent_total", "type" => "query").increment(1); + histogram!("dht_udp_query_size_bytes").record(len as f64); + } + #[cfg(not(feature = "metrics"))] + let _ = len; + true + } + Err(_) => false, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::protocol::DhtMessage; + + #[test] + fn manual_find_node_encoding_round_trips() { + let mut buffer = BytesMut::new(); + let tid = [1; 8]; + let target = [2; 20]; + let sender = [3; 20]; + encode_find_node_query(&mut buffer, &tid, &target, &sender); + let message: DhtMessage = serde_bencode::from_bytes(&buffer).unwrap(); + assert_eq!(message.t.as_ref(), &tid); + assert_eq!(message.q.as_deref(), Some("find_node")); + assert_eq!(message.a.unwrap().target.unwrap().as_ref(), &target); + } + + #[test] + fn manual_get_peers_encoding_round_trips() { + let mut buffer = BytesMut::new(); + let tid = [1; 8]; + let info_hash = [2; 20]; + let sender = [3; 20]; + encode_get_peers_query(&mut buffer, &tid, &info_hash, &sender); + let message: DhtMessage = serde_bencode::from_bytes(&buffer).unwrap(); + assert_eq!(message.t.as_ref(), &tid); + assert_eq!(message.q.as_deref(), Some("get_peers")); + assert_eq!(message.a.unwrap().info_hash.unwrap().as_ref(), &info_hash); + } + + #[test] + fn manual_sample_infohashes_encoding_round_trips() { + let mut buffer = BytesMut::new(); + let tid = [1; 8]; + let target = [2; 20]; + let sender = [3; 20]; + encode_sample_infohashes_query(&mut buffer, &tid, &target, &sender); + let message: DhtMessage = serde_bencode::from_bytes(&buffer).unwrap(); + assert_eq!(message.t.as_ref(), &tid); + assert_eq!(message.q.as_deref(), Some("sample_infohashes")); + assert_eq!(message.a.unwrap().target.unwrap().as_ref(), &target); + } + + #[test] + fn manual_response_encoding_round_trips() { + let mut buffer = BytesMut::new(); + let nodes = [NodeTuple { + id: [4; 20], + addr: "8.8.8.8:6881".parse().unwrap(), + }]; + encode_response(&mut buffer, &[1, 2], &[2; 20], &[3; 8], &nodes, false); + let message: DhtMessage = serde_bencode::from_bytes(&buffer).unwrap(); + let response = message.r.unwrap(); + assert_eq!(response.nodes.unwrap().len(), 26); + } + + #[test] + fn compact_get_peers_values_are_validated() { + let response = DhtResponse { + id: None, + nodes: None, + nodes6: None, + values: Some(vec![ + serde_bytes::ByteBuf::from(vec![8, 8, 8, 8, 0x1a, 0xe1]), + serde_bytes::ByteBuf::from(vec![10, 0, 0, 1, 0x1a, 0xe1]), + serde_bytes::ByteBuf::from(vec![1, 2, 3]), + ]), + samples: None, + num: None, + interval: None, + }; + let mut peers = Vec::new(); + assert_eq!( + for_each_response_peer(&response, NetMode::Ipv4Only, |peer| peers.push(peer)), + 1 + ); + assert_eq!(peers[0], "8.8.8.8:6881".parse().unwrap()); + } +} diff --git a/dht-crawler/src/lib.rs b/dht-crawler/src/lib.rs new file mode 100644 index 0000000..b571365 --- /dev/null +++ b/dht-crawler/src/lib.rs @@ -0,0 +1,62 @@ +// 负责导出可复用 DHT 抓取 Metadata 和运行观测能力 + +//! High-throughput BitTorrent DHT crawler with bounded crawl and Metadata pipelines. +//! +//! [`DHTServer`] is the primary entry point. Configure it with [`DHTOptions`], register +//! callbacks, then await [`DHTServer::start`] until another task calls [`DHTServer::shutdown`]. +//! Runtime counters are available through [`DHTServer::runtime_stats`] without enabling any +//! exporter. See the repository README for scheduling, backpressure and migration details. + +mod addr; +mod bootstrap; +mod budget; +mod crawl_config; +mod crawl_engine; +mod error; +mod krpc; +/// BEP-9 Metadata fetch support. +pub mod metadata; +mod node_id; +mod node_pool; +mod peer_lookup; +/// Serializable BEP-5 KRPC wire types. +pub mod protocol; +mod routing_snapshot; +mod runtime_stats; +mod sample_infohashes; +/// Bounded, deduplicating Metadata scheduler. +pub mod scheduler; +mod server; +/// Public configuration, callback payload and network types. +pub mod types; +mod udp_buffer; +mod udp_ingress; + +pub use error::{DHTError, Result}; +pub use runtime_stats::{ + DhtObservabilitySnapshot, DhtRuntimeSnapshot, DhtRuntimeStats, FixedHistogramSnapshot, +}; +pub use scheduler::{MetadataScheduler, MetadataSchedulerCallbacks, MetadataSchedulerLimits}; +pub use server::{DHTServer, HashDiscovered}; +pub use types::{ + BootstrapOptions, CrawlOptions, DHTOptions, FileInfo, MetadataFetchCompletion, + MetadataFetchCompletionStatus, MetadataOptions, NetMode, NodeTuple, PeerLookupOptions, + PeerLookupResult, PoolOptions, RateLimitOptions, SampleInfohashesOptions, SchedulerOptions, + TargetOptions, TorrentInfo, +}; + +/// Common server, configuration and callback payload imports. +pub mod prelude { + pub use crate::error::{DHTError, Result}; + pub use crate::runtime_stats::{DhtRuntimeSnapshot, DhtRuntimeStats}; + pub use crate::scheduler::{ + MetadataScheduler, MetadataSchedulerCallbacks, MetadataSchedulerLimits, + }; + pub use crate::server::DHTServer; + pub use crate::types::{ + BootstrapOptions, CrawlOptions, DHTOptions, FileInfo, MetadataFetchCompletion, + MetadataFetchCompletionStatus, MetadataOptions, NetMode, NodeTuple, PeerLookupOptions, + PeerLookupResult, PoolOptions, RateLimitOptions, SampleInfohashesOptions, SchedulerOptions, + TargetOptions, TorrentInfo, + }; +} diff --git a/dht-crawler/src/metadata.rs b/dht-crawler/src/metadata.rs new file mode 100644 index 0000000..91c336c --- /dev/null +++ b/dht-crawler/src/metadata.rs @@ -0,0 +1,720 @@ +// 负责执行 BitTorrent 握手 Metadata 下载校验解析和 Peer 失败缓存 + +use crate::runtime_stats::DhtRuntimeStats; +use crate::types::FileInfo; +use ahash::AHashMap; +use bytes::Bytes; +#[cfg(feature = "metrics")] +use metrics::{counter, gauge, histogram}; +use rbit::peer::ExtensionMessage; +use rbit::{ + ExtensionHandshake, Message, MetadataMessage, MetadataMessageType, PeerConnection, PeerId, + metadata_piece_count, +}; +use sha1::{Digest, Sha1}; +use std::collections::{BTreeMap, VecDeque}; +use std::net::SocketAddr; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; +use tokio::time::timeout; + +pub(crate) type FetchedMetadata = (String, u64, Vec, u64); + +pub(crate) enum MetadataFetchOutcome { + Fetched(FetchedMetadata), + Failed, + SkippedCached, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum MetadataFetchFailure { + Connect, + NoExtension, + Send, + SizeLimit, + Sha1, + Parse, + Other, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum PeerFailureReason { + Timeout, + 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 { + match self { + Self::Timeout => "timeout", + Self::ConnectFailed => "connect_failed", + } + } +} + +#[derive(Debug, Clone, Copy)] +struct PeerFailureEntry { + expires_at: Instant, + reason: PeerFailureReason, +} + +#[derive(Default)] +struct PeerFailureCacheInner { + entries: AHashMap, + expiry: VecDeque<(Instant, SocketAddr)>, +} + +struct PeerFailureCache { + inner: Mutex, + capacity: usize, + ttl: Duration, +} + +impl PeerFailureCache { + fn new(capacity: usize, ttl: Duration) -> Self { + Self { + inner: Mutex::new(PeerFailureCacheInner { + entries: AHashMap::with_capacity(capacity.min(16_384)), + expiry: VecDeque::with_capacity(capacity.min(16_384)), + }), + capacity, + ttl, + } + } + + fn get(&self, addr: SocketAddr, now: Instant) -> (Option, usize) { + if self.capacity == 0 || self.ttl.is_zero() { + return (None, 0); + } + let mut inner = self + .inner + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + Self::expire(&mut inner, now); + ( + inner.entries.get(&addr).map(|entry| entry.reason), + inner.entries.len(), + ) + } + + fn insert(&self, addr: SocketAddr, reason: PeerFailureReason, now: Instant) -> usize { + if self.capacity == 0 || self.ttl.is_zero() { + return 0; + } + let mut inner = self + .inner + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + Self::expire(&mut inner, now); + while inner.entries.len() >= self.capacity && !inner.entries.contains_key(&addr) { + let Some((expires_at, oldest_addr)) = inner.expiry.pop_front() else { + break; + }; + if inner + .entries + .get(&oldest_addr) + .is_some_and(|entry| entry.expires_at == expires_at) + { + inner.entries.remove(&oldest_addr); + } + } + + let expires_at = now + self.ttl; + inner + .entries + .insert(addr, PeerFailureEntry { expires_at, reason }); + inner.expiry.push_back((expires_at, addr)); + inner.entries.len() + } + + fn remove(&self, addr: &SocketAddr, now: Instant) -> usize { + if self.capacity == 0 || self.ttl.is_zero() { + return 0; + } + let mut inner = self + .inner + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + Self::expire(&mut inner, now); + inner.entries.remove(addr); + inner.entries.len() + } + + fn expire(inner: &mut PeerFailureCacheInner, now: Instant) { + while let Some((expires_at, addr)) = inner.expiry.front().copied() { + if expires_at > now { + break; + } + inner.expiry.pop_front(); + if inner + .entries + .get(&addr) + .is_some_and(|entry| entry.expires_at == expires_at) + { + inner.entries.remove(&addr); + } + } + } +} + +#[derive(Clone)] +/// BEP-9 Metadata fetcher with an end-to-end timeout and shared Peer failure cache. +pub struct RbitFetcher { + total_timeout: Duration, + runtime_stats: DhtRuntimeStats, + peer_failure_cache: Arc, + connect_rate_limiter: Arc, +} + +impl RbitFetcher { + /// Creates a standalone fetcher with the default failure-cache capacity and TTL. + /// + /// [`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, 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, + ) -> Self { + Self { + total_timeout: Duration::from_secs(if timeout_secs == 0 { 15 } else { timeout_secs }), + runtime_stats, + peer_failure_cache: Arc::new(PeerFailureCache::new( + peer_failure_cache_capacity, + Duration::from_secs(peer_failure_ttl_secs), + )), + connect_rate_limiter: Arc::new(ConnectRateLimiter::per_second(max_connects_per_second)), + } + } + + /// Fetch metadata from one peer under a single end-to-end deadline. + /// + /// The deadline covers TCP connect, both BitTorrent handshakes, all metadata + /// piece I/O, hash validation and bencode parsing. Inner library timeouts can + /// therefore never stack on top of the configured metadata timeout. + #[cfg(test)] + pub(crate) async fn fetch( + &self, + info_hash: &[u8; 20], + peer_addr: SocketAddr, + ) -> MetadataFetchOutcome { + self.fetch_with_attempt_observer(info_hash, peer_addr, || {}) + .await + } + + pub(crate) async fn fetch_with_attempt_observer( + &self, + info_hash: &[u8; 20], + peer_addr: SocketAddr, + on_attempt: F, + ) -> MetadataFetchOutcome + where + F: FnOnce() + Send, + { + let (cached_reason, cache_entries) = self.peer_failure_cache.get(peer_addr, Instant::now()); + self.set_peer_failure_cache_entries(cache_entries); + if let Some(reason) = cached_reason { + self.runtime_stats.metadata_peer_failure_cache_hit(); + match reason { + PeerFailureReason::Timeout => self.runtime_stats.peer_cache_hit_timeout(), + PeerFailureReason::ConnectFailed => self.runtime_stats.peer_cache_hit_connect(), + } + #[cfg(feature = "metrics")] + counter!("dht_metadata_peer_failure_cache_hits_total", "reason" => reason.as_str()) + .increment(1); + #[cfg(not(feature = "metrics"))] + let _ = reason; + return MetadataFetchOutcome::SkippedCached; + } + + self.connect_rate_limiter.acquire().await; + on_attempt(); + self.runtime_stats.metadata_peer_attempt(); + #[cfg(feature = "metrics")] + { + counter!("dht_metadata_fetch_attempts_total").increment(1); + counter!("dht_metadata_peer_attempts_total").increment(1); + } + + let started = Instant::now(); + let result = timeout( + self.total_timeout, + self.fetch_with_peer(info_hash, peer_addr), + ) + .await; + + #[cfg(feature = "metrics")] + histogram!("dht_metadata_fetch_duration_seconds").record(started.elapsed().as_secs_f64()); + self.runtime_stats.observe_metadata_fetch_duration( + started.elapsed().as_millis().min(u128::from(u64::MAX)) as u64, + ); + + match result { + Ok(Ok(metadata)) => { + let cache_entries = self.peer_failure_cache.remove(&peer_addr, Instant::now()); + self.set_peer_failure_cache_entries(cache_entries); + self.runtime_stats.metadata_peer_succeeded(); + #[cfg(feature = "metrics")] + { + counter!("dht_metadata_fetch_success_total").increment(1); + counter!("dht_metadata_fetch_result_total", "result" => "success").increment(1); + } + MetadataFetchOutcome::Fetched(metadata) + } + Ok(Err(reason)) => { + self.runtime_stats.metadata_peer_failed(); + match reason { + MetadataFetchFailure::Connect => self.runtime_stats.metadata_failure_connect(), + MetadataFetchFailure::NoExtension => { + self.runtime_stats.metadata_failure_no_extension() + } + MetadataFetchFailure::Send => self.runtime_stats.metadata_failure_send(), + MetadataFetchFailure::SizeLimit => { + self.runtime_stats.metadata_failure_size_limit() + } + MetadataFetchFailure::Sha1 => self.runtime_stats.metadata_failure_sha1(), + MetadataFetchFailure::Parse => self.runtime_stats.metadata_failure_parse(), + MetadataFetchFailure::Other => self.runtime_stats.metadata_failure_other(), + } + #[cfg(feature = "metrics")] + counter!("dht_metadata_fetch_result_total", "result" => "failed").increment(1); + MetadataFetchOutcome::Failed + } + Err(_) => { + self.record_peer_failure(peer_addr, PeerFailureReason::Timeout); + self.runtime_stats.metadata_peer_failed(); + self.runtime_stats.metadata_peer_timeout(); + self.runtime_stats.metadata_failure_timeout(); + #[cfg(feature = "metrics")] + { + counter!("dht_metadata_fetch_fail_total", "reason" => "timeout").increment(1); + counter!("dht_metadata_fetch_result_total", "result" => "timeout").increment(1); + } + MetadataFetchOutcome::Failed + } + } + } + + pub(crate) async fn verify_handshake( + &self, + info_hash: &[u8; 20], + peer_addr: SocketAddr, + ) -> bool { + let (cached_reason, _) = self.peer_failure_cache.get(peer_addr, Instant::now()); + if cached_reason.is_some() { + return false; + } + self.connect_rate_limiter.acquire().await; + let peer_id = PeerId::generate(); + match timeout( + self.total_timeout, + PeerConnection::connect(peer_addr, *info_hash, *peer_id.as_bytes()), + ) + .await + { + Ok(Ok(_)) => { + self.peer_failure_cache.remove(&peer_addr, Instant::now()); + true + } + Ok(Err(_)) => { + self.record_peer_failure(peer_addr, PeerFailureReason::ConnectFailed); + false + } + Err(_) => { + self.record_peer_failure(peer_addr, PeerFailureReason::Timeout); + false + } + } + } + + fn record_peer_failure(&self, peer_addr: SocketAddr, reason: PeerFailureReason) { + let cache_entries = self + .peer_failure_cache + .insert(peer_addr, reason, Instant::now()); + self.set_peer_failure_cache_entries(cache_entries); + #[cfg(feature = "metrics")] + counter!("dht_metadata_peer_failure_cache_inserts_total", "reason" => reason.as_str()) + .increment(1); + } + + fn set_peer_failure_cache_entries(&self, count: usize) { + self.runtime_stats + .set_metadata_peer_failure_cache_entries(count); + #[cfg(feature = "metrics")] + gauge!("dht_metadata_peer_failure_cache_entries").set(count as f64); + } + + async fn fetch_with_peer( + &self, + info_hash: &[u8; 20], + peer_addr: SocketAddr, + ) -> Result { + let peer_id = PeerId::generate(); + let mut conn = match PeerConnection::connect(peer_addr, *info_hash, *peer_id.as_bytes()) + .await + { + Ok(conn) => { + #[cfg(feature = "metrics")] + counter!("dht_metadata_connection_result_total", "result" => "success") + .increment(1); + conn + } + Err(_) => { + self.record_peer_failure(peer_addr, PeerFailureReason::ConnectFailed); + self.runtime_stats.metadata_connect_failed(); + #[cfg(feature = "metrics")] + counter!("dht_metadata_connection_result_total", "result" => "failed").increment(1); + return Err(MetadataFetchFailure::Connect); + } + }; + + if !conn.supports_extension { + self.runtime_stats.metadata_no_extension(); + #[cfg(feature = "metrics")] + counter!("dht_metadata_handshake_result_total", "result" => "no_extension_support") + .increment(1); + return Err(MetadataFetchFailure::NoExtension); + } + + let my_ut_metadata_id = 1; + let handshake = ExtensionHandshake::with_extensions(&[("ut_metadata", my_ut_metadata_id)]); + let handshake_bytes = handshake.encode().map_err(|_| MetadataFetchFailure::Send)?; + if conn + .send(Message::Extended { + id: 0, + payload: handshake_bytes, + }) + .await + .is_err() + { + #[cfg(feature = "metrics")] + counter!("dht_metadata_fetch_fail_total", "reason" => "send_error").increment(1); + return Err(MetadataFetchFailure::Send); + } + + let mut metadata_size = 0; + let mut remote_ut_metadata_id = 0; + let mut pieces: BTreeMap = BTreeMap::new(); + let mut total_received = 0usize; + let mut request_sent = false; + + let info_bytes = loop { + let msg = conn + .receive() + .await + .map_err(|_| MetadataFetchFailure::Other)?; + let Message::Extended { id, payload } = msg else { + continue; + }; + + if id == 0 { + if let Ok(ExtensionMessage::Handshake(remote_hs)) = + ExtensionMessage::decode(id, &payload) + { + if let Some(size) = remote_hs.metadata_size { + metadata_size = size as u32; + } + if let Some(ext_id) = remote_hs.get_extension_id("ut_metadata") { + remote_ut_metadata_id = ext_id; + } + } + + if metadata_size > 0 && remote_ut_metadata_id > 0 && !request_sent { + if metadata_size > 10 * 1024 * 1024 { + #[cfg(feature = "metrics")] + counter!("dht_metadata_fetch_fail_total", "reason" => "size_limit") + .increment(1); + return Err(MetadataFetchFailure::SizeLimit); + } + + let count = metadata_piece_count(metadata_size as usize); + for piece in 0..count { + let encoded = MetadataMessage::request(piece as u32) + .encode() + .map_err(|_| MetadataFetchFailure::Send)?; + if conn + .send(Message::Extended { + id: remote_ut_metadata_id, + payload: encoded, + }) + .await + .is_err() + { + #[cfg(feature = "metrics")] + counter!("dht_metadata_fetch_fail_total", "reason" => "send_error") + .increment(1); + return Err(MetadataFetchFailure::Send); + } + } + request_sent = true; + } + continue; + } + + if id != my_ut_metadata_id { + continue; + } + let Ok(meta_msg) = MetadataMessage::decode(&payload) else { + continue; + }; + if meta_msg.msg_type != MetadataMessageType::Data { + continue; + } + let Some(data) = meta_msg.data else { + continue; + }; + + #[cfg(feature = "metrics")] + counter!("dht_metadata_bytes_downloaded_total").increment(data.len() as u64); + self.runtime_stats.metadata_bytes_downloaded(data.len()); + + let data_len = data.len(); + if let Some(previous) = pieces.insert(meta_msg.piece, data) { + total_received = total_received.saturating_sub(previous.len()); + } + total_received = total_received.saturating_add(data_len); + + if metadata_size == 0 || total_received < metadata_size as usize { + continue; + } + + let count = metadata_piece_count(metadata_size as usize); + let mut full_data = Vec::with_capacity(metadata_size as usize); + for piece in 0..count { + let data = pieces + .get(&(piece as u32)) + .ok_or(MetadataFetchFailure::Other)?; + full_data.extend_from_slice(data); + } + + let info_hash_copy = *info_hash; + let validated = tokio::task::spawn_blocking(move || { + let mut hasher = Sha1::new(); + hasher.update(&full_data); + let digest: [u8; 20] = hasher.finalize().into(); + (digest == info_hash_copy).then_some(full_data) + }) + .await + .ok() + .flatten(); + + match validated { + Some(data) => { + #[cfg(feature = "metrics")] + counter!("dht_metadata_handshake_result_total", "result" => "success") + .increment(1); + break data; + } + None => { + #[cfg(feature = "metrics")] + counter!("dht_metadata_fetch_fail_total", "reason" => "sha1_mismatch") + .increment(1); + return Err(MetadataFetchFailure::Sha1); + } + } + }; + + self.runtime_stats.observe_metadata_size(info_bytes.len()); + match parse_metadata(&info_bytes) { + Some(metadata) => { + #[cfg(feature = "metrics")] + histogram!("dht_metadata_size_bytes").record(info_bytes.len() as f64); + Ok(metadata) + } + None => { + #[cfg(feature = "metrics")] + counter!("dht_metadata_fetch_fail_total", "reason" => "parse_error").increment(1); + Err(MetadataFetchFailure::Parse) + } + } + } +} + +fn parse_metadata(info_bytes: &[u8]) -> Option { + let value = rbit::decode(info_bytes).ok()?; + let dict = value.as_dict()?; + let name = dict + .get(&b"name"[..]) + .and_then(|value| value.as_str()) + .unwrap_or("Unknown") + .to_string(); + let piece_length = dict + .get(&b"piece length"[..]) + .and_then(|value| value.as_integer()) + .unwrap_or(0) as u64; + + let mut total_size = 0; + let mut file_list = Vec::new(); + if let Some(files) = dict.get(&b"files"[..]).and_then(|value| value.as_list()) { + for file in files { + let Some(file_dict) = file.as_dict() else { + continue; + }; + let Some(length) = file_dict + .get(&b"length"[..]) + .and_then(|value| value.as_integer()) + else { + continue; + }; + let length = length as u64; + total_size += length; + let path = file_dict + .get(&b"path"[..]) + .and_then(|value| value.as_list()) + .map(|parts| { + parts + .iter() + .filter_map(|part| part.as_str()) + .collect::>() + .join("/") + }) + .unwrap_or_default(); + file_list.push(FileInfo { path, size: length }); + } + } else if let Some(length) = dict + .get(&b"length"[..]) + .and_then(|value| value.as_integer()) + { + total_size = length as u64; + file_list.push(FileInfo { + path: name.clone(), + size: total_size, + }); + } + + (total_size > 0).then_some((name, total_size, file_list, piece_length)) +} + +#[cfg(test)] +mod tests { + use super::*; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + 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(); + let cache = PeerFailureCache::new(10, Duration::from_secs(60)); + let first: SocketAddr = "127.0.0.1:1000".parse().unwrap(); + let same_ip_other_port: SocketAddr = "127.0.0.1:1001".parse().unwrap(); + + assert_eq!(cache.insert(first, PeerFailureReason::Timeout, start), 1); + assert_eq!(cache.get(first, start).0, Some(PeerFailureReason::Timeout)); + assert_eq!(cache.get(same_ip_other_port, start).0, None); + assert_eq!(cache.get(first, start + Duration::from_secs(61)), (None, 0)); + } + + #[test] + fn peer_failure_cache_evicts_oldest_at_capacity() { + let start = Instant::now(); + let cache = PeerFailureCache::new(1, Duration::from_secs(60)); + let first: SocketAddr = "127.0.0.1:1000".parse().unwrap(); + let second: SocketAddr = "127.0.0.1:1001".parse().unwrap(); + + cache.insert(first, PeerFailureReason::Timeout, start); + cache.insert(second, PeerFailureReason::ConnectFailed, start); + + assert_eq!(cache.get(first, start).0, None); + assert_eq!( + cache.get(second, start).0, + Some(PeerFailureReason::ConnectFailed) + ); + } + + #[tokio::test] + async fn total_timeout_covers_peer_handshake() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let accept_task = tokio::spawn(async move { + let (_stream, _) = listener.accept().await.unwrap(); + std::future::pending::<()>().await; + }); + + let stats = DhtRuntimeStats::default(); + 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, + MetadataFetchOutcome::Failed + )); + assert!(started.elapsed() < Duration::from_secs(2)); + + let cached_started = Instant::now(); + assert!(matches!( + fetcher.fetch(&[8; 20], addr).await, + MetadataFetchOutcome::SkippedCached + )); + assert!(cached_started.elapsed() < Duration::from_millis(100)); + + let snapshot = stats.snapshot(); + assert_eq!(snapshot.metadata_peer_attempts, 1); + assert_eq!(snapshot.metadata_peer_failed, 1); + assert_eq!(snapshot.metadata_peer_timeouts, 1); + assert_eq!(snapshot.metadata_peer_failure_cache_hits, 1); + assert_eq!(snapshot.metadata_peer_failure_cache_entries, 1); + + accept_task.abort(); + } + + #[tokio::test] + async fn handshake_verifier_accepts_matching_bittorrent_peer() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let peer = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut handshake = [0_u8; 68]; + stream.read_exact(&mut handshake).await.unwrap(); + stream.write_all(&handshake).await.unwrap(); + }); + let fetcher = + RbitFetcher::new_with_runtime_stats(1, 10, 10, 60, DhtRuntimeStats::default()); + assert!(fetcher.verify_handshake(&[7; 20], addr).await); + peer.await.unwrap(); + } +} diff --git a/dht-crawler/src/node_id.rs b/dht-crawler/src/node_id.rs new file mode 100644 index 0000000..18c546a --- /dev/null +++ b/dht-crawler/src/node_id.rs @@ -0,0 +1,82 @@ +use rand::RngExt; + +pub(crate) type TransactionId = [u8; 8]; + +pub(crate) fn transaction_id_from_bytes(bytes: &[u8]) -> Option { + if bytes.len() != 8 { + return None; + } + + let mut tid = [0u8; 8]; + tid.copy_from_slice(bytes); + Some(tid) +} + +pub(crate) fn random_node_id() -> [u8; 20] { + let mut id = [0u8; 20]; + rand::rng().fill(&mut id); + id +} + +pub(crate) fn neighbor_node_id(remote_id: &[u8], local_id: &[u8]) -> Vec { + let mut id = Vec::with_capacity(20); + let prefix_len = remote_id.len().min(6); + id.extend_from_slice(&remote_id[..prefix_len]); + + if local_id.len() > prefix_len { + id.extend_from_slice(&local_id[prefix_len..]); + } + while id.len() < 20 { + id.push(rand::random()); + } + id.truncate(20); + id +} + +pub(crate) fn bucket_index(id: &[u8], local_id: &[u8; 20]) -> usize { + for bit in 0..160 { + let byte = bit / 8; + if byte >= id.len() { + break; + } + let mask = 1 << (7 - (bit % 8)); + if (id[byte] ^ local_id[byte]) & mask != 0 { + return bit; + } + } + 159 +} + +pub(crate) fn target_for_bucket(local_id: &[u8; 20], bucket: usize) -> [u8; 20] { + let mut id = *local_id; + let bucket = bucket.min(159); + let byte = bucket / 8; + let bit = 7 - (bucket % 8); + id[byte] ^= 1 << bit; + for item in id.iter_mut().skip(byte + 1) { + *item = rand::random(); + } + id +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn transaction_ids_are_eight_bytes() { + let first = 1u64.to_be_bytes(); + assert_eq!(transaction_id_from_bytes(&first), Some(first)); + assert!(transaction_id_from_bytes(&[1, 2]).is_none()); + } + + #[test] + fn neighbor_id_keeps_remote_prefix_and_local_suffix() { + let remote = [1u8; 20]; + let local = [2u8; 20]; + let id = neighbor_node_id(&remote, &local); + + assert_eq!(&id[..6], &[1u8; 6]); + assert_eq!(&id[6..], &[2u8; 14]); + } +} diff --git a/dht-crawler/src/node_pool.rs b/dht-crawler/src/node_pool.rs new file mode 100644 index 0000000..3a57abc --- /dev/null +++ b/dht-crawler/src/node_pool.rs @@ -0,0 +1,390 @@ +use crate::addr::is_valid_node_addr; +use crate::budget::RateBucket; +use crate::types::NodeTuple; +use ahash::{AHashMap, AHashSet}; +use std::collections::VecDeque; +use std::net::{IpAddr, SocketAddr}; +use std::time::{Duration, Instant}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum AdmissionOutcome { + Admitted, + Replaced, + Duplicate, + RateLimited, + Invalid, +} + +#[derive(Debug, Clone, Copy)] +struct QueuedNode { + node: NodeTuple, + #[cfg_attr(not(feature = "metrics"), allow(dead_code))] + queued_at: Instant, +} + +pub(crate) struct NodePool { + queue: VecDeque, + queued: AHashSet, + recent: AHashMap, + recent_expiry: VecDeque<(Instant, SocketAddr)>, + replacement_budget: RateBucket, + recent_ttl: Duration, + capacity: usize, + warmed: bool, +} + +impl NodePool { + pub(crate) fn new( + capacity: usize, + replacements_per_minute: u32, + recent_ttl: Duration, + now: Instant, + ) -> Self { + let capacity = capacity.max(1); + let replacement_burst = replacements_per_minute.div_ceil(60).max(1); + Self { + queue: VecDeque::with_capacity(capacity), + queued: AHashSet::with_capacity(capacity), + recent: AHashMap::with_capacity(capacity), + recent_expiry: VecDeque::with_capacity(capacity), + replacement_budget: RateBucket::per_minute( + replacements_per_minute, + replacement_burst, + true, + now, + ), + recent_ttl, + capacity, + warmed: false, + } + } + + pub(crate) fn admit(&mut self, node: NodeTuple, now: Instant) -> AdmissionOutcome { + if !is_valid_node_addr(&node.addr) { + return AdmissionOutcome::Invalid; + } + self.expire_recent(now); + if self.queued.contains(&node.addr) || self.recent.contains_key(&node.addr) { + return AdmissionOutcome::Duplicate; + } + if self.warmed && !self.replacement_budget.try_take_one(now) { + return AdmissionOutcome::RateLimited; + } + + let replaced = if self.queue.len() >= self.capacity { + self.pop_front_internal().is_some() + } else { + false + }; + self.queued.insert(node.addr); + self.queue.push_back(QueuedNode { + node, + queued_at: now, + }); + if self.queue.len() >= self.capacity { + self.warmed = true; + } + if replaced { + AdmissionOutcome::Replaced + } else { + AdmissionOutcome::Admitted + } + } + + pub(crate) fn front(&self) -> Option { + self.queue.front().map(|entry| entry.node) + } + + /// Move the FIFO head behind the remaining queued nodes without marking + /// it as probed. The address stays in `queued` and is not added to `recent`. + pub(crate) fn rotate_front_to_back(&mut self) -> bool { + if self.queue.len() <= 1 { + return false; + } + self.queue.rotate_left(1); + true + } + + pub(crate) fn take_front_for_probe(&mut self, now: Instant) -> Option { + let entry = self.pop_front_internal()?; + let expires_at = now + self.recent_ttl; + self.recent.insert(entry.node.addr, expires_at); + self.recent_expiry.push_back((expires_at, entry.node.addr)); + Some(entry.node) + } + + pub(crate) fn restore_front(&mut self, node: NodeTuple, queued_at: Instant) { + self.recent.remove(&node.addr); + self.queued.insert(node.addr); + self.queue.push_front(QueuedNode { node, queued_at }); + } + + pub(crate) fn contains_recent(&mut self, addr: &SocketAddr, now: Instant) -> bool { + self.expire_recent(now); + self.recent.contains_key(addr) + } + + pub(crate) fn record_probe(&mut self, addr: SocketAddr, now: Instant) { + let expires_at = now + self.recent_ttl; + self.recent.insert(addr, expires_at); + self.recent_expiry.push_back((expires_at, addr)); + } + + pub(crate) fn len(&self) -> usize { + self.queue.len() + } + + #[cfg_attr(not(feature = "metrics"), allow(dead_code))] + pub(crate) fn oldest_age(&self, now: Instant) -> Duration { + self.queue + .front() + .and_then(|entry| now.checked_duration_since(entry.queued_at)) + .unwrap_or_default() + } + + #[cfg(test)] + fn is_warmed(&self) -> bool { + self.warmed + } + + fn pop_front_internal(&mut self) -> Option { + let entry = self.queue.pop_front()?; + self.queued.remove(&entry.node.addr); + Some(entry) + } + + fn expire_recent(&mut self, now: Instant) { + while let Some((expires_at, addr)) = self.recent_expiry.front().copied() { + if expires_at > now { + break; + } + self.recent_expiry.pop_front(); + if self.recent.get(&addr).copied() == Some(expires_at) { + self.recent.remove(&addr); + } + } + } +} + +#[derive(Debug, Clone, Copy)] +struct ResponsiveEntry { + node: NodeTuple, + expires_at: Instant, +} + +/// Fixed-size responsive-node ring. The crawl actor is the only writer. +pub(crate) struct ResponsiveReservoir { + slots: Vec>, + index: AHashMap, + write_cursor: usize, + revisit_cursor: usize, + ttl: Duration, +} + +impl ResponsiveReservoir { + pub(crate) fn new(capacity: usize, ttl: Duration) -> Self { + let capacity = capacity.max(1); + Self { + slots: vec![None; capacity], + index: AHashMap::with_capacity(capacity), + write_cursor: 0, + revisit_cursor: 0, + ttl, + } + } + + pub(crate) fn record(&mut self, node: NodeTuple, now: Instant) { + let entry = ResponsiveEntry { + node, + expires_at: now + self.ttl, + }; + if let Some(slot) = self.index.get(&node.addr).copied() { + self.slots[slot] = Some(entry); + return; + } + + let slot = self.write_cursor; + if let Some(old) = self.slots[slot] + && self.index.get(&old.node.addr).copied() == Some(slot) + { + self.index.remove(&old.node.addr); + } + self.slots[slot] = Some(entry); + self.index.insert(node.addr, slot); + self.write_cursor = (self.write_cursor + 1) % self.slots.len(); + } + + pub(crate) fn next_revisit(&mut self, now: Instant) -> Option { + for _ in 0..self.slots.len() { + let slot = self.revisit_cursor; + self.revisit_cursor = (self.revisit_cursor + 1) % self.slots.len(); + let Some(entry) = self.slots[slot] else { + continue; + }; + if entry.expires_at <= now { + if self.index.get(&entry.node.addr).copied() == Some(slot) { + self.index.remove(&entry.node.addr); + } + self.slots[slot] = None; + continue; + } + return Some(entry.node); + } + None + } + + pub(crate) fn snapshot(&self, limit: usize, now: Instant) -> Vec { + self.slots + .iter() + .filter_map(|entry| { + entry + .filter(|entry| entry.expires_at > now) + .map(|entry| entry.node) + }) + .take(limit) + .collect() + } +} + +#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)] +pub(crate) enum SubnetKey { + V4([u8; 3]), + V6([u8; 8]), +} + +impl SubnetKey { + pub(crate) fn from_addr(addr: &SocketAddr) -> Self { + match addr.ip() { + IpAddr::V4(ip) => { + let octets = ip.octets(); + Self::V4([octets[0], octets[1], octets[2]]) + } + IpAddr::V6(ip) => { + let octets = ip.octets(); + Self::V6(octets[..8].try_into().expect("IPv6 prefix has eight bytes")) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::net::{IpAddr, Ipv4Addr}; + + fn node(id: u8, addr: &str) -> NodeTuple { + NodeTuple { + id: [id; 20], + addr: addr.parse().unwrap(), + } + } + + #[test] + fn strict_fifo_replaces_oldest_after_warmup() { + let start = Instant::now(); + let mut pool = NodePool::new(2, 600, Duration::from_secs(60), start); + let first = node(1, "8.8.8.8:1"); + let second = node(2, "1.1.1.1:2"); + let third = node(3, "9.9.9.9:3"); + assert_eq!(pool.admit(first, start), AdmissionOutcome::Admitted); + assert_eq!(pool.admit(second, start), AdmissionOutcome::Admitted); + assert!(pool.is_warmed()); + assert_eq!(pool.admit(third, start), AdmissionOutcome::Replaced); + assert_eq!(pool.front(), Some(second)); + } + + #[test] + fn duplicate_does_not_reorder_fifo() { + let start = Instant::now(); + let mut pool = NodePool::new(3, 600, Duration::from_secs(60), start); + let first = node(1, "8.8.8.8:1"); + let second = node(2, "1.1.1.1:2"); + pool.admit(first, start); + pool.admit(second, start); + assert_eq!(pool.admit(first, start), AdmissionOutcome::Duplicate); + assert_eq!(pool.front(), Some(first)); + } + + #[test] + fn rotating_front_preserves_queued_dedup_and_recent_state() { + let start = Instant::now(); + let mut pool = NodePool::new(3, 600, Duration::from_secs(60), start); + let first = node(1, "8.8.8.8:1"); + let second = node(2, "1.1.1.1:2"); + + assert_eq!(pool.admit(first, start), AdmissionOutcome::Admitted); + assert_eq!(pool.admit(second, start), AdmissionOutcome::Admitted); + assert!(pool.rotate_front_to_back()); + + assert_eq!(pool.front(), Some(second)); + assert_eq!(pool.admit(first, start), AdmissionOutcome::Duplicate); + assert!(!pool.contains_recent(&first.addr, start)); + } + + #[test] + fn warmed_pool_enforces_replacement_rate() { + let start = Instant::now(); + let mut pool = NodePool::new(2, 60, Duration::from_secs(60), start); + pool.admit(node(1, "8.8.8.8:1"), start); + pool.admit(node(2, "1.1.1.1:2"), start); + assert_eq!( + pool.admit(node(3, "9.9.9.9:3"), start), + AdmissionOutcome::Replaced + ); + assert_eq!( + pool.admit(node(4, "208.67.222.222:4"), start), + AdmissionOutcome::RateLimited + ); + assert_eq!( + pool.admit(node(4, "208.67.222.222:4"), start + Duration::from_secs(1)), + AdmissionOutcome::Replaced + ); + } + + #[test] + fn recent_probe_blocks_readmission_until_expiry() { + let start = Instant::now(); + let mut pool = NodePool::new(3, 600, Duration::from_secs(10), start); + let first = node(1, "8.8.8.8:1"); + pool.admit(first, start); + assert_eq!(pool.take_front_for_probe(start), Some(first)); + assert_eq!(pool.admit(first, start), AdmissionOutcome::Duplicate); + assert_eq!( + pool.admit(first, start + Duration::from_secs(11)), + AdmissionOutcome::Admitted + ); + } + + #[test] + fn responsive_ring_overwrites_without_growing() { + let start = Instant::now(); + let mut reservoir = ResponsiveReservoir::new(2, Duration::from_secs(10)); + reservoir.record(node(1, "8.8.8.8:1"), start); + reservoir.record(node(2, "1.1.1.1:2"), start); + reservoir.record(node(3, "9.9.9.9:3"), start); + let snapshot = reservoir.snapshot(10, start); + assert_eq!(snapshot.len(), 2); + assert!(!snapshot.iter().any(|entry| entry.id == [1; 20])); + } + + #[test] + #[ignore = "release-only FIFO throughput smoke test"] + fn million_fifo_operations() { + let start = Instant::now(); + let mut pool = NodePool::new(100_000, u32::MAX, Duration::from_secs(600), start); + for value in 0..1_000_000u32 { + let octets = value.to_be_bytes(); + let node = NodeTuple { + id: [octets[3]; 20], + addr: SocketAddr::new( + IpAddr::V4(Ipv4Addr::new(11, octets[1], octets[2], octets[3])), + (value % 65_534 + 1) as u16, + ), + }; + let outcome = pool.admit(node, start); + assert!(!matches!(outcome, AdmissionOutcome::RateLimited)); + } + assert_eq!(pool.len(), 100_000); + eprintln!("1,000,000 FIFO admissions in {:?}", start.elapsed()); + } +} diff --git a/dht-crawler/src/peer_lookup.rs b/dht-crawler/src/peer_lookup.rs new file mode 100644 index 0000000..ba3b3d6 --- /dev/null +++ b/dht-crawler/src/peer_lookup.rs @@ -0,0 +1,783 @@ +// 负责执行有界主动 Peer 查找并向采集和调用方返回结果 + +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; +use crate::routing_snapshot::{RoutingSnapshot, xor_distance_cmp}; +use crate::runtime_stats::DhtRuntimeStats; +use crate::server::HashDiscovered; +use crate::types::{NetMode, NodeTuple, PeerLookupOptions, PeerLookupResult}; +use ahash::{AHashMap, AHashSet}; +use arc_swap::ArcSwap; +use bytes::BytesMut; +#[cfg(feature = "metrics")] +use metrics::counter; +use std::collections::VecDeque; +use std::net::SocketAddr; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use tokio::net::UdpSocket; +use tokio::sync::{mpsc, oneshot}; +use tokio_util::sync::CancellationToken; + +const LOOKUP_TID_TAG: u8 = 0xa5; +const MAX_QUERIES_PER_LOOKUP: usize = 12; +const MAX_CONCURRENT_QUERIES_PER_LOOKUP: usize = 4; +const MAX_FRONTIER_NODES: usize = 64; +const MAX_PEERS_PER_LOOKUP: usize = 12; +const LOOKUP_TIMEOUT: Duration = Duration::from_secs(2); +const QUERY_TIMEOUT: Duration = Duration::from_millis(500); +const MAINTENANCE_INTERVAL: Duration = Duration::from_millis(25); +const REQUEST_CHANNEL_CAPACITY: usize = 16_384; +const RESPONSE_CHANNEL_CAPACITY: usize = 4_096; + +#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)] +struct PendingKey { + addr: SocketAddr, + tid: TransactionId, +} + +#[derive(Debug, Clone, Copy)] +struct PendingQuery { + lookup_id: u64, + deadline: Instant, + preferred_phase: bool, +} + +struct LookupState { + info_hash: [u8; 20], + info_hash_hex: String, + frontier: Vec, + preferred: Option, + seen_nodes: AHashSet, + peers: AHashSet, + queried: usize, + outstanding: usize, + deadline: Instant, + preferred_phase: bool, + completion: Option>, +} + +impl LookupState { + fn pop_closest(&mut self) -> Option { + if let Some(preferred) = self.preferred.take() { + return Some(preferred); + } + let index = self + .frontier + .iter() + .enumerate() + .min_by(|(_, left), (_, right)| xor_distance_cmp(&left.id, &right.id, &self.info_hash)) + .map(|(index, _)| index)?; + Some(self.frontier.swap_remove(index)) + } + + fn is_complete(&self, now: Instant) -> bool { + self.deadline <= now + || self.peers.len() >= MAX_PEERS_PER_LOOKUP + || (self.outstanding == 0 + && (self.queried >= MAX_QUERIES_PER_LOOKUP || self.frontier.is_empty())) + } +} + +struct LookupResponse { + remote_addr: SocketAddr, + tid: TransactionId, + response: DhtResponse, +} + +pub(crate) struct PeerLookupRequest { + pub(crate) info_hash: [u8; 20], + pub(crate) preferred_node: Option, + pub(crate) completion: Option>, +} + +impl PeerLookupRequest { + pub(crate) fn new(info_hash: [u8; 20]) -> Self { + Self { + info_hash, + preferred_node: None, + completion: None, + } + } +} + +#[derive(Clone)] +pub(crate) struct PeerLookupHandle { + request_tx: mpsc::Sender, + response_tx: mpsc::Sender, + runtime_stats: DhtRuntimeStats, +} + +pub(crate) struct PeerLookupRuntime { + pub(crate) options: PeerLookupOptions, + pub(crate) stats: DhtRuntimeStats, + pub(crate) outbound_query_budget: SharedRateBudget, + pub(crate) shutdown: CancellationToken, +} + +impl PeerLookupHandle { + pub(crate) fn request_sender(&self) -> mpsc::Sender { + self.request_tx.clone() + } + + pub(crate) async fn lookup( + &self, + info_hash: [u8; 20], + ) -> Result { + let (completion, receiver) = oneshot::channel(); + self.request_tx + .send(PeerLookupRequest { + info_hash, + preferred_node: None, + completion: Some(completion), + }) + .await + .map_err(|_| "Peer Lookup 已停止")?; + receiver.await.map_err(|_| "Peer Lookup 未返回结果") + } + + pub(crate) fn route_response( + &self, + remote_addr: SocketAddr, + tid: TransactionId, + response: DhtResponse, + ) { + if self + .response_tx + .try_send(LookupResponse { + remote_addr, + tid, + response, + }) + .is_err() + { + self.runtime_stats.peer_lookup_response_dropped(); + #[cfg(feature = "metrics")] + counter!("dht_peer_lookup_dropped_total", "reason" => "response_queue_full") + .increment(1); + } + } +} + +pub(crate) fn is_peer_lookup_tid(tid: &TransactionId) -> bool { + tid[0] == LOOKUP_TID_TAG +} + +pub(crate) fn spawn_peer_lookup( + netmode: NetMode, + local_id: [u8; 20], + sockets: &std::collections::HashMap>, + snapshot: Arc>, + hash_tx: mpsc::Sender, + runtime: PeerLookupRuntime, +) -> PeerLookupHandle { + let PeerLookupRuntime { + options, + stats, + outbound_query_budget, + shutdown, + } = runtime; + let (request_tx, request_rx) = mpsc::channel(REQUEST_CHANNEL_CAPACITY); + let (response_tx, response_rx) = mpsc::channel(RESPONSE_CHANNEL_CAPACITY); + let socket_v4 = sockets + .iter() + .find_map(|(addr, socket)| addr.is_ipv4().then(|| socket.clone())); + let socket_v6 = sockets + .iter() + .find_map(|(addr, socket)| addr.is_ipv6().then(|| socket.clone())); + let actor = PeerLookupActor { + netmode, + local_id, + socket_v4, + socket_v6, + snapshot, + hash_tx, + request_rx, + response_rx, + request_budget: RateBucket::per_second( + options.max_lookups_per_second, + options.burst, + true, + Instant::now(), + ), + max_active_lookups: options.max_active_lookups, + enabled: options.max_lookups_per_second > 0 && options.max_active_lookups > 0, + queued: VecDeque::new(), + queued_hashes: AHashSet::new(), + active_hashes: AHashSet::new(), + active: AHashMap::new(), + pending: AHashMap::new(), + pending_expiry: VecDeque::new(), + next_lookup_id: 1, + next_tid: 1, + runtime_stats: stats.clone(), + outbound_query_budget, + shutdown, + }; + tokio::spawn(actor.run()); + PeerLookupHandle { + request_tx, + response_tx, + runtime_stats: stats, + } +} + +struct PeerLookupActor { + netmode: NetMode, + local_id: [u8; 20], + socket_v4: Option>, + socket_v6: Option>, + snapshot: Arc>, + hash_tx: mpsc::Sender, + request_rx: mpsc::Receiver, + response_rx: mpsc::Receiver, + request_budget: RateBucket, + max_active_lookups: usize, + enabled: bool, + queued: VecDeque, + queued_hashes: AHashSet<[u8; 20]>, + active_hashes: AHashSet<[u8; 20]>, + active: AHashMap, + pending: AHashMap, + pending_expiry: VecDeque<(Instant, PendingKey)>, + next_lookup_id: u64, + next_tid: u64, + runtime_stats: DhtRuntimeStats, + outbound_query_budget: SharedRateBudget, + shutdown: CancellationToken, +} + +impl PeerLookupActor { + async fn run(mut self) { + let mut maintenance = tokio::time::interval(MAINTENANCE_INTERVAL); + maintenance.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + tokio::select! { + biased; + _ = self.shutdown.cancelled() => break, + response = self.response_rx.recv() => { + let Some(response) = response else { break }; + self.handle_response(response, Instant::now()).await; + } + _ = maintenance.tick() => self.expire(Instant::now()).await, + request = self.request_rx.recv() => { + let Some(request) = request else { break }; + self.queue_request(request); + } + } + } + } + + fn queue_request(&mut self, request: PeerLookupRequest) { + self.runtime_stats.peer_lookup_requested(); + #[cfg(feature = "metrics")] + counter!("dht_peer_lookup_requests_total").increment(1); + if !self.enabled { + if let Some(completion) = request.completion { + let _ = completion.send(PeerLookupResult { + peers: Vec::new(), + queries: 0, + }); + } + return; + } + if self.queued_hashes.contains(&request.info_hash) + || self.active_hashes.contains(&request.info_hash) + { + if let Some(completion) = request.completion { + let _ = completion.send(PeerLookupResult { + peers: Vec::new(), + queries: 0, + }); + } + return; + } + if self.queued.len() >= REQUEST_CHANNEL_CAPACITY { + self.runtime_stats.peer_lookup_rate_limited(); + #[cfg(feature = "metrics")] + counter!("dht_peer_lookup_dropped_total", "reason" => "request_queue_full") + .increment(1); + if let Some(completion) = request.completion { + let _ = completion.send(PeerLookupResult { + peers: Vec::new(), + queries: 0, + }); + } + return; + } + self.queued_hashes.insert(request.info_hash); + self.queued.push_back(request); + } + + async fn start_queued(&mut self, now: Instant) { + while self.active.len() < self.max_active_lookups + && !self.queued.is_empty() + && self.request_budget.try_take_one(now) + { + let request = self.queued.pop_front().expect("queued request exists"); + self.queued_hashes.remove(&request.info_hash); + if !self.start_lookup(request, now).await { + self.request_budget.refund_one(); + } + } + } + + async fn start_lookup(&mut self, request: PeerLookupRequest, now: Instant) -> bool { + let info_hash = request.info_hash; + + let filter_ipv6 = match (self.socket_v4.is_some(), self.socket_v6.is_some()) { + (true, false) => Some(false), + (false, true) => Some(true), + _ => None, + }; + let mut frontier = + self.snapshot + .load() + .closest_nodes(&info_hash, MAX_QUERIES_PER_LOOKUP, filter_ipv6); + let preferred = request.preferred_node; + if let Some(preferred) = preferred { + frontier.retain(|node| node.addr != preferred.addr); + } + if frontier.is_empty() && preferred.is_none() { + self.runtime_stats.peer_lookup_empty(); + return false; + } + + let lookup_id = self.next_lookup_id; + self.next_lookup_id = self.next_lookup_id.wrapping_add(1).max(1); + let mut seen_nodes: AHashSet<_> = frontier.iter().map(|node| node.addr).collect(); + if let Some(preferred) = preferred { + seen_nodes.insert(preferred.addr); + } + self.active.insert( + lookup_id, + LookupState { + info_hash, + info_hash_hex: hex::encode(info_hash), + frontier, + preferred, + seen_nodes, + peers: AHashSet::new(), + queried: 0, + outstanding: 0, + deadline: now + LOOKUP_TIMEOUT, + preferred_phase: preferred.is_some(), + completion: request.completion, + }, + ); + self.active_hashes.insert(info_hash); + self.runtime_stats.peer_lookup_started(); + #[cfg(feature = "metrics")] + counter!("dht_peer_lookup_started_total").increment(1); + self.dispatch_more(lookup_id, now).await; + true + } + + async fn dispatch_more(&mut self, lookup_id: u64, now: Instant) { + loop { + if !self.outbound_query_budget.try_take_one(now) { + break; + } + 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, + tid, + }; + let mut buffer = BytesMut::with_capacity(128); + encode_get_peers_query(&mut buffer, &tid, &info_hash, &self.local_id); + let socket = if node.addr.is_ipv4() { + self.socket_v4.clone() + } else { + self.socket_v6.clone() + }; + let sent = match socket { + Some(socket) => socket.send_to(&buffer, node.addr).await.is_ok(), + 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; + } + + let deadline = now + QUERY_TIMEOUT; + self.pending.insert( + key, + PendingQuery { + lookup_id, + deadline, + preferred_phase, + }, + ); + self.pending_expiry.push_back((deadline, key)); + if let Some(state) = self.active.get_mut(&lookup_id) { + state.outstanding += 1; + } + self.runtime_stats.udp_sent(buffer.len()); + self.runtime_stats.peer_lookup_query(); + #[cfg(feature = "metrics")] + counter!("dht_peer_lookup_queries_total").increment(1); + } + self.finish_if_complete(lookup_id, now); + } + + async fn handle_response(&mut self, event: LookupResponse, now: Instant) { + let key = PendingKey { + addr: event.remote_addr, + tid: event.tid, + }; + let Some(pending) = self.pending.remove(&key) else { + return; + }; + let Some(state) = self.active.get_mut(&pending.lookup_id) else { + 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(); + for_each_response_peer(&event.response, self.netmode, |peer| { + if state.peers.len() < MAX_PEERS_PER_LOOKUP && state.peers.insert(peer) { + discovered.push(peer); + } + }); + let mut response_nodes = Vec::new(); + for_each_response_node(&event.response, self.netmode, |node| { + response_nodes.push(node) + }); + for node in response_nodes { + if state.frontier.len() >= MAX_FRONTIER_NODES { + break; + } + if state.seen_nodes.insert(node.addr) { + state.frontier.push(node); + } + } + 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 { + let event = HashDiscovered { + info_hash: hash.clone(), + peer_addr: peer, + discovered_at: now, + }; + if self.hash_tx.try_send(event).is_ok() { + self.runtime_stats.peer_lookup_peer_found(); + #[cfg(feature = "metrics")] + counter!("dht_peer_lookup_peers_found_total").increment(1); + } else { + self.runtime_stats.peer_lookup_output_dropped(); + #[cfg(feature = "metrics")] + counter!("dht_peer_lookup_dropped_total", "reason" => "hash_queue_full") + .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; + } + + async fn expire(&mut self, now: Instant) { + let mut affected = AHashSet::new(); + while let Some((deadline, key)) = self.pending_expiry.front().copied() { + if deadline > now { + break; + } + self.pending_expiry.pop_front(); + let should_remove = self + .pending + .get(&key) + .is_some_and(|pending| pending.deadline == deadline); + if !should_remove { + continue; + } + 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(); + } + for lookup_id in affected { + self.dispatch_more(lookup_id, now).await; + } + let expired: Vec<_> = self + .active + .iter() + .filter_map(|(lookup_id, state)| (state.deadline <= now).then_some(*lookup_id)) + .collect(); + for lookup_id in expired { + self.finish_lookup(lookup_id); + } + self.start_queued(now).await; + } + + fn finish_if_complete(&mut self, lookup_id: u64, now: Instant) { + if self + .active + .get(&lookup_id) + .is_some_and(|state| state.is_complete(now)) + { + self.finish_lookup(lookup_id); + } + } + + fn finish_lookup(&mut self, lookup_id: u64) { + if let Some(mut state) = self.active.remove(&lookup_id) { + self.active_hashes.remove(&state.info_hash); + if let Some(completion) = state.completion.take() { + let mut peers: Vec<_> = state.peers.into_iter().collect(); + peers.sort_unstable(); + let _ = completion.send(PeerLookupResult { + peers, + queries: state.queried, + }); + } + } + self.pending + .retain(|_, pending| pending.lookup_id != lookup_id); + } + + fn next_transaction_id(&mut self) -> TransactionId { + let mut tid = self.next_tid.to_be_bytes(); + tid[0] = LOOKUP_TID_TAG; + self.next_tid = self.next_tid.wrapping_add(1).max(1); + tid + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::protocol::DhtMessage; + + #[test] + fn lookup_transaction_ids_have_a_reserved_tag() { + let tid = [LOOKUP_TID_TAG, 1, 2, 3, 4, 5, 6, 7]; + assert!(is_peer_lookup_tid(&tid)); + assert!(!is_peer_lookup_tid(&[0; 8])); + } + + #[tokio::test] + 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: [8; 20], + addr: fallback_addr, + }], + 1, + ))); + let (hash_tx, mut 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(), + }, + ); + let (completion, result_rx) = oneshot::channel(); + handle + .request_tx + .send(PeerLookupRequest { + info_hash: [3; 20], + preferred_node: Some(NodeTuple { + id: [9; 20], + addr: remote_addr, + }), + completion: Some(completion), + }) + .await + .unwrap(); + + let mut buffer = [0u8; 512]; + let (len, source) = + tokio::time::timeout(Duration::from_secs(1), remote_socket.recv_from(&mut buffer)) + .await + .unwrap() + .unwrap(); + assert_eq!(source, local_addr); + 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, + DhtResponse { + id: Some(serde_bytes::ByteBuf::from(vec![9; 20])), + nodes: None, + nodes6: None, + values: Some(vec![serde_bytes::ByteBuf::from(vec![ + 8, 8, 4, 4, 0x1a, 0xe1, + ])]), + samples: None, + num: None, + interval: None, + }, + ); + + let event = tokio::time::timeout(Duration::from_secs(1), hash_rx.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!(event.info_hash, hex::encode([3; 20])); + assert_eq!(event.peer_addr, "8.8.4.4:6881".parse().unwrap()); + let result = tokio::time::timeout(Duration::from_secs(1), result_rx) + .await + .unwrap() + .unwrap(); + assert_eq!(result.peers, vec!["8.8.4.4:6881".parse().unwrap()]); + assert_eq!(result.queries, 1); + let snapshot = stats.snapshot(); + assert_eq!(snapshot.peer_lookup_started, 1); + 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, + }), + completion: None, + }) + .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/protocol.rs b/dht-crawler/src/protocol.rs new file mode 100644 index 0000000..ff291d5 --- /dev/null +++ b/dht-crawler/src/protocol.rs @@ -0,0 +1,63 @@ +use serde::Deserialize; + +#[derive(Deserialize, Debug, Clone)] +#[allow(dead_code)] +/// Decoded KRPC envelope. +pub struct DhtMessage { + /// Transaction ID bytes. + pub t: serde_bytes::ByteBuf, + #[allow(dead_code)] + /// Message kind (`q`, `r`, or `e`). + pub y: String, + #[allow(dead_code)] + /// Query method when `y == q`. + pub q: Option, + /// Query arguments. + pub a: Option, + /// Response dictionary. + pub r: Option, +} + +#[derive(Deserialize, Debug, Clone)] +/// Supported BEP-5 query arguments. +pub struct DhtArgs { + /// Sender node ID. + pub id: Option, + /// find_node target ID. + pub target: Option, + /// get_peers/announce InfoHash. + pub info_hash: Option, + /// announce validation token. + pub token: Option, + /// Explicit announced Peer port. + pub port: Option, + /// Non-zero means use the UDP source port. + pub implied_port: Option, +} + +#[derive(Deserialize, Debug, Clone)] +/// Supported BEP-5 response fields. +pub struct DhtResponse { + #[serde(default)] + #[allow(dead_code)] + /// Responder node ID. + pub id: Option, + #[serde(default)] + /// Compact IPv4 node tuples. + pub nodes: Option, + #[serde(default)] + /// Compact IPv6 node tuples. + pub nodes6: Option, + #[serde(default)] + /// Compact Peer endpoints returned by `get_peers`. + pub values: Option>, + #[serde(default)] + /// BEP-51 concatenated 20-byte sampled InfoHashes. + pub samples: Option, + #[serde(default)] + /// BEP-51 estimated number of InfoHashes held by the responder. + pub num: Option, + #[serde(default)] + /// BEP-51 requested delay before sampling this node again, in seconds. + pub interval: Option, +} diff --git a/dht-crawler/src/routing_snapshot.rs b/dht-crawler/src/routing_snapshot.rs new file mode 100644 index 0000000..16c2af9 --- /dev/null +++ b/dht-crawler/src/routing_snapshot.rs @@ -0,0 +1,107 @@ +use crate::types::NodeTuple; +use rand::seq::IndexedRandom; + +#[derive(Default)] +pub(crate) struct RoutingSnapshot { + v4: Vec, + v6: Vec, +} + +impl RoutingSnapshot { + pub(crate) fn from_nodes(nodes: Vec, limit: usize) -> Self { + let mut v4 = Vec::new(); + let mut v6 = Vec::new(); + for node in nodes.into_iter().take(limit) { + if node.addr.is_ipv6() { + v6.push(node); + } else { + v4.push(node); + } + } + Self { v4, v6 } + } + + pub(crate) fn random_nodes(&self, count: usize, filter_ipv6: Option) -> Vec { + let mut rng = rand::rng(); + match filter_ipv6 { + Some(true) => self.v6.sample(&mut rng, count).cloned().collect(), + Some(false) => self.v4.sample(&mut rng, count).cloned().collect(), + None => { + let mut all = Vec::with_capacity(self.v4.len() + self.v6.len()); + all.extend_from_slice(&self.v4); + all.extend_from_slice(&self.v6); + all.sample(&mut rng, count).cloned().collect() + } + } + } + + pub(crate) fn closest_nodes( + &self, + target: &[u8; 20], + count: usize, + filter_ipv6: Option, + ) -> Vec { + if count == 0 { + return Vec::new(); + } + let mut nodes = match filter_ipv6 { + Some(true) => self.v6.clone(), + Some(false) => self.v4.clone(), + None => { + let mut all = Vec::with_capacity(self.v4.len() + self.v6.len()); + all.extend_from_slice(&self.v4); + all.extend_from_slice(&self.v6); + all + } + }; + let compare = + |left: &NodeTuple, right: &NodeTuple| xor_distance_cmp(&left.id, &right.id, target); + if nodes.len() > count { + nodes.select_nth_unstable_by(count, compare); + nodes.truncate(count); + } + nodes.sort_unstable_by(compare); + nodes + } +} + +pub(crate) fn xor_distance_cmp( + left: &[u8; 20], + right: &[u8; 20], + target: &[u8; 20], +) -> std::cmp::Ordering { + for index in 0..20 { + let ordering = (left[index] ^ target[index]).cmp(&(right[index] ^ target[index])); + if ordering != std::cmp::Ordering::Equal { + return ordering; + } + } + std::cmp::Ordering::Equal +} + +#[cfg(test)] +mod tests { + use super::*; + use std::net::SocketAddr; + + fn node(id: u8, port: u16) -> NodeTuple { + NodeTuple { + id: [id; 20], + addr: SocketAddr::from(([8, 8, 8, 8], port)), + } + } + + #[test] + fn closest_nodes_orders_by_xor_distance_and_limits_results() { + let snapshot = + RoutingSnapshot::from_nodes(vec![node(0xf0, 1), node(0x01, 2), node(0x10, 3)], 3); + let closest = snapshot.closest_nodes(&[0; 20], 2, Some(false)); + assert_eq!( + closest + .iter() + .map(|node| node.addr.port()) + .collect::>(), + vec![2, 3] + ); + } +} diff --git a/dht-crawler/src/runtime_stats.rs b/dht-crawler/src/runtime_stats.rs new file mode 100644 index 0000000..8d2bd44 --- /dev/null +++ b/dht-crawler/src/runtime_stats.rs @@ -0,0 +1,1415 @@ +use std::sync::Arc; +use std::sync::atomic::{AtomicU32, AtomicU64, AtomicUsize, Ordering}; + +pub const METADATA_QUEUE_WAIT_BUCKETS_MS: [u64; 8] = [10, 50, 100, 250, 500, 1_000, 2_000, 5_000]; +pub const METADATA_FETCH_BUCKETS_MS: [u64; 7] = [250, 500, 1_000, 2_000, 4_000, 6_000, 10_000]; +pub const METADATA_SIZE_BUCKETS_BYTES: [u64; 8] = [ + 16 * 1024, + 32 * 1024, + 64 * 1024, + 128 * 1024, + 256 * 1024, + 512 * 1024, + 1024 * 1024, + 10 * 1024 * 1024, +]; + +#[derive(Debug, Clone, PartialEq, Eq)] +/// Snapshot of a non-cumulative fixed-bucket histogram. +pub struct FixedHistogramSnapshot { + /// Inclusive upper bound for each bucket. + pub bounds: Vec, + /// Per-bucket counts; values are not cumulative. + pub counts: Vec, + /// Values larger than the final bound. + pub overflow: u64, + /// Total observations including overflow. + pub count: u64, + /// Wrapping sum of all observed values. + pub sum: u64, +} + +impl FixedHistogramSnapshot { + /// Returns the inclusive bucket bound containing the requested approximate percentile. + /// + /// `percentile` is clamped to `0.0..=1.0`. Overflow observations return the final finite + /// bound because the histogram intentionally stores no dynamic maximum. + pub fn percentile(&self, percentile: f64) -> Option { + if self.count == 0 { + return None; + } + let rank = ((self.count as f64 * percentile.clamp(0.0, 1.0)).ceil() as u64).max(1); + let mut cumulative = 0u64; + for (bound, count) in self.bounds.iter().zip(&self.counts) { + cumulative = cumulative.saturating_add(*count); + if cumulative >= rank { + return Some(*bound); + } + } + self.bounds.last().copied() + } +} + +struct AtomicFixedHistogram { + bounds: [u64; N], + counts: [AtomicU64; N], + overflow: AtomicU64, + count: AtomicU64, + sum: AtomicU64, +} + +impl AtomicFixedHistogram { + fn new(bounds: [u64; N]) -> Self { + Self { + bounds, + counts: std::array::from_fn(|_| AtomicU64::new(0)), + overflow: AtomicU64::new(0), + count: AtomicU64::new(0), + sum: AtomicU64::new(0), + } + } + + fn record(&self, value: u64) { + if let Some(index) = self.bounds.iter().position(|bound| value <= *bound) { + self.counts[index].fetch_add(1, Ordering::Relaxed); + } else { + self.overflow.fetch_add(1, Ordering::Relaxed); + } + self.count.fetch_add(1, Ordering::Relaxed); + self.sum.fetch_add(value, Ordering::Relaxed); + } + + fn snapshot(&self) -> FixedHistogramSnapshot { + FixedHistogramSnapshot { + bounds: self.bounds.to_vec(), + counts: self + .counts + .iter() + .map(|count| count.load(Ordering::Relaxed)) + .collect(), + overflow: self.overflow.load(Ordering::Relaxed), + count: self.count.load(Ordering::Relaxed), + sum: self.sum.load(Ordering::Relaxed), + } + } +} + +impl Default for AtomicFixedHistogram { + fn default() -> Self { + Self::new([0; N]) + } +} + +/// Transport-neutral counters and fixed-bucket histograms used by dashboards. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DhtObservabilitySnapshot { + /// UDP datagrams received before validation. + pub udp_rx_packets: u64, + /// UDP bytes received before validation. + pub udp_rx_bytes: u64, + /// Successfully sent UDP query and response datagrams. + pub udp_tx_packets: u64, + /// Successfully sent UDP query and response bytes. + pub udp_tx_bytes: u64, + /// Inbound ping queries. + pub inbound_ping: u64, + /// Inbound find_node queries. + pub inbound_find_node: u64, + /// Inbound get_peers queries. + pub inbound_get_peers: u64, + /// Inbound announce_peer queries. + pub inbound_announce_peer: u64, + /// Other or invalid inbound query names. + pub inbound_other: u64, + /// Replies admitted by the regular response budget. + pub response_normal: u64, + /// Priority ping/get_peers replies admitted by the reserve. + pub response_fallback: u64, + /// Replies rejected by final response limits. + pub response_rate_limited: u64, + /// Admitted replies whose `send_to` failed. + pub response_send_failed: u64, + /// Valid, unfiltered announces with a usable Peer port. + pub announce_accepted: u64, + /// Announces with a missing or invalid token. + pub announce_invalid_token: u64, + /// Announces rejected by the application Hash filter. + pub announce_filtered: u64, + /// New nodes admitted to the FIFO pool. + pub node_admitted: u64, + /// Full-pool replacements. + pub node_replaced: u64, + /// Nodes rejected as queued or recently probed duplicates. + pub node_dropped_duplicate: u64, + /// Nodes rejected by the replacement budget. + pub node_dropped_rate_limited: u64, + /// Nodes rejected because their endpoint is not usable. + pub node_dropped_invalid: u64, + /// BEP-9 Metadata piece payload bytes received. + pub metadata_bytes_downloaded: u64, + /// Metadata Peer attempts reaching the end-to-end timeout. + pub metadata_failure_timeout: u64, + /// Metadata Peer connection failures. + pub metadata_failure_connect: u64, + /// Peers without extension-protocol support. + pub metadata_failure_no_extension: u64, + /// Extension handshake or piece-request send failures. + pub metadata_failure_send: u64, + /// Metadata payloads exceeding the 10 MiB limit. + pub metadata_failure_size_limit: u64, + /// Metadata payloads failing InfoHash SHA1 validation. + pub metadata_failure_sha1: u64, + /// Validated payloads that could not be parsed as torrent info. + pub metadata_failure_parse: u64, + /// Other receive, piece or worker failures. + pub metadata_failure_other: u64, + /// Failure-cache hits for previously timed-out Peers. + pub peer_cache_timeout_hits: u64, + /// Failure-cache hits for previous connection failures. + pub peer_cache_connect_hits: u64, + /// Queue wait observations in milliseconds. + pub queue_wait_ms: FixedHistogramSnapshot, + /// End-to-end Peer attempt observations in milliseconds. + pub fetch_duration_ms: FixedHistogramSnapshot, + /// Complete bencoded info payload observations in bytes. + pub metadata_size_bytes: FixedHistogramSnapshot, +} + +/// Cheap, cloneable handle for reading live DHT runtime statistics. +#[derive(Clone, Default)] +pub struct DhtRuntimeStats { + inner: Arc, +} + +/// Point-in-time view returned by [`DhtRuntimeStats::snapshot`]. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct DhtRuntimeSnapshot { + /// Valid announce hashes observed before ingress publication. + pub hashes_received: u64, + /// Hashes dropped because bounded ingress was full. + pub hash_ingress_dropped: u64, + /// Current hash-ingress queue depth. + pub hash_ingress_queue_depth: usize, + /// Configured hash-ingress queue capacity. + pub hash_ingress_queue_capacity: usize, + + /// Current priority crawl-event depth. + pub crawl_priority_queue_depth: usize, + /// Configured priority crawl-event capacity. + pub crawl_priority_queue_capacity: usize, + /// Current discovered-node event depth. + pub crawl_discovery_queue_depth: usize, + /// Configured discovered-node event capacity. + pub crawl_discovery_queue_capacity: usize, + + /// Deduplicated pending Metadata Hashes. + pub metadata_queue_depth: usize, + /// Configured pending Metadata capacity. + pub metadata_queue_max: usize, + /// Current Metadata jobs. + pub metadata_in_flight: usize, + /// Unique Hash insertions, including insertions that evicted an older Hash. + pub metadata_queue_inserted: u64, + /// Duplicate Hash updates merged into existing entries. + pub metadata_queue_deduplicated: u64, + /// Oldest Hashes evicted by newer events. + pub metadata_queue_evicted: u64, + /// Events rejected or expired because they were stale. + pub metadata_queue_stale: u64, + /// Metadata Peer races started after admission. + pub metadata_races_started: u64, + /// Distinct Peer candidates launched into Metadata races. + pub metadata_peer_candidates: u64, + /// Peer candidates added while a Metadata race was already running. + pub metadata_live_peer_joins: u64, + /// Losing Peer attempts canceled after another Peer succeeded. + pub metadata_peer_canceled: u64, + /// Active `get_peers` lookup requests received from Metadata jobs. + pub peer_lookup_requested: u64, + /// Active `get_peers` lookups admitted by the bounded rate limiter. + pub peer_lookup_started: u64, + /// Active lookup requests rejected by concurrency or rate limits. + pub peer_lookup_rate_limited: u64, + /// Active lookup requests skipped because the routing snapshot was empty. + pub peer_lookup_empty: u64, + /// Active `get_peers` UDP queries sent. + pub peer_lookup_queries: u64, + /// Matched active `get_peers` responses. + pub peer_lookup_responses: u64, + /// Active `get_peers` queries that timed out. + pub peer_lookup_timeouts: u64, + /// Active `get_peers` UDP sends that failed immediately. + pub peer_lookup_send_failures: u64, + /// Unique Peer endpoints fed back into Metadata scheduling. + pub peer_lookup_peers_found: u64, + /// Tagged lookup responses dropped before reaching the lookup actor. + 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. + pub sample_infohashes_responses: u64, + /// BEP-51 requests that timed out. + pub sample_infohashes_timeouts: u64, + /// BEP-51 UDP sends that failed immediately. + pub sample_infohashes_send_failures: u64, + /// BEP-51 responses dropped before reaching the sampler actor. + 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. + pub sample_infohashes_hashes_dropped: u64, + /// Real Peer network attempts. + pub metadata_peer_attempts: u64, + /// Successful Peer downloads and parses. + pub metadata_peer_succeeded: u64, + /// Failed real Peer attempts. + pub metadata_peer_failed: u64, + /// End-to-end Peer timeouts. + pub metadata_peer_timeouts: u64, + /// Peer connection failures. + pub metadata_connect_failed: u64, + /// Peers without extension-protocol support. + pub metadata_no_extension: u64, + /// Peer failure-cache skips. + pub metadata_peer_failure_cache_hits: u64, + /// Current Peer failure-cache entries. + pub metadata_peer_failure_cache_entries: usize, + + /// Current strict FIFO crawl-pool size. + pub node_pool_size: usize, + /// Configured FIFO capacity. + pub node_pool_capacity: usize, + /// Bootstrap low-water mark. + pub node_pool_low_watermark: usize, + + /// Current pending find_node transactions. + pub find_node_in_flight: usize, + /// Configured total find_node in-flight limit. + pub find_node_in_flight_max: usize, + /// Metadata-pressure-adjusted find_node budget per second. + pub find_node_effective_rate_per_sec: u32, + /// Queries sent to never-before-probed destinations. + pub queries_new: u64, + /// Queries sent to responsive revisit nodes. + pub queries_revisit: u64, + /// Queries sent to bootstrap endpoints. + pub queries_bootstrap: u64, + /// Replies matched to pending transactions. + pub responses: u64, + /// Replies without a matching pending transaction. + pub unmatched_responses: u64, + /// Pending transactions that expired. + pub timeouts: u64, + /// Outbound find_node send failures. + pub send_failures: u64, + /// Response events dropped before reaching the actor. + pub crawl_events_dropped_response: u64, + /// Discovered-node events dropped before reaching the actor. + pub crawl_events_dropped_discovered: u64, + + /// UDP datagrams received before validation. + pub udp_received: u64, + /// Valid datagrams dropped because every worker queue was full. + pub udp_queue_full: u64, + /// Empty, oversized or non-bencoded UDP datagrams. + pub udp_invalid: u64, + /// DHT replies rejected by final rate limits. + pub udp_responses_rate_limited: u64, + /// Priority replies admitted by the ping/get_peers reserve. + pub udp_responses_priority_reserved: u64, +} + +struct DhtRuntimeStatsInner { + hashes_received: AtomicU64, + hash_ingress_dropped: AtomicU64, + hash_ingress_queue_depth: AtomicUsize, + hash_ingress_queue_capacity: usize, + + crawl_priority_queue_depth: AtomicUsize, + crawl_priority_queue_capacity: usize, + crawl_discovery_queue_depth: AtomicUsize, + crawl_discovery_queue_capacity: usize, + + metadata_queue_depth: AtomicUsize, + metadata_queue_max: usize, + metadata_in_flight: AtomicUsize, + metadata_queue_inserted: AtomicU64, + metadata_queue_deduplicated: AtomicU64, + metadata_queue_evicted: AtomicU64, + metadata_queue_stale: AtomicU64, + metadata_races_started: AtomicU64, + metadata_peer_candidates: AtomicU64, + metadata_live_peer_joins: AtomicU64, + metadata_peer_canceled: AtomicU64, + peer_lookup_requested: AtomicU64, + peer_lookup_started: AtomicU64, + peer_lookup_rate_limited: AtomicU64, + peer_lookup_empty: AtomicU64, + peer_lookup_queries: AtomicU64, + peer_lookup_responses: AtomicU64, + peer_lookup_timeouts: AtomicU64, + peer_lookup_send_failures: AtomicU64, + 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, + metadata_peer_succeeded: AtomicU64, + metadata_peer_failed: AtomicU64, + metadata_peer_timeouts: AtomicU64, + metadata_connect_failed: AtomicU64, + metadata_no_extension: AtomicU64, + metadata_peer_failure_cache_hits: AtomicU64, + metadata_peer_failure_cache_entries: AtomicUsize, + + node_pool_size: AtomicUsize, + node_pool_capacity: usize, + node_pool_low_watermark: usize, + + find_node_in_flight: AtomicUsize, + find_node_in_flight_max: usize, + find_node_effective_rate_per_sec: AtomicU32, + queries_new: AtomicU64, + queries_revisit: AtomicU64, + queries_bootstrap: AtomicU64, + responses: AtomicU64, + unmatched_responses: AtomicU64, + timeouts: AtomicU64, + send_failures: AtomicU64, + crawl_events_dropped_response: AtomicU64, + crawl_events_dropped_discovered: AtomicU64, + + udp_received: AtomicU64, + udp_queue_full: AtomicU64, + udp_invalid: AtomicU64, + udp_responses_rate_limited: AtomicU64, + udp_responses_priority_reserved: AtomicU64, + + udp_rx_bytes: AtomicU64, + udp_tx_packets: AtomicU64, + udp_tx_bytes: AtomicU64, + inbound_ping: AtomicU64, + inbound_find_node: AtomicU64, + inbound_get_peers: AtomicU64, + inbound_announce_peer: AtomicU64, + inbound_other: AtomicU64, + response_normal: AtomicU64, + response_send_failed: AtomicU64, + announce_accepted: AtomicU64, + announce_invalid_token: AtomicU64, + announce_filtered: AtomicU64, + node_admitted: AtomicU64, + node_replaced: AtomicU64, + node_dropped_duplicate: AtomicU64, + node_dropped_rate_limited: AtomicU64, + node_dropped_invalid: AtomicU64, + metadata_bytes_downloaded: AtomicU64, + metadata_failure_timeout: AtomicU64, + metadata_failure_connect: AtomicU64, + metadata_failure_no_extension: AtomicU64, + metadata_failure_send: AtomicU64, + metadata_failure_size_limit: AtomicU64, + metadata_failure_sha1: AtomicU64, + metadata_failure_parse: AtomicU64, + metadata_failure_other: AtomicU64, + peer_cache_timeout_hits: AtomicU64, + peer_cache_connect_hits: AtomicU64, + queue_wait_ms: AtomicFixedHistogram<8>, + fetch_duration_ms: AtomicFixedHistogram<7>, + metadata_size_bytes: AtomicFixedHistogram<8>, +} + +impl Default for DhtRuntimeStatsInner { + fn default() -> Self { + Self { + queue_wait_ms: AtomicFixedHistogram::new(METADATA_QUEUE_WAIT_BUCKETS_MS), + fetch_duration_ms: AtomicFixedHistogram::new(METADATA_FETCH_BUCKETS_MS), + metadata_size_bytes: AtomicFixedHistogram::new(METADATA_SIZE_BUCKETS_BYTES), + hashes_received: AtomicU64::new(0), + hash_ingress_dropped: AtomicU64::new(0), + hash_ingress_queue_depth: AtomicUsize::new(0), + hash_ingress_queue_capacity: 0, + crawl_priority_queue_depth: AtomicUsize::new(0), + crawl_priority_queue_capacity: 0, + crawl_discovery_queue_depth: AtomicUsize::new(0), + crawl_discovery_queue_capacity: 0, + metadata_queue_depth: AtomicUsize::new(0), + metadata_queue_max: 0, + metadata_in_flight: AtomicUsize::new(0), + metadata_queue_inserted: AtomicU64::new(0), + metadata_queue_deduplicated: AtomicU64::new(0), + metadata_queue_evicted: AtomicU64::new(0), + metadata_queue_stale: AtomicU64::new(0), + metadata_races_started: AtomicU64::new(0), + metadata_peer_candidates: AtomicU64::new(0), + metadata_live_peer_joins: AtomicU64::new(0), + metadata_peer_canceled: AtomicU64::new(0), + peer_lookup_requested: AtomicU64::new(0), + peer_lookup_started: AtomicU64::new(0), + peer_lookup_rate_limited: AtomicU64::new(0), + peer_lookup_empty: AtomicU64::new(0), + peer_lookup_queries: AtomicU64::new(0), + peer_lookup_responses: AtomicU64::new(0), + peer_lookup_timeouts: AtomicU64::new(0), + peer_lookup_send_failures: AtomicU64::new(0), + 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), + metadata_peer_succeeded: AtomicU64::new(0), + metadata_peer_failed: AtomicU64::new(0), + metadata_peer_timeouts: AtomicU64::new(0), + metadata_connect_failed: AtomicU64::new(0), + metadata_no_extension: AtomicU64::new(0), + metadata_peer_failure_cache_hits: AtomicU64::new(0), + metadata_peer_failure_cache_entries: AtomicUsize::new(0), + node_pool_size: AtomicUsize::new(0), + node_pool_capacity: 0, + node_pool_low_watermark: 0, + find_node_in_flight: AtomicUsize::new(0), + find_node_in_flight_max: 0, + find_node_effective_rate_per_sec: AtomicU32::new(0), + queries_new: AtomicU64::new(0), + queries_revisit: AtomicU64::new(0), + queries_bootstrap: AtomicU64::new(0), + responses: AtomicU64::new(0), + unmatched_responses: AtomicU64::new(0), + timeouts: AtomicU64::new(0), + send_failures: AtomicU64::new(0), + crawl_events_dropped_response: AtomicU64::new(0), + crawl_events_dropped_discovered: AtomicU64::new(0), + udp_received: AtomicU64::new(0), + udp_queue_full: AtomicU64::new(0), + udp_invalid: AtomicU64::new(0), + udp_responses_rate_limited: AtomicU64::new(0), + udp_responses_priority_reserved: AtomicU64::new(0), + udp_rx_bytes: AtomicU64::new(0), + udp_tx_packets: AtomicU64::new(0), + udp_tx_bytes: AtomicU64::new(0), + inbound_ping: AtomicU64::new(0), + inbound_find_node: AtomicU64::new(0), + inbound_get_peers: AtomicU64::new(0), + inbound_announce_peer: AtomicU64::new(0), + inbound_other: AtomicU64::new(0), + response_normal: AtomicU64::new(0), + response_send_failed: AtomicU64::new(0), + announce_accepted: AtomicU64::new(0), + announce_invalid_token: AtomicU64::new(0), + announce_filtered: AtomicU64::new(0), + node_admitted: AtomicU64::new(0), + node_replaced: AtomicU64::new(0), + node_dropped_duplicate: AtomicU64::new(0), + node_dropped_rate_limited: AtomicU64::new(0), + node_dropped_invalid: AtomicU64::new(0), + metadata_bytes_downloaded: AtomicU64::new(0), + metadata_failure_timeout: AtomicU64::new(0), + metadata_failure_connect: AtomicU64::new(0), + metadata_failure_no_extension: AtomicU64::new(0), + metadata_failure_send: AtomicU64::new(0), + metadata_failure_size_limit: AtomicU64::new(0), + metadata_failure_sha1: AtomicU64::new(0), + metadata_failure_parse: AtomicU64::new(0), + metadata_failure_other: AtomicU64::new(0), + peer_cache_timeout_hits: AtomicU64::new(0), + peer_cache_connect_hits: AtomicU64::new(0), + } + } +} + +#[derive(Debug, Clone, Copy, Default)] +pub(crate) struct DhtRuntimeLimits { + pub metadata_queue: usize, + pub node_pool: usize, + pub node_pool_low_watermark: usize, + pub find_node_in_flight: usize, + pub initial_find_node_rate: u32, + pub hash_ingress_queue: usize, + pub crawl_priority_queue: usize, + pub crawl_discovery_queue: usize, +} + +impl DhtRuntimeStats { + pub(crate) fn with_limits(limits: DhtRuntimeLimits) -> Self { + Self { + inner: Arc::new(DhtRuntimeStatsInner { + metadata_queue_max: limits.metadata_queue, + node_pool_capacity: limits.node_pool, + node_pool_low_watermark: limits.node_pool_low_watermark, + find_node_in_flight_max: limits.find_node_in_flight, + find_node_effective_rate_per_sec: AtomicU32::new(limits.initial_find_node_rate), + hash_ingress_queue_capacity: limits.hash_ingress_queue, + crawl_priority_queue_capacity: limits.crawl_priority_queue, + crawl_discovery_queue_capacity: limits.crawl_discovery_queue, + ..DhtRuntimeStatsInner::default() + }), + } + } + + /// Load every exposed value once using relaxed atomics. + /// Loads the core runtime counters and gauges using relaxed atomics. + pub fn snapshot(&self) -> DhtRuntimeSnapshot { + let inner = &self.inner; + DhtRuntimeSnapshot { + hashes_received: inner.hashes_received.load(Ordering::Relaxed), + hash_ingress_dropped: inner.hash_ingress_dropped.load(Ordering::Relaxed), + hash_ingress_queue_depth: inner.hash_ingress_queue_depth.load(Ordering::Relaxed), + hash_ingress_queue_capacity: inner.hash_ingress_queue_capacity, + crawl_priority_queue_depth: inner.crawl_priority_queue_depth.load(Ordering::Relaxed), + crawl_priority_queue_capacity: inner.crawl_priority_queue_capacity, + crawl_discovery_queue_depth: inner.crawl_discovery_queue_depth.load(Ordering::Relaxed), + crawl_discovery_queue_capacity: inner.crawl_discovery_queue_capacity, + metadata_queue_depth: inner.metadata_queue_depth.load(Ordering::Relaxed), + metadata_queue_max: inner.metadata_queue_max, + metadata_in_flight: inner.metadata_in_flight.load(Ordering::Relaxed), + metadata_queue_inserted: inner.metadata_queue_inserted.load(Ordering::Relaxed), + metadata_queue_deduplicated: inner.metadata_queue_deduplicated.load(Ordering::Relaxed), + metadata_queue_evicted: inner.metadata_queue_evicted.load(Ordering::Relaxed), + metadata_queue_stale: inner.metadata_queue_stale.load(Ordering::Relaxed), + metadata_races_started: inner.metadata_races_started.load(Ordering::Relaxed), + metadata_peer_candidates: inner.metadata_peer_candidates.load(Ordering::Relaxed), + metadata_live_peer_joins: inner.metadata_live_peer_joins.load(Ordering::Relaxed), + metadata_peer_canceled: inner.metadata_peer_canceled.load(Ordering::Relaxed), + peer_lookup_requested: inner.peer_lookup_requested.load(Ordering::Relaxed), + peer_lookup_started: inner.peer_lookup_started.load(Ordering::Relaxed), + peer_lookup_rate_limited: inner.peer_lookup_rate_limited.load(Ordering::Relaxed), + peer_lookup_empty: inner.peer_lookup_empty.load(Ordering::Relaxed), + peer_lookup_queries: inner.peer_lookup_queries.load(Ordering::Relaxed), + peer_lookup_responses: inner.peer_lookup_responses.load(Ordering::Relaxed), + peer_lookup_timeouts: inner.peer_lookup_timeouts.load(Ordering::Relaxed), + peer_lookup_send_failures: inner.peer_lookup_send_failures.load(Ordering::Relaxed), + peer_lookup_peers_found: inner.peer_lookup_peers_found.load(Ordering::Relaxed), + peer_lookup_response_dropped: inner + .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), + sample_infohashes_send_failures: inner + .sample_infohashes_send_failures + .load(Ordering::Relaxed), + sample_infohashes_response_dropped: inner + .sample_infohashes_response_dropped + .load(Ordering::Relaxed), + 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), + sample_infohashes_hashes_dropped: inner + .sample_infohashes_hashes_dropped + .load(Ordering::Relaxed), + metadata_peer_attempts: inner.metadata_peer_attempts.load(Ordering::Relaxed), + metadata_peer_succeeded: inner.metadata_peer_succeeded.load(Ordering::Relaxed), + metadata_peer_failed: inner.metadata_peer_failed.load(Ordering::Relaxed), + metadata_peer_timeouts: inner.metadata_peer_timeouts.load(Ordering::Relaxed), + metadata_connect_failed: inner.metadata_connect_failed.load(Ordering::Relaxed), + metadata_no_extension: inner.metadata_no_extension.load(Ordering::Relaxed), + metadata_peer_failure_cache_hits: inner + .metadata_peer_failure_cache_hits + .load(Ordering::Relaxed), + metadata_peer_failure_cache_entries: inner + .metadata_peer_failure_cache_entries + .load(Ordering::Relaxed), + node_pool_size: inner.node_pool_size.load(Ordering::Relaxed), + node_pool_capacity: inner.node_pool_capacity, + node_pool_low_watermark: inner.node_pool_low_watermark, + find_node_in_flight: inner.find_node_in_flight.load(Ordering::Relaxed), + find_node_in_flight_max: inner.find_node_in_flight_max, + find_node_effective_rate_per_sec: inner + .find_node_effective_rate_per_sec + .load(Ordering::Relaxed), + queries_new: inner.queries_new.load(Ordering::Relaxed), + queries_revisit: inner.queries_revisit.load(Ordering::Relaxed), + queries_bootstrap: inner.queries_bootstrap.load(Ordering::Relaxed), + responses: inner.responses.load(Ordering::Relaxed), + unmatched_responses: inner.unmatched_responses.load(Ordering::Relaxed), + timeouts: inner.timeouts.load(Ordering::Relaxed), + send_failures: inner.send_failures.load(Ordering::Relaxed), + crawl_events_dropped_response: inner + .crawl_events_dropped_response + .load(Ordering::Relaxed), + crawl_events_dropped_discovered: inner + .crawl_events_dropped_discovered + .load(Ordering::Relaxed), + udp_received: inner.udp_received.load(Ordering::Relaxed), + udp_queue_full: inner.udp_queue_full.load(Ordering::Relaxed), + udp_invalid: inner.udp_invalid.load(Ordering::Relaxed), + udp_responses_rate_limited: inner.udp_responses_rate_limited.load(Ordering::Relaxed), + udp_responses_priority_reserved: inner + .udp_responses_priority_reserved + .load(Ordering::Relaxed), + } + } + + /// Loads detailed transport counters, failure categories and fixed histograms. + pub fn observability_snapshot(&self) -> DhtObservabilitySnapshot { + let inner = &self.inner; + DhtObservabilitySnapshot { + udp_rx_packets: inner.udp_received.load(Ordering::Relaxed), + udp_rx_bytes: inner.udp_rx_bytes.load(Ordering::Relaxed), + udp_tx_packets: inner.udp_tx_packets.load(Ordering::Relaxed), + udp_tx_bytes: inner.udp_tx_bytes.load(Ordering::Relaxed), + inbound_ping: inner.inbound_ping.load(Ordering::Relaxed), + inbound_find_node: inner.inbound_find_node.load(Ordering::Relaxed), + inbound_get_peers: inner.inbound_get_peers.load(Ordering::Relaxed), + inbound_announce_peer: inner.inbound_announce_peer.load(Ordering::Relaxed), + inbound_other: inner.inbound_other.load(Ordering::Relaxed), + response_normal: inner.response_normal.load(Ordering::Relaxed), + response_fallback: inner + .udp_responses_priority_reserved + .load(Ordering::Relaxed), + response_rate_limited: inner.udp_responses_rate_limited.load(Ordering::Relaxed), + response_send_failed: inner.response_send_failed.load(Ordering::Relaxed), + announce_accepted: inner.announce_accepted.load(Ordering::Relaxed), + announce_invalid_token: inner.announce_invalid_token.load(Ordering::Relaxed), + announce_filtered: inner.announce_filtered.load(Ordering::Relaxed), + node_admitted: inner.node_admitted.load(Ordering::Relaxed), + node_replaced: inner.node_replaced.load(Ordering::Relaxed), + node_dropped_duplicate: inner.node_dropped_duplicate.load(Ordering::Relaxed), + node_dropped_rate_limited: inner.node_dropped_rate_limited.load(Ordering::Relaxed), + node_dropped_invalid: inner.node_dropped_invalid.load(Ordering::Relaxed), + metadata_bytes_downloaded: inner.metadata_bytes_downloaded.load(Ordering::Relaxed), + metadata_failure_timeout: inner.metadata_failure_timeout.load(Ordering::Relaxed), + metadata_failure_connect: inner.metadata_failure_connect.load(Ordering::Relaxed), + metadata_failure_no_extension: inner + .metadata_failure_no_extension + .load(Ordering::Relaxed), + metadata_failure_send: inner.metadata_failure_send.load(Ordering::Relaxed), + metadata_failure_size_limit: inner.metadata_failure_size_limit.load(Ordering::Relaxed), + metadata_failure_sha1: inner.metadata_failure_sha1.load(Ordering::Relaxed), + metadata_failure_parse: inner.metadata_failure_parse.load(Ordering::Relaxed), + metadata_failure_other: inner.metadata_failure_other.load(Ordering::Relaxed), + peer_cache_timeout_hits: inner.peer_cache_timeout_hits.load(Ordering::Relaxed), + peer_cache_connect_hits: inner.peer_cache_connect_hits.load(Ordering::Relaxed), + queue_wait_ms: inner.queue_wait_ms.snapshot(), + fetch_duration_ms: inner.fetch_duration_ms.snapshot(), + metadata_size_bytes: inner.metadata_size_bytes.snapshot(), + } + } + + pub(crate) fn hash_received(&self) { + self.inner.hashes_received.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn hash_ingress_dropped(&self) { + self.inner + .hash_ingress_dropped + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn set_hash_ingress_queue_depth(&self, depth: usize) { + self.inner + .hash_ingress_queue_depth + .store(depth, Ordering::Relaxed); + } + + pub(crate) fn set_crawl_priority_queue_depth(&self, depth: usize) { + self.inner + .crawl_priority_queue_depth + .store(depth, Ordering::Relaxed); + } + + pub(crate) fn set_crawl_discovery_queue_depth(&self, depth: usize) { + self.inner + .crawl_discovery_queue_depth + .store(depth, Ordering::Relaxed); + } + + pub(crate) fn set_metadata_queue(&self, depth: usize, in_flight: usize) { + self.inner + .metadata_queue_depth + .store(depth, Ordering::Relaxed); + self.inner + .metadata_in_flight + .store(in_flight, Ordering::Relaxed); + } + + pub(crate) fn metadata_queue_deduplicated(&self) { + self.inner + .metadata_queue_deduplicated + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn metadata_queue_inserted(&self) { + self.inner + .metadata_queue_inserted + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn metadata_queue_evicted(&self) { + self.inner + .metadata_queue_evicted + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn metadata_queue_stale(&self, count: usize) { + self.inner + .metadata_queue_stale + .fetch_add(count as u64, Ordering::Relaxed); + } + + pub(crate) fn metadata_race_started(&self) { + self.inner + .metadata_races_started + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn metadata_peer_candidate(&self) { + self.inner + .metadata_peer_candidates + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn metadata_live_peer_join(&self) { + self.inner + .metadata_live_peer_joins + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn metadata_peer_canceled(&self, count: usize) { + self.inner + .metadata_peer_canceled + .fetch_add(count as u64, Ordering::Relaxed); + } + + pub(crate) fn peer_lookup_requested(&self) { + self.inner + .peer_lookup_requested + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn peer_lookup_started(&self) { + self.inner + .peer_lookup_started + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn peer_lookup_rate_limited(&self) { + self.inner + .peer_lookup_rate_limited + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn peer_lookup_empty(&self) { + self.inner.peer_lookup_empty.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn peer_lookup_query(&self) { + self.inner + .peer_lookup_queries + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn peer_lookup_response(&self) { + self.inner + .peer_lookup_responses + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn peer_lookup_timeout(&self) { + self.inner + .peer_lookup_timeouts + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn peer_lookup_send_failed(&self) { + self.inner + .peer_lookup_send_failures + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn peer_lookup_peer_found(&self) { + self.inner + .peer_lookup_peers_found + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn peer_lookup_response_dropped(&self) { + self.inner + .peer_lookup_response_dropped + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn peer_lookup_output_dropped(&self) { + self.inner + .peer_lookup_output_dropped + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn peer_lookup_preferred_succeeded(&self) { + self.inner + .peer_lookup_preferred_succeeded + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn peer_lookup_fallback(&self) { + self.inner + .peer_lookup_fallbacks + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn sample_query(&self) { + self.inner + .sample_infohashes_queries + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn sample_response(&self) { + self.inner + .sample_infohashes_responses + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn sample_timeout(&self) { + self.inner + .sample_infohashes_timeouts + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn sample_send_failed(&self) { + self.inner + .sample_infohashes_send_failures + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn sample_response_dropped(&self) { + self.inner + .sample_infohashes_response_dropped + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn sample_hash_discovered(&self) { + self.inner + .sample_infohashes_hashes_discovered + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn sample_hash_filtered(&self, count: usize) { + self.inner + .sample_infohashes_hashes_filtered + .fetch_add(count as u64, Ordering::Relaxed); + } + + pub(crate) fn sample_hash_duplicate(&self) { + self.inner + .sample_infohashes_hashes_duplicate + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn sample_hash_dropped(&self) { + self.inner + .sample_infohashes_hashes_dropped + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn metadata_peer_attempt(&self) { + self.inner + .metadata_peer_attempts + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn metadata_peer_succeeded(&self) { + self.inner + .metadata_peer_succeeded + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn metadata_peer_failed(&self) { + self.inner + .metadata_peer_failed + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn metadata_peer_timeout(&self) { + self.inner + .metadata_peer_timeouts + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn metadata_connect_failed(&self) { + self.inner + .metadata_connect_failed + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn metadata_no_extension(&self) { + self.inner + .metadata_no_extension + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn metadata_peer_failure_cache_hit(&self) { + self.inner + .metadata_peer_failure_cache_hits + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn set_metadata_peer_failure_cache_entries(&self, count: usize) { + self.inner + .metadata_peer_failure_cache_entries + .store(count, Ordering::Relaxed); + } + + pub(crate) fn set_node_pool_size(&self, size: usize) { + self.inner.node_pool_size.store(size, Ordering::Relaxed); + } + + pub(crate) fn set_find_node_in_flight(&self, count: usize) { + self.inner + .find_node_in_flight + .store(count, Ordering::Relaxed); + } + + pub(crate) fn set_find_node_effective_rate(&self, rate: u32) { + self.inner + .find_node_effective_rate_per_sec + .store(rate, Ordering::Relaxed); + } + + pub(crate) fn query_new(&self) { + self.inner.queries_new.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn query_revisit(&self) { + self.inner.queries_revisit.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn query_bootstrap(&self) { + self.inner.queries_bootstrap.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn response(&self) { + self.inner.responses.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn unmatched_response(&self) { + self.inner + .unmatched_responses + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn timeout(&self) { + self.inner.timeouts.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn send_failure(&self) { + self.inner.send_failures.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn crawl_event_dropped_response(&self) { + self.inner + .crawl_events_dropped_response + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn crawl_event_dropped_discovered(&self) { + self.inner + .crawl_events_dropped_discovered + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn udp_received(&self) { + self.inner.udp_received.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn udp_received_bytes(&self, bytes: usize) { + self.inner + .udp_rx_bytes + .fetch_add(bytes as u64, Ordering::Relaxed); + } + + pub(crate) fn udp_sent(&self, bytes: usize) { + self.inner.udp_tx_packets.fetch_add(1, Ordering::Relaxed); + self.inner + .udp_tx_bytes + .fetch_add(bytes as u64, Ordering::Relaxed); + } + + pub(crate) fn inbound_query(&self, query: &str) { + let counter = match query { + "ping" => &self.inner.inbound_ping, + "find_node" => &self.inner.inbound_find_node, + "get_peers" => &self.inner.inbound_get_peers, + "announce_peer" => &self.inner.inbound_announce_peer, + _ => &self.inner.inbound_other, + }; + counter.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn response_normal(&self) { + self.inner.response_normal.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn response_send_failed(&self) { + self.inner + .response_send_failed + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn announce_accepted(&self) { + self.inner.announce_accepted.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn announce_invalid_token(&self) { + self.inner + .announce_invalid_token + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn announce_filtered(&self) { + self.inner.announce_filtered.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn node_admitted(&self) { + self.inner.node_admitted.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn node_replaced(&self) { + self.inner.node_replaced.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn node_dropped_duplicate(&self) { + self.inner + .node_dropped_duplicate + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn node_dropped_rate_limited(&self) { + self.inner + .node_dropped_rate_limited + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn node_dropped_invalid(&self) { + self.inner + .node_dropped_invalid + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn metadata_bytes_downloaded(&self, bytes: usize) { + self.inner + .metadata_bytes_downloaded + .fetch_add(bytes as u64, Ordering::Relaxed); + } + + pub(crate) fn metadata_failure_timeout(&self) { + self.inner + .metadata_failure_timeout + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn metadata_failure_connect(&self) { + self.inner + .metadata_failure_connect + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn metadata_failure_no_extension(&self) { + self.inner + .metadata_failure_no_extension + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn metadata_failure_send(&self) { + self.inner + .metadata_failure_send + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn metadata_failure_size_limit(&self) { + self.inner + .metadata_failure_size_limit + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn metadata_failure_sha1(&self) { + self.inner + .metadata_failure_sha1 + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn metadata_failure_parse(&self) { + self.inner + .metadata_failure_parse + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn metadata_failure_other(&self) { + self.inner + .metadata_failure_other + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn peer_cache_hit_timeout(&self) { + self.inner + .peer_cache_timeout_hits + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn peer_cache_hit_connect(&self) { + self.inner + .peer_cache_connect_hits + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn observe_metadata_queue_wait(&self, millis: u64) { + self.inner.queue_wait_ms.record(millis); + } + + pub(crate) fn observe_metadata_fetch_duration(&self, millis: u64) { + self.inner.fetch_duration_ms.record(millis); + } + + pub(crate) fn observe_metadata_size(&self, bytes: usize) { + self.inner.metadata_size_bytes.record(bytes as u64); + } + + pub(crate) fn udp_queue_full(&self) { + self.inner.udp_queue_full.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn udp_invalid(&self) { + self.inner.udp_invalid.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn udp_response_rate_limited(&self) { + self.inner + .udp_responses_rate_limited + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn udp_response_priority_reserved(&self) { + self.inner + .udp_responses_priority_reserved + .fetch_add(1, Ordering::Relaxed); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn cloned_handle_updates_one_snapshot() { + let stats = DhtRuntimeStats::with_limits(DhtRuntimeLimits { + metadata_queue: 100, + node_pool: 1_000, + node_pool_low_watermark: 10, + find_node_in_flight: 512, + initial_find_node_rate: 200, + hash_ingress_queue: 20, + crawl_priority_queue: 30, + crawl_discovery_queue: 40, + }); + let writer = stats.clone(); + + writer.hash_received(); + writer.hash_ingress_dropped(); + writer.set_hash_ingress_queue_depth(11); + writer.set_crawl_priority_queue_depth(12); + writer.set_crawl_discovery_queue_depth(13); + writer.set_metadata_queue(42, 7); + writer.metadata_queue_inserted(); + writer.metadata_queue_deduplicated(); + writer.metadata_queue_evicted(); + writer.metadata_queue_stale(3); + writer.metadata_race_started(); + writer.metadata_peer_candidate(); + writer.metadata_live_peer_join(); + writer.metadata_peer_canceled(2); + writer.peer_lookup_requested(); + writer.peer_lookup_started(); + writer.peer_lookup_rate_limited(); + writer.peer_lookup_empty(); + writer.peer_lookup_query(); + writer.peer_lookup_response(); + writer.peer_lookup_timeout(); + writer.peer_lookup_send_failed(); + writer.peer_lookup_peer_found(); + writer.peer_lookup_response_dropped(); + writer.peer_lookup_output_dropped(); + writer.peer_lookup_preferred_succeeded(); + writer.peer_lookup_fallback(); + writer.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(); + writer.metadata_peer_succeeded(); + writer.metadata_peer_failed(); + writer.metadata_peer_timeout(); + writer.metadata_connect_failed(); + writer.metadata_no_extension(); + writer.metadata_peer_failure_cache_hit(); + writer.set_metadata_peer_failure_cache_entries(9); + writer.set_node_pool_size(321); + writer.set_find_node_in_flight(12); + writer.set_find_node_effective_rate(150); + writer.query_new(); + writer.query_revisit(); + writer.query_bootstrap(); + writer.response(); + writer.unmatched_response(); + writer.timeout(); + writer.send_failure(); + writer.crawl_event_dropped_response(); + writer.crawl_event_dropped_discovered(); + writer.udp_received(); + writer.udp_queue_full(); + writer.udp_invalid(); + writer.udp_response_rate_limited(); + writer.udp_response_priority_reserved(); + + assert_eq!( + stats.snapshot(), + DhtRuntimeSnapshot { + hashes_received: 1, + hash_ingress_dropped: 1, + hash_ingress_queue_depth: 11, + hash_ingress_queue_capacity: 20, + crawl_priority_queue_depth: 12, + crawl_priority_queue_capacity: 30, + crawl_discovery_queue_depth: 13, + crawl_discovery_queue_capacity: 40, + metadata_queue_depth: 42, + metadata_queue_max: 100, + metadata_in_flight: 7, + metadata_queue_inserted: 1, + metadata_queue_deduplicated: 1, + metadata_queue_evicted: 1, + metadata_queue_stale: 3, + metadata_races_started: 1, + metadata_peer_candidates: 1, + metadata_live_peer_joins: 1, + metadata_peer_canceled: 2, + peer_lookup_requested: 1, + peer_lookup_started: 1, + peer_lookup_rate_limited: 1, + peer_lookup_empty: 1, + peer_lookup_queries: 1, + peer_lookup_responses: 1, + peer_lookup_timeouts: 1, + peer_lookup_send_failures: 1, + peer_lookup_peers_found: 1, + peer_lookup_response_dropped: 1, + peer_lookup_output_dropped: 1, + peer_lookup_preferred_succeeded: 1, + peer_lookup_fallbacks: 1, + sample_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, + metadata_peer_succeeded: 1, + metadata_peer_failed: 1, + metadata_peer_timeouts: 1, + metadata_connect_failed: 1, + metadata_no_extension: 1, + metadata_peer_failure_cache_hits: 1, + metadata_peer_failure_cache_entries: 9, + node_pool_size: 321, + node_pool_capacity: 1_000, + node_pool_low_watermark: 10, + find_node_in_flight: 12, + find_node_in_flight_max: 512, + find_node_effective_rate_per_sec: 150, + queries_new: 1, + queries_revisit: 1, + queries_bootstrap: 1, + responses: 1, + unmatched_responses: 1, + timeouts: 1, + send_failures: 1, + crawl_events_dropped_response: 1, + crawl_events_dropped_discovered: 1, + udp_received: 1, + udp_queue_full: 1, + udp_invalid: 1, + udp_responses_rate_limited: 1, + udp_responses_priority_reserved: 1, + } + ); + } + + #[test] + fn observability_snapshot_tracks_fixed_histograms_and_categories() { + let stats = DhtRuntimeStats::default(); + stats.udp_received(); + stats.udp_received_bytes(128); + stats.udp_sent(64); + stats.inbound_query("ping"); + stats.inbound_query("unknown"); + stats.node_admitted(); + stats.metadata_bytes_downloaded(1_024); + stats.metadata_failure_parse(); + stats.observe_metadata_queue_wait(75); + stats.observe_metadata_queue_wait(600); + stats.observe_metadata_fetch_duration(1_500); + stats.observe_metadata_size(70_000); + + let snapshot = stats.observability_snapshot(); + assert_eq!(snapshot.udp_rx_packets, 1); + assert_eq!(snapshot.udp_rx_bytes, 128); + assert_eq!(snapshot.udp_tx_packets, 1); + assert_eq!(snapshot.udp_tx_bytes, 64); + assert_eq!(snapshot.inbound_ping, 1); + assert_eq!(snapshot.inbound_other, 1); + assert_eq!(snapshot.node_admitted, 1); + assert_eq!(snapshot.metadata_bytes_downloaded, 1_024); + assert_eq!(snapshot.metadata_failure_parse, 1); + assert_eq!(snapshot.queue_wait_ms.percentile(0.50), Some(100)); + assert_eq!(snapshot.queue_wait_ms.percentile(0.95), Some(1_000)); + assert_eq!(snapshot.fetch_duration_ms.percentile(0.95), Some(2_000)); + assert_eq!(snapshot.metadata_size_bytes.percentile(0.50), Some(131_072)); + } +} diff --git a/dht-crawler/src/sample_infohashes.rs b/dht-crawler/src/sample_infohashes.rs new file mode 100644 index 0000000..b7c5ee7 --- /dev/null +++ b/dht-crawler/src/sample_infohashes.rs @@ -0,0 +1,547 @@ +// 负责执行有界 BEP-51 采样并将准入后的 infohash 交给 Peer 查找 + +use crate::budget::{RateBucket, SharedRateBudget}; +use crate::crawl_engine::CrawlEngine; +use crate::krpc::{encode_sample_infohashes_query, for_each_response_node}; +use crate::node_id::{TransactionId, random_node_id}; +use crate::peer_lookup::{PeerLookupHandle, PeerLookupRequest}; +use crate::protocol::DhtResponse; +use crate::routing_snapshot::RoutingSnapshot; +use crate::runtime_stats::DhtRuntimeStats; +use crate::types::{NetMode, NodeTuple, SampleInfohashesOptions}; +use ahash::{AHashMap, AHashSet}; +use arc_swap::{ArcSwap, ArcSwapOption}; +use bytes::BytesMut; +#[cfg(feature = "metrics")] +use metrics::{counter, gauge}; +use std::collections::VecDeque; +use std::future::Future; +use std::net::SocketAddr; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use tokio::net::UdpSocket; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; + +const SAMPLE_TID_TAG: u8 = 0x51; +const RESPONSE_CHANNEL_CAPACITY: usize = 4_096; +const ADMISSION_BATCH_CHANNEL_CAPACITY: usize = 64; +const MAINTENANCE_INTERVAL: Duration = Duration::from_millis(25); +const PRODUCTIVE_REVISIT: Duration = Duration::from_secs(60); +const MAX_PROTOCOL_INTERVAL: Duration = Duration::from_secs(6 * 60 * 60); + +#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)] +struct PendingKey { + addr: SocketAddr, + tid: TransactionId, +} + +#[derive(Debug, Clone, Copy)] +struct PendingRequest { + deadline: Instant, +} + +struct SampleResponse { + remote_addr: SocketAddr, + tid: TransactionId, + 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, + runtime_stats: DhtRuntimeStats, +} + +impl SampleInfohashesHandle { + pub(crate) fn route_response( + &self, + remote_addr: SocketAddr, + tid: TransactionId, + response: DhtResponse, + ) { + if self + .response_tx + .try_send(SampleResponse { + remote_addr, + tid, + response, + }) + .is_err() + { + self.runtime_stats.sample_response_dropped(); + #[cfg(feature = "metrics")] + counter!("dht_sample_infohashes_dropped_total", "reason" => "response_queue_full") + .increment(1); + } + } +} + +pub(crate) fn is_sample_infohashes_tid(tid: &TransactionId) -> bool { + tid[0] == SAMPLE_TID_TAG +} + +pub(crate) struct SampleInfohashesRuntime { + pub(crate) options: SampleInfohashesOptions, + pub(crate) stats: DhtRuntimeStats, + pub(crate) outbound_query_budget: SharedRateBudget, + pub(crate) hash_admission: Arc>, + pub(crate) shutdown: CancellationToken, +} + +pub(crate) fn spawn_sample_infohashes( + netmode: NetMode, + local_id: [u8; 20], + sockets: &std::collections::HashMap>, + snapshot: Arc>, + crawl_engine: Arc, + peer_lookup: PeerLookupHandle, + runtime: SampleInfohashesRuntime, +) -> SampleInfohashesHandle { + let SampleInfohashesRuntime { + options, + stats, + outbound_query_budget, + hash_admission, + shutdown, + } = runtime; + let (response_tx, response_rx) = mpsc::channel(RESPONSE_CHANNEL_CAPACITY); + let (admission_tx, admission_rx) = mpsc::channel(ADMISSION_BATCH_CHANNEL_CAPACITY); + let socket_v4 = sockets + .iter() + .find_map(|(addr, socket)| addr.is_ipv4().then(|| socket.clone())); + let socket_v6 = sockets + .iter() + .find_map(|(addr, socket)| addr.is_ipv6().then(|| socket.clone())); + let now = Instant::now(); + spawn_sample_admission( + admission_rx, + hash_admission, + peer_lookup.request_sender(), + stats.clone(), + shutdown.clone(), + ); + let actor = SampleInfohashesActor { + netmode, + local_id, + socket_v4, + socket_v6, + snapshot, + crawl_engine, + admission_tx, + response_rx, + query_budget: RateBucket::per_second( + options.max_queries_per_second, + options.burst, + true, + now, + ), + max_in_flight: options.max_in_flight, + request_timeout: Duration::from_millis(options.request_timeout_millis.max(100)), + unsupported_backoff: Duration::from_secs(options.unsupported_backoff_secs.max(1)), + dedup_capacity: options.dedup_capacity, + seen_hashes: AHashSet::new(), + seen_order: VecDeque::new(), + next_allowed: AHashMap::new(), + pending_addrs: AHashSet::new(), + pending: AHashMap::new(), + pending_expiry: VecDeque::new(), + next_tid: 1, + runtime_stats: stats.clone(), + outbound_query_budget, + shutdown, + }; + tokio::spawn(actor.run()); + SampleInfohashesHandle { + response_tx, + runtime_stats: stats, + } +} + +struct SampleInfohashesActor { + netmode: NetMode, + local_id: [u8; 20], + socket_v4: Option>, + socket_v6: Option>, + snapshot: Arc>, + crawl_engine: Arc, + admission_tx: mpsc::Sender, + response_rx: mpsc::Receiver, + query_budget: RateBucket, + max_in_flight: usize, + request_timeout: Duration, + unsupported_backoff: Duration, + dedup_capacity: usize, + seen_hashes: AHashSet<[u8; 20]>, + seen_order: VecDeque<[u8; 20]>, + next_allowed: AHashMap, + pending_addrs: AHashSet, + pending: AHashMap, + 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), + completion: None, + }; + let sent = tokio::select! { + _ = shutdown.cancelled() => false, + result = peer_lookup_tx.send(request) => result.is_ok(), + }; + if !sent { + return; + } + runtime_stats.sample_hash_discovered(); + } + } + }); +} + +impl SampleInfohashesActor { + async fn run(mut self) { + let mut maintenance = tokio::time::interval(MAINTENANCE_INTERVAL); + maintenance.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + tokio::select! { + biased; + _ = self.shutdown.cancelled() => break, + response = self.response_rx.recv() => { + let Some(response) = response else { break }; + self.handle_response(response, Instant::now()); + } + _ = maintenance.tick() => { + let now = Instant::now(); + self.expire(now); + self.dispatch(now).await; + } + } + } + } + + async fn dispatch(&mut self, now: Instant) { + if self.max_in_flight == 0 + || self.pending.len() >= self.max_in_flight + || self.admission_tx.capacity() == 0 + { + return; + } + let available = self.max_in_flight.saturating_sub(self.pending.len()); + let budget = self.query_budget.try_take(available, now); + if budget == 0 { + return; + } + let filter_ipv6 = match (self.socket_v4.is_some(), self.socket_v6.is_some()) { + (true, false) => Some(false), + (false, true) => Some(true), + _ => None, + }; + let candidates = self + .snapshot + .load() + .random_nodes((budget * 8).max(64), filter_ipv6); + let mut sent_count = 0usize; + for node in candidates { + if sent_count >= budget { + break; + } + if self.pending_addrs.contains(&node.addr) + || self + .next_allowed + .get(&node.addr) + .is_some_and(|deadline| *deadline > now) + { + continue; + } + if self.send_query(node, now).await { + sent_count += 1; + } + } + if sent_count < budget { + self.query_budget.refund(budget - sent_count); + } + } + + async fn send_query(&mut self, node: NodeTuple, now: Instant) -> bool { + let socket = if node.addr.is_ipv4() { + self.socket_v4.clone() + } else { + self.socket_v6.clone() + }; + let Some(socket) = socket else { + return false; + }; + if !self.outbound_query_budget.try_take_one(now) { + return false; + } + let tid = self.next_transaction_id(); + let mut buffer = BytesMut::with_capacity(128); + encode_sample_infohashes_query(&mut buffer, &tid, &random_node_id(), &self.local_id); + if socket.send_to(&buffer, node.addr).await.is_err() { + self.outbound_query_budget.refund_one(); + self.runtime_stats.sample_send_failed(); + self.next_allowed + .insert(node.addr, now + self.unsupported_backoff); + return false; + } + let key = PendingKey { + addr: node.addr, + tid, + }; + let deadline = now + self.request_timeout; + self.pending.insert(key, PendingRequest { deadline }); + self.pending_expiry.push_back((deadline, key)); + self.pending_addrs.insert(node.addr); + self.runtime_stats.udp_sent(buffer.len()); + self.runtime_stats.sample_query(); + #[cfg(feature = "metrics")] + { + counter!("dht_sample_infohashes_queries_total").increment(1); + gauge!("dht_sample_infohashes_in_flight").set(self.pending.len() as f64); + } + true + } + + fn handle_response(&mut self, event: SampleResponse, now: Instant) { + let key = PendingKey { + addr: event.remote_addr, + tid: event.tid, + }; + if self.pending.remove(&key).is_none() { + return; + } + self.pending_addrs.remove(&event.remote_addr); + self.runtime_stats.sample_response(); + + for_each_response_node(&event.response, self.netmode, |node| { + self.crawl_engine.route_discovered(node) + }); + + let responder_id = event + .response + .id + .as_deref() + .and_then(|id| <[u8; 20]>::try_from(id.as_slice()).ok()) + .unwrap_or([0; 20]); + let preferred_node = NodeTuple { + id: responder_id, + addr: event.remote_addr, + }; + let mut hashes = Vec::new(); + let mut batch_hashes = AHashSet::new(); + if let Some(samples) = event.response.samples.as_deref() + && samples.len() % 20 == 0 + { + for chunk in samples.chunks_exact(20) { + let hash: [u8; 20] = chunk.try_into().expect("sample hash is 20 bytes"); + if self.seen_hashes.contains(&hash) { + self.runtime_stats.sample_hash_duplicate(); + continue; + } + if batch_hashes.insert(hash) { + hashes.push(hash); + } + } + } + + let candidate_count = hashes.len(); + let admitted_to_triage = if hashes.is_empty() { + false + } else { + let hashes_to_remember = hashes.clone(); + if self + .admission_tx + .try_send(SampleAdmissionBatch { + preferred_node, + hashes, + }) + .is_ok() + { + for hash in hashes_to_remember { + self.remember_hash(hash); + } + true + } else { + for _ in 0..candidate_count { + self.runtime_stats.sample_hash_dropped(); + } + #[cfg(feature = "metrics")] + counter!("dht_sample_infohashes_dropped_total", "reason" => "admission_queue_full") + .increment(candidate_count as u64); + false + } + }; + let discovered = usize::from(admitted_to_triage) * candidate_count; + + let protocol_interval = Duration::from_secs(event.response.interval.unwrap_or(300)) + .clamp(Duration::from_secs(10), MAX_PROTOCOL_INTERVAL); + let delay = if discovered > 0 { + protocol_interval.min(PRODUCTIVE_REVISIT) + } else { + protocol_interval.saturating_add(self.unsupported_backoff) + }; + self.next_allowed.insert(event.remote_addr, now + delay); + #[cfg(feature = "metrics")] + { + counter!("dht_sample_infohashes_responses_total").increment(1); + counter!("dht_sample_infohashes_hashes_total").increment(discovered as u64); + gauge!("dht_sample_infohashes_in_flight").set(self.pending.len() as f64); + } + } + + fn remember_hash(&mut self, hash: [u8; 20]) { + if self.dedup_capacity == 0 { + return; + } + if self.seen_hashes.insert(hash) { + self.seen_order.push_back(hash); + } + while self.seen_order.len() > self.dedup_capacity { + if let Some(expired) = self.seen_order.pop_front() { + self.seen_hashes.remove(&expired); + } + } + } + + fn expire(&mut self, now: Instant) { + while let Some((deadline, key)) = self.pending_expiry.front().copied() { + if deadline > now { + break; + } + self.pending_expiry.pop_front(); + if !self + .pending + .get(&key) + .is_some_and(|pending| pending.deadline == deadline) + { + continue; + } + self.pending.remove(&key); + self.pending_addrs.remove(&key.addr); + self.next_allowed + .insert(key.addr, now + self.unsupported_backoff); + self.runtime_stats.sample_timeout(); + #[cfg(feature = "metrics")] + counter!("dht_sample_infohashes_timeouts_total").increment(1); + } + #[cfg(feature = "metrics")] + gauge!("dht_sample_infohashes_in_flight").set(self.pending.len() as f64); + } + + fn next_transaction_id(&mut self) -> TransactionId { + let mut tid = self.next_tid.to_be_bytes(); + tid[0] = SAMPLE_TID_TAG; + self.next_tid = self.next_tid.wrapping_add(1).max(1); + tid + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::protocol::DhtMessage; + + #[test] + fn sample_transaction_ids_have_a_reserved_tag() { + assert!(is_sample_infohashes_tid(&[ + SAMPLE_TID_TAG, + 1, + 2, + 3, + 4, + 5, + 6, + 7 + ])); + assert!(!is_sample_infohashes_tid(&[0; 8])); + } + + #[test] + fn bep51_response_fields_decode() { + let bytes = b"d1:rd2:id20:aaaaaaaaaaaaaaaaaaaa8:intervali60e3:numi2e7:samples40:bbbbbbbbbbbbbbbbbbbbcccccccccccccccccccce1:t8:123456781:y1:re"; + let message: DhtMessage = serde_bencode::from_bytes(bytes).unwrap(); + let response = message.r.unwrap(); + assert_eq!(response.interval, Some(60)); + assert_eq!(response.num, Some(2)); + assert_eq!(response.samples.unwrap().len(), 40); + } + + #[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/scheduler.rs b/dht-crawler/src/scheduler.rs new file mode 100644 index 0000000..4d10064 --- /dev/null +++ b/dht-crawler/src/scheduler.rs @@ -0,0 +1,1335 @@ +use crate::metadata::{FetchedMetadata, MetadataFetchOutcome, RbitFetcher}; +use crate::peer_lookup::PeerLookupRequest; +#[cfg(test)] +use crate::runtime_stats::DhtRuntimeLimits; +use crate::runtime_stats::DhtRuntimeStats; +use crate::server::HashDiscovered; +use crate::types::{MetadataFetchCompletion, MetadataFetchCompletionStatus, TorrentInfo}; +use arc_swap::ArcSwapOption; +#[cfg(feature = "metrics")] +use metrics::{counter, gauge, histogram}; +use std::collections::{BTreeMap, HashMap, HashSet}; +use std::future::Future; +use std::net::SocketAddr; +use std::panic::{AssertUnwindSafe, catch_unwind}; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::time::{Duration, Instant}; +use tokio::sync::mpsc; +use tokio::task::JoinSet; +use tokio_util::sync::CancellationToken; + +/// Maximum distinct Peer endpoints retained and raced for one InfoHash. +pub const MAX_METADATA_PEERS_PER_HASH: usize = 12; +const HASH_QUEUE_TTL: Duration = Duration::from_secs(60); +const PEER_COALESCE_WINDOW: Duration = Duration::from_millis(25); +const LIVE_PEER_IDLE_GRACE: Duration = Duration::from_millis(300); +const ACTIVE_PEER_LOOKUP_DELAY: Duration = Duration::from_millis(100); +const DISPATCH_TICK: Duration = Duration::from_millis(25); + +/// Callback returning whether the application accepted a downloaded torrent. +pub type TorrentAckCallback = Box bool + Send + Sync + 'static>; +/// Callback invoked once when an admitted Metadata job reaches a terminal state. +pub type MetadataCompletionCallback = Box; +/// Asynchronous InfoHash admission callback. +pub type MetadataFetchCallback = Box< + dyn Fn(String) -> std::pin::Pin + Send>> + + Send + + Sync + + 'static, +>; + +#[derive(Debug, Clone, Copy)] +/// Pending-queue capacity and maximum concurrent Metadata jobs. +pub struct MetadataSchedulerLimits { + /// Maximum deduplicated pending InfoHashes. + pub queue_size: usize, + /// Maximum concurrently spawned jobs. + pub concurrency: usize, +} + +#[derive(Clone)] +/// Atomically replaceable callbacks shared with [`MetadataScheduler`]. +pub struct MetadataSchedulerCallbacks { + /// Torrent delivery callback. + pub torrent: Arc>, + /// Pre-download admission callback. + pub fetch_gate: Arc>, + /// Terminal completion callback. + pub completion: Arc>, +} + +pub(crate) struct MetadataSchedulerRuntime { + pub(crate) stats: DhtRuntimeStats, + pub(crate) peer_lookup_tx: Option>, +} + +#[derive(Debug, Clone, Copy)] +struct PeerCandidate { + addr: SocketAddr, + discovered_at: Instant, +} + +#[derive(Debug)] +struct QueuedHash { + info_hash: String, + peers: Vec, + queued_at: Instant, + ready_at: Instant, + latest_at: Instant, + order_key: (Instant, u64), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum QueuePushKind { + Inserted, + Updated, + EvictedOldest, + Stale, +} + +#[derive(Debug)] +struct PendingHashQueue { + capacity: usize, + ttl: Duration, + entries: HashMap, + order: BTreeMap<(Instant, u64), String>, + sequence: u64, +} + +impl PendingHashQueue { + fn new(capacity: usize, ttl: Duration) -> Self { + Self { + capacity: capacity.max(1), + ttl, + entries: HashMap::new(), + order: BTreeMap::new(), + sequence: 0, + } + } + + fn len(&self) -> usize { + self.entries.len() + } + + fn is_empty(&self) -> bool { + self.entries.is_empty() + } + + #[cfg(test)] + fn contains(&self, info_hash: &str) -> bool { + self.entries.contains_key(info_hash) + } + + fn next_order_key(&mut self, at: Instant) -> (Instant, u64) { + self.sequence = self.sequence.wrapping_add(1); + (at, self.sequence) + } + + fn push(&mut self, event: HashDiscovered, now: Instant) -> QueuePushKind { + if now + .checked_duration_since(event.discovered_at) + .unwrap_or_default() + > self.ttl + { + return QueuePushKind::Stale; + } + + if let Some(mut entry) = self.entries.remove(&event.info_hash) { + self.order.remove(&entry.order_key); + let discovered_at = entry + .peers + .iter() + .find(|peer| peer.addr == event.peer_addr) + .map(|peer| peer.discovered_at.max(event.discovered_at)) + .unwrap_or(event.discovered_at); + entry.peers.retain(|peer| peer.addr != event.peer_addr); + entry.peers.push(PeerCandidate { + addr: event.peer_addr, + discovered_at, + }); + entry + .peers + .sort_unstable_by_key(|peer| std::cmp::Reverse(peer.discovered_at)); + entry.peers.truncate(MAX_METADATA_PEERS_PER_HASH); + entry.latest_at = entry.latest_at.max(event.discovered_at); + entry.order_key = self.next_order_key(entry.latest_at); + self.order.insert(entry.order_key, entry.info_hash.clone()); + self.entries.insert(entry.info_hash.clone(), entry); + return QueuePushKind::Updated; + } + + let mut result = QueuePushKind::Inserted; + if self.entries.len() >= self.capacity { + let Some((&oldest_key, oldest_hash)) = self.order.first_key_value() else { + return QueuePushKind::Stale; + }; + if event.discovered_at <= oldest_key.0 { + return QueuePushKind::Stale; + } + let oldest_hash = oldest_hash.clone(); + self.order.remove(&oldest_key); + self.entries.remove(&oldest_hash); + result = QueuePushKind::EvictedOldest; + } + + let order_key = self.next_order_key(event.discovered_at); + let info_hash = event.info_hash; + self.order.insert(order_key, info_hash.clone()); + self.entries.insert( + info_hash.clone(), + QueuedHash { + info_hash, + peers: vec![PeerCandidate { + addr: event.peer_addr, + discovered_at: event.discovered_at, + }], + queued_at: now, + ready_at: now + PEER_COALESCE_WINDOW, + latest_at: event.discovered_at, + order_key, + }, + ); + result + } + + fn remove(&mut self, info_hash: &str) -> Option { + let entry = self.entries.remove(info_hash)?; + self.order.remove(&entry.order_key); + Some(entry) + } + + fn pop_newest_ready( + &mut self, + now: Instant, + in_flight: &HashMap, + ) -> Option { + let info_hash = self.order.iter().rev().find_map(|(_, hash)| { + let entry = self.entries.get(hash)?; + (!in_flight.contains_key(hash) && entry.ready_at <= now).then(|| hash.clone()) + })?; + self.remove(&info_hash) + } + + fn expire(&mut self, now: Instant) -> usize { + let mut expired = 0; + while let Some((&oldest_key, oldest_hash)) = self.order.first_key_value() { + if now.checked_duration_since(oldest_key.0).unwrap_or_default() <= self.ttl { + break; + } + let oldest_hash = oldest_hash.clone(); + self.order.remove(&oldest_key); + self.entries.remove(&oldest_hash); + expired += 1; + } + expired + } +} + +#[derive(Debug)] +struct MetadataJob { + info_hash: String, + peers: Vec, + peer_rx: mpsc::Receiver, +} + +#[derive(Debug)] +struct InFlightState { + peer_tx: mpsc::Sender, + scheduled_peers: HashSet, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum JobOutcome { + GateRejected, + Completed(MetadataFetchCompletionStatus), +} + +#[derive(Debug)] +struct JobResult { + info_hash: String, + attempts: usize, + outcome: JobOutcome, +} + +/// Bounded, deduplicating, freshness-aware Metadata job scheduler. +pub struct MetadataScheduler { + hash_rx: mpsc::Receiver, + max_queue_size: usize, + max_concurrent: usize, + fetcher: Arc, + callback: Arc>, + on_metadata_fetch: Arc>, + completion_callback: Arc>, + total_received: Arc, + total_dropped: Arc, + total_dispatched: Arc, + total_completed: Arc, + queue_len: Arc, + runtime_stats: DhtRuntimeStats, + peer_lookup_tx: Option>, + shutdown: CancellationToken, +} + +impl MetadataScheduler { + /// Creates a scheduler using local runtime statistics. + pub fn new( + hash_rx: mpsc::Receiver, + fetcher: Arc, + limits: MetadataSchedulerLimits, + callbacks: MetadataSchedulerCallbacks, + queue_len: Arc, + shutdown: CancellationToken, + ) -> Self { + Self::new_with_runtime_stats( + hash_rx, + fetcher, + limits, + callbacks, + queue_len, + shutdown, + DhtRuntimeStats::default(), + ) + } + + pub(crate) fn new_with_runtime_stats( + hash_rx: mpsc::Receiver, + fetcher: Arc, + limits: MetadataSchedulerLimits, + callbacks: MetadataSchedulerCallbacks, + queue_len: Arc, + shutdown: CancellationToken, + runtime_stats: DhtRuntimeStats, + ) -> Self { + Self::new_with_runtime( + hash_rx, + fetcher, + limits, + callbacks, + queue_len, + shutdown, + MetadataSchedulerRuntime { + stats: runtime_stats, + peer_lookup_tx: None, + }, + ) + } + + pub(crate) fn new_with_runtime( + hash_rx: mpsc::Receiver, + fetcher: Arc, + limits: MetadataSchedulerLimits, + callbacks: MetadataSchedulerCallbacks, + queue_len: Arc, + shutdown: CancellationToken, + runtime: MetadataSchedulerRuntime, + ) -> Self { + Self { + hash_rx, + max_queue_size: limits.queue_size.max(1), + max_concurrent: limits.concurrency.max(1), + fetcher, + callback: callbacks.torrent, + on_metadata_fetch: callbacks.fetch_gate, + completion_callback: callbacks.completion, + total_received: Arc::new(AtomicU64::new(0)), + total_dropped: Arc::new(AtomicU64::new(0)), + total_dispatched: Arc::new(AtomicU64::new(0)), + total_completed: Arc::new(AtomicU64::new(0)), + queue_len, + runtime_stats: runtime.stats, + peer_lookup_tx: runtime.peer_lookup_tx, + shutdown, + } + } + + /// Replaces the torrent delivery callback. + pub fn set_callback(&mut self, callback: Arc) { + self.callback.store(Some(callback)); + } + + /// Replaces the pre-download admission callback. + pub fn set_metadata_fetch_callback(&mut self, callback: Arc) { + self.on_metadata_fetch.store(Some(callback)); + } + + /// Replaces the terminal completion callback. + pub fn set_completion_callback(&mut self, callback: Arc) { + self.completion_callback.store(Some(callback)); + } + + /// Runs until cancellation and drains completed jobs before returning when input closes. + pub async fn run(mut self) { + let mut queue = PendingHashQueue::new(self.max_queue_size, HASH_QUEUE_TTL); + let mut in_flight = HashMap::::new(); + let mut tasks = JoinSet::::new(); + let mut maintenance = tokio::time::interval(Duration::from_secs(1)); + maintenance.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + let mut dispatch_tick = tokio::time::interval(DISPATCH_TICK); + dispatch_tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + let mut stats_interval = tokio::time::interval(Duration::from_secs(60)); + stats_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + let mut input_closed = false; + + loop { + self.dispatch_ready(&mut queue, &mut in_flight, &mut tasks); + self.sync_queue_len(queue.len(), in_flight.len()); + + if input_closed && queue.is_empty() && tasks.is_empty() { + break; + } + + tokio::select! { + _ = self.shutdown.cancelled() => break, + result = self.hash_rx.recv(), if !input_closed => { + self.runtime_stats + .set_hash_ingress_queue_depth(self.hash_rx.len()); + match result { + Some(hash) => self.enqueue(&mut queue, &mut in_flight, hash), + None => input_closed = true, + } + } + result = tasks.join_next(), if !tasks.is_empty() => { + match result { + Some(Ok(result)) => { + self.handle_job_result(result, &mut queue, &mut in_flight); + } + Some(Err(error)) => { + #[cfg(feature = "metrics")] + counter!("dht_metadata_worker_join_error_total").increment(1); + log::warn!("Metadata worker task failed: {error}"); + } + None => {} + } + } + _ = dispatch_tick.tick() => {} + _ = maintenance.tick() => { + let expired = queue.expire(Instant::now()); + if expired > 0 { + self.total_dropped.fetch_add(expired as u64, Ordering::Relaxed); + self.runtime_stats.metadata_queue_stale(expired); + #[cfg(feature = "metrics")] + counter!("dht_metadata_queue_events_total", "result" => "expired") + .increment(expired as u64); + } + } + _ = stats_interval.tick() => self.print_stats_inline(), + } + } + + tasks.abort_all(); + while tasks.join_next().await.is_some() {} + self.queue_len.store(0, Ordering::Relaxed); + self.runtime_stats.set_hash_ingress_queue_depth(0); + self.runtime_stats.set_metadata_queue(0, 0); + #[cfg(feature = "metrics")] + gauge!("dht_metadata_queue_depth").set(0.0); + } + + fn enqueue( + &self, + queue: &mut PendingHashQueue, + in_flight: &mut HashMap, + hash: HashDiscovered, + ) { + self.total_received.fetch_add(1, Ordering::Relaxed); + let now = Instant::now(); + if let Some(state) = in_flight.get_mut(&hash.info_hash) { + let kind = if now + .checked_duration_since(hash.discovered_at) + .unwrap_or_default() + > HASH_QUEUE_TTL + { + QueuePushKind::Stale + } else if state.scheduled_peers.contains(&hash.peer_addr) + || state.scheduled_peers.len() >= MAX_METADATA_PEERS_PER_HASH + { + QueuePushKind::Updated + } else { + let peer = PeerCandidate { + addr: hash.peer_addr, + discovered_at: hash.discovered_at, + }; + match state.peer_tx.try_send(peer) { + Ok(()) => { + state.scheduled_peers.insert(hash.peer_addr); + self.runtime_stats.metadata_live_peer_join(); + #[cfg(feature = "metrics")] + counter!("dht_metadata_live_peer_joins_total").increment(1); + QueuePushKind::Updated + } + Err(_) => QueuePushKind::Updated, + } + }; + self.record_queue_push(kind); + return; + } + + let kind = queue.push(hash, now); + self.record_queue_push(kind); + } + + fn record_queue_push(&self, kind: QueuePushKind) { + let result = match kind { + QueuePushKind::Inserted => { + self.runtime_stats.metadata_queue_inserted(); + "inserted" + } + QueuePushKind::Updated => { + self.runtime_stats.metadata_queue_deduplicated(); + "deduplicated" + } + QueuePushKind::EvictedOldest => { + self.total_dropped.fetch_add(1, Ordering::Relaxed); + self.runtime_stats.metadata_queue_inserted(); + self.runtime_stats.metadata_queue_evicted(); + "evicted_oldest" + } + QueuePushKind::Stale => { + self.total_dropped.fetch_add(1, Ordering::Relaxed); + self.runtime_stats.metadata_queue_stale(1); + "stale" + } + }; + #[cfg(feature = "metrics")] + counter!("dht_metadata_queue_events_total", "result" => result).increment(1); + #[cfg(not(feature = "metrics"))] + let _ = result; + } + + fn dispatch_ready( + &self, + queue: &mut PendingHashQueue, + in_flight: &mut HashMap, + tasks: &mut JoinSet, + ) { + while tasks.len() < self.max_concurrent { + let Some(mut entry) = queue.pop_newest_ready(Instant::now(), in_flight) else { + break; + }; + let queue_wait = Instant::now() + .checked_duration_since(entry.queued_at) + .unwrap_or_default(); + self.runtime_stats.observe_metadata_queue_wait( + queue_wait.as_millis().min(u128::from(u64::MAX)) as u64, + ); + #[cfg(feature = "metrics")] + histogram!("dht_metadata_queue_wait_seconds").record(queue_wait.as_secs_f64()); + let now = Instant::now(); + entry.peers.retain(|peer| { + now.checked_duration_since(peer.discovered_at) + .unwrap_or_default() + <= HASH_QUEUE_TTL + }); + entry.peers.truncate(MAX_METADATA_PEERS_PER_HASH); + if entry.peers.is_empty() { + self.total_dropped.fetch_add(1, Ordering::Relaxed); + self.runtime_stats.metadata_queue_stale(1); + continue; + } + let (peer_tx, peer_rx) = mpsc::channel(MAX_METADATA_PEERS_PER_HASH); + let state = InFlightState { + peer_tx, + scheduled_peers: entry.peers.iter().map(|peer| peer.addr).collect(), + }; + in_flight.insert(entry.info_hash.clone(), state); + self.spawn_job(entry, peer_rx, tasks); + } + } + + fn spawn_job( + &self, + entry: QueuedHash, + peer_rx: mpsc::Receiver, + tasks: &mut JoinSet, + ) { + self.total_dispatched.fetch_add(1, Ordering::Relaxed); + #[cfg(feature = "metrics")] + counter!("dht_metadata_jobs_dispatched_total").increment(1); + + let job = MetadataJob { + info_hash: entry.info_hash, + peers: entry.peers, + peer_rx, + }; + let fetcher = self.fetcher.clone(); + let callback = self.callback.clone(); + let on_metadata_fetch = self.on_metadata_fetch.clone(); + let runtime_stats = self.runtime_stats.clone(); + let peer_lookup_tx = self.peer_lookup_tx.clone(); + tasks.spawn(async move { + Self::process_hash( + job, + &fetcher, + &callback, + &on_metadata_fetch, + &runtime_stats, + peer_lookup_tx, + ) + .await + }); + } + + fn handle_job_result( + &self, + result: JobResult, + queue: &mut PendingHashQueue, + in_flight: &mut HashMap, + ) { + if !in_flight.contains_key(&result.info_hash) { + return; + } + + match result.outcome { + JobOutcome::GateRejected => { + queue.remove(&result.info_hash); + in_flight.remove(&result.info_hash); + #[cfg(feature = "metrics")] + counter!("dht_metadata_jobs_completed_total", "result" => "gate_rejected") + .increment(1); + } + JobOutcome::Completed(status) => { + queue.remove(&result.info_hash); + in_flight.remove(&result.info_hash); + self.finish(result.info_hash, status, result.attempts); + } + } + } + + fn finish(&self, info_hash: String, status: MetadataFetchCompletionStatus, attempts: usize) { + self.total_completed.fetch_add(1, Ordering::Relaxed); + let result = match status { + MetadataFetchCompletionStatus::Accepted => "accepted", + MetadataFetchCompletionStatus::FetchFailed => "fetch_failed", + MetadataFetchCompletionStatus::DeliveryRejected => "delivery_rejected", + }; + #[cfg(feature = "metrics")] + counter!("dht_metadata_jobs_completed_total", "result" => result).increment(1); + #[cfg(not(feature = "metrics"))] + let _ = result; + + if let Some(callback) = self.completion_callback.load_full() { + let completion = MetadataFetchCompletion { + info_hash, + status, + attempts, + }; + if catch_unwind(AssertUnwindSafe(|| callback(completion))).is_err() { + #[cfg(feature = "metrics")] + counter!("dht_metadata_completion_callback_panics_total").increment(1); + log::warn!("Metadata completion callback panicked"); + } + } + } + + async fn process_hash( + mut job: MetadataJob, + fetcher: &Arc, + callback: &Arc>, + on_metadata_fetch: &Arc>, + runtime_stats: &DhtRuntimeStats, + peer_lookup_tx: Option>, + ) -> JobResult { + if let Some(gate) = on_metadata_fetch.load_full() + && !gate(job.info_hash.clone()).await + { + return JobResult { + info_hash: job.info_hash, + attempts: 0, + outcome: JobOutcome::GateRejected, + }; + } + + let info_hash_bytes: [u8; 20] = match hex::decode(&job.info_hash) { + Ok(bytes) if bytes.len() == 20 => { + let mut hash = [0u8; 20]; + hash.copy_from_slice(&bytes); + hash + } + _ => { + return JobResult { + info_hash: job.info_hash, + attempts: 0, + outcome: JobOutcome::Completed(MetadataFetchCompletionStatus::FetchFailed), + }; + } + }; + + let attempts = Arc::new(AtomicUsize::new(0)); + let attempts_for_fetch = attempts.clone(); + let fetcher_for_race = fetcher.clone(); + let lookup_task = peer_lookup_tx.map(|peer_lookup_tx| { + tokio::spawn(async move { + tokio::time::sleep(ACTIVE_PEER_LOOKUP_DELAY).await; + let _ = peer_lookup_tx.try_send(PeerLookupRequest::new(info_hash_bytes)); + }) + }); + let fetched = race_peer_fetches( + job.peers, + &mut job.peer_rx, + move |peer| { + let attempts = attempts_for_fetch.clone(); + let fetcher = fetcher_for_race.clone(); + async move { + let outcome = fetcher + .fetch_with_attempt_observer(&info_hash_bytes, peer.addr, move || { + attempts.fetch_add(1, Ordering::Relaxed); + }) + .await; + (peer, outcome) + } + }, + runtime_stats, + ) + .await; + if let Some(task) = lookup_task { + task.abort(); + } + let attempts = attempts.load(Ordering::Relaxed); + + if let Some((peer, (name, total_size, files, piece_length))) = fetched { + let metadata = TorrentInfo { + info_hash: job.info_hash.clone(), + name, + total_size, + files, + magnet_link: format!("magnet:?xt=urn:btih:{}", job.info_hash), + peers: vec![peer.addr.to_string()], + piece_length, + timestamp: std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs(), + }; + + let accepted = callback.load_full().is_some_and(|callback| { + catch_unwind(AssertUnwindSafe(|| callback(metadata))).unwrap_or(false) + }); + let status = if accepted { + MetadataFetchCompletionStatus::Accepted + } else { + MetadataFetchCompletionStatus::DeliveryRejected + }; + return JobResult { + info_hash: job.info_hash, + attempts, + outcome: JobOutcome::Completed(status), + }; + } + + JobResult { + info_hash: job.info_hash, + attempts, + outcome: JobOutcome::Completed(MetadataFetchCompletionStatus::FetchFailed), + } + } + + fn sync_queue_len(&self, queued: usize, in_flight: usize) { + let depth = queued.saturating_add(in_flight); + self.queue_len.store(depth, Ordering::Relaxed); + self.runtime_stats.set_metadata_queue(depth, in_flight); + #[cfg(feature = "metrics")] + { + gauge!("dht_metadata_queue_depth").set(depth as f64); + gauge!("dht_metadata_in_flight").set(in_flight as f64); + } + } + + fn print_stats_inline(&self) { + #[cfg(debug_assertions)] + { + let received = self.total_received.load(Ordering::Relaxed); + let dropped = self.total_dropped.load(Ordering::Relaxed); + let dispatched = self.total_dispatched.load(Ordering::Relaxed); + let completed = self.total_completed.load(Ordering::Relaxed); + let queue_len = self.queue_len.load(Ordering::Relaxed); + let queue_pressure = queue_len as f64 / self.max_queue_size as f64 * 100.0; + log::info!( + "Metadata scheduler: depth={}/{}, pressure={:.1}%, received={}, dispatched={}, completed={}, evicted_or_expired={}", + queue_len, + self.max_queue_size, + queue_pressure, + received, + dispatched, + completed, + dropped, + ); + } + } +} + +async fn race_peer_fetches( + peers: Vec, + peer_rx: &mut mpsc::Receiver, + fetch: F, + runtime_stats: &DhtRuntimeStats, +) -> Option<(PeerCandidate, FetchedMetadata)> +where + F: Fn(PeerCandidate) -> Fut + Clone + Send + Sync + 'static, + Fut: Future + Send + 'static, +{ + runtime_stats.metadata_race_started(); + #[cfg(feature = "metrics")] + counter!("dht_metadata_peer_races_total").increment(1); + + let mut tasks = JoinSet::new(); + let mut candidates = 0usize; + for peer in peers { + let fetch = fetch.clone(); + tasks.spawn(fetch(peer)); + candidates += 1; + runtime_stats.metadata_peer_candidate(); + #[cfg(feature = "metrics")] + counter!("dht_metadata_peer_candidates_total", "source" => "initial").increment(1); + } + + let mut accepting_live_peers = candidates < MAX_METADATA_PEERS_PER_HASH; + loop { + if tasks.is_empty() { + if !accepting_live_peers { + return None; + } + match tokio::time::timeout(LIVE_PEER_IDLE_GRACE, peer_rx.recv()).await { + Ok(Some(peer)) => { + let fetch = fetch.clone(); + tasks.spawn(fetch(peer)); + candidates += 1; + runtime_stats.metadata_peer_candidate(); + #[cfg(feature = "metrics")] + counter!("dht_metadata_peer_candidates_total", "source" => "live").increment(1); + accepting_live_peers = candidates < MAX_METADATA_PEERS_PER_HASH; + } + Ok(None) | Err(_) => return None, + } + continue; + } + + tokio::select! { + result = tasks.join_next() => { + let Some(Ok((peer, outcome))) = result else { + continue; + }; + if let MetadataFetchOutcome::Fetched(metadata) = outcome { + let canceled = tasks.len(); + runtime_stats.metadata_peer_canceled(canceled); + #[cfg(feature = "metrics")] + counter!("dht_metadata_peer_canceled_total").increment(canceled as u64); + peer_rx.close(); + tasks.abort_all(); + while tasks.join_next().await.is_some() {} + return Some((peer, metadata)); + } + } + peer = peer_rx.recv(), if accepting_live_peers => { + match peer { + Some(peer) => { + let fetch = fetch.clone(); + tasks.spawn(fetch(peer)); + candidates += 1; + runtime_stats.metadata_peer_candidate(); + #[cfg(feature = "metrics")] + counter!("dht_metadata_peer_candidates_total", "source" => "live") + .increment(1); + accepting_live_peers = candidates < MAX_METADATA_PEERS_PER_HASH; + } + None => accepting_live_peers = false, + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::AtomicBool; + + struct DropSignal(Arc); + + impl Drop for DropSignal { + fn drop(&mut self) { + self.0.store(true, Ordering::Relaxed); + } + } + + fn empty_torrent_callback() -> Arc> { + Arc::new(ArcSwapOption::empty()) + } + + fn fetch_gate(callback: F) -> Arc> + where + F: Fn(String) -> Fut + Send + Sync + 'static, + Fut: std::future::Future + Send + 'static, + { + let holder = Arc::new(ArcSwapOption::empty()); + let callback: Arc = + Arc::new(Box::new(move |hash| Box::pin(callback(hash)))); + holder.store(Some(callback)); + holder + } + + fn completion_callback(callback: F) -> Arc> + where + F: Fn(MetadataFetchCompletion) + Send + Sync + 'static, + { + let holder = Arc::new(ArcSwapOption::empty()); + let callback: Arc = Arc::new(Box::new(callback)); + holder.store(Some(callback)); + holder + } + + fn event(hash: &str, port: u16, discovered_at: Instant) -> HashDiscovered { + HashDiscovered { + info_hash: hash.to_string(), + peer_addr: SocketAddr::from(([127, 0, 0, 1], port)), + discovered_at, + } + } + + #[test] + fn runtime_stats_track_queue_outcomes_and_depth() { + let (_hash_tx, hash_rx) = mpsc::channel(4); + let stats = DhtRuntimeStats::with_limits(DhtRuntimeLimits { + metadata_queue: 2, + ..DhtRuntimeLimits::default() + }); + let scheduler = MetadataScheduler::new_with_runtime_stats( + hash_rx, + Arc::new(RbitFetcher::new(1)), + MetadataSchedulerLimits { + queue_size: 2, + concurrency: 1, + }, + MetadataSchedulerCallbacks { + torrent: empty_torrent_callback(), + fetch_gate: Arc::new(ArcSwapOption::empty()), + completion: Arc::new(ArcSwapOption::empty()), + }, + Arc::new(AtomicUsize::new(0)), + CancellationToken::new(), + stats.clone(), + ); + let start = Instant::now(); + let mut queue = PendingHashQueue::new(2, HASH_QUEUE_TTL); + let mut in_flight = HashMap::new(); + + scheduler.enqueue(&mut queue, &mut in_flight, event("old", 1000, start)); + scheduler.enqueue( + &mut queue, + &mut in_flight, + event("old", 1001, start + Duration::from_secs(1)), + ); + scheduler.enqueue( + &mut queue, + &mut in_flight, + event("middle", 1002, start + Duration::from_secs(2)), + ); + scheduler.enqueue( + &mut queue, + &mut in_flight, + event("new", 1003, start + Duration::from_secs(3)), + ); + scheduler.enqueue( + &mut queue, + &mut in_flight, + event( + "stale", + 1004, + start + .checked_sub(HASH_QUEUE_TTL + Duration::from_secs(1)) + .unwrap(), + ), + ); + scheduler.sync_queue_len(queue.len(), 1); + + let snapshot = stats.snapshot(); + assert_eq!(snapshot.metadata_queue_inserted, 3); + assert_eq!(snapshot.metadata_queue_deduplicated, 1); + assert_eq!(snapshot.metadata_queue_evicted, 1); + assert_eq!(snapshot.metadata_queue_stale, 1); + assert_eq!(snapshot.metadata_queue_depth, 3); + assert_eq!(snapshot.metadata_queue_max, 2); + assert_eq!(snapshot.metadata_in_flight, 1); + } + + #[test] + fn queue_deduplicates_hash_and_keeps_twelve_newest_unique_peers() { + let start = Instant::now(); + let mut queue = PendingHashQueue::new(10, Duration::from_secs(60)); + for offset in 0..13 { + assert_ne!( + queue.push( + event( + "hash", + 1000 + offset, + start + Duration::from_secs(offset as u64) + ), + start + Duration::from_secs(offset as u64), + ), + QueuePushKind::Stale + ); + } + queue.push( + event("hash", 1012, start + Duration::from_secs(20)), + start + Duration::from_secs(20), + ); + + assert_eq!(queue.len(), 1); + let entry = queue.remove("hash").unwrap(); + let ports: Vec<_> = entry.peers.iter().map(|peer| peer.addr.port()).collect(); + assert_eq!(ports, (1001..=1012).rev().collect::>()); + } + + #[test] + fn coalescing_window_is_fixed_from_first_enqueue() { + let start = Instant::now(); + let mut queue = PendingHashQueue::new(10, Duration::from_secs(60)); + let in_flight = HashMap::new(); + queue.push(event("hash", 1000, start), start); + queue.push( + event("hash", 1001, start + Duration::from_millis(15)), + start + Duration::from_millis(15), + ); + + assert!( + queue + .pop_newest_ready(start + Duration::from_millis(24), &in_flight) + .is_none() + ); + let entry = queue + .pop_newest_ready(start + Duration::from_millis(25), &in_flight) + .expect("duplicate arrivals must not extend the coalescing deadline"); + assert_eq!(entry.peers.len(), 2); + } + + #[test] + fn in_flight_hash_routes_at_most_twelve_unique_peers_to_live_race() { + let (_hash_tx, hash_rx) = mpsc::channel(4); + let stats = DhtRuntimeStats::default(); + let scheduler = MetadataScheduler::new_with_runtime_stats( + hash_rx, + Arc::new(RbitFetcher::new(1)), + MetadataSchedulerLimits { + queue_size: 4, + concurrency: 1, + }, + MetadataSchedulerCallbacks { + torrent: empty_torrent_callback(), + fetch_gate: Arc::new(ArcSwapOption::empty()), + completion: Arc::new(ArcSwapOption::empty()), + }, + Arc::new(AtomicUsize::new(0)), + CancellationToken::new(), + stats.clone(), + ); + let (peer_tx, mut peer_rx) = mpsc::channel(MAX_METADATA_PEERS_PER_HASH); + let mut in_flight = HashMap::from([( + "hash".to_string(), + InFlightState { + peer_tx, + scheduled_peers: HashSet::from([SocketAddr::from(([127, 0, 0, 1], 1000))]), + }, + )]); + let mut queue = PendingHashQueue::new(4, HASH_QUEUE_TTL); + let now = Instant::now(); + + for port in 1001..=1013 { + scheduler.enqueue(&mut queue, &mut in_flight, event("hash", port, now)); + } + scheduler.enqueue(&mut queue, &mut in_flight, event("hash", 1001, now)); + + let mut received = Vec::new(); + while let Ok(peer) = peer_rx.try_recv() { + received.push(peer.addr.port()); + } + assert_eq!(received, (1001..=1011).collect::>()); + assert_eq!( + in_flight["hash"].scheduled_peers.len(), + MAX_METADATA_PEERS_PER_HASH + ); + assert!(queue.is_empty()); + let snapshot = stats.snapshot(); + assert_eq!(snapshot.metadata_live_peer_joins, 11); + assert_eq!(snapshot.metadata_queue_deduplicated, 14); + } + + #[tokio::test] + async fn peer_race_returns_first_success_and_cancels_remaining_fetches() { + let now = Instant::now(); + let peers = vec![ + PeerCandidate { + addr: SocketAddr::from(([127, 0, 0, 1], 1000)), + discovered_at: now, + }, + PeerCandidate { + addr: SocketAddr::from(([127, 0, 0, 1], 1001)), + discovered_at: now, + }, + PeerCandidate { + addr: SocketAddr::from(([127, 0, 0, 1], 1002)), + discovered_at: now, + }, + ]; + let started = Arc::new(AtomicUsize::new(0)); + let cancelled = Arc::new(AtomicBool::new(false)); + let barrier = Arc::new(tokio::sync::Barrier::new(peers.len())); + let started_for_race = started.clone(); + let cancelled_for_race = cancelled.clone(); + let (_peer_tx, mut peer_rx) = mpsc::channel(MAX_METADATA_PEERS_PER_HASH); + let stats = DhtRuntimeStats::default(); + + let result = race_peer_fetches( + peers, + &mut peer_rx, + move |peer| { + let started = started_for_race.clone(); + let cancelled = cancelled_for_race.clone(); + let barrier = barrier.clone(); + async move { + started.fetch_add(1, Ordering::Relaxed); + barrier.wait().await; + match peer.addr.port() { + 1000 => { + tokio::time::sleep(Duration::from_millis(10)).await; + ( + peer, + MetadataFetchOutcome::Fetched(( + "winner".to_string(), + 1, + Vec::new(), + 0, + )), + ) + } + 1001 => { + tokio::time::sleep(Duration::from_secs(1)).await; + (peer, MetadataFetchOutcome::Failed) + } + _ => { + let _signal = DropSignal(cancelled); + std::future::pending::<()>().await; + unreachable!() + } + } + } + }, + &stats, + ) + .await + .expect("one peer should win the race"); + + assert_eq!(result.0.addr.port(), 1000); + assert_eq!(result.1.0, "winner"); + assert!(cancelled.load(Ordering::Relaxed)); + let snapshot = stats.snapshot(); + assert_eq!(snapshot.metadata_races_started, 1); + assert_eq!(snapshot.metadata_peer_candidates, 3); + assert_eq!(snapshot.metadata_peer_canceled, 2); + } + + #[tokio::test] + async fn live_peer_joins_running_race_and_cancels_slow_initial_peer() { + let now = Instant::now(); + let initial = vec![PeerCandidate { + addr: SocketAddr::from(([127, 0, 0, 1], 1000)), + discovered_at: now, + }]; + let live = PeerCandidate { + addr: SocketAddr::from(([127, 0, 0, 1], 1001)), + discovered_at: now, + }; + let (peer_tx, mut peer_rx) = mpsc::channel(MAX_METADATA_PEERS_PER_HASH); + let initial_started = Arc::new(tokio::sync::Notify::new()); + let initial_started_for_fetch = initial_started.clone(); + let initial_cancelled = Arc::new(AtomicBool::new(false)); + let initial_cancelled_for_fetch = initial_cancelled.clone(); + let stats = DhtRuntimeStats::default(); + let sender = tokio::spawn(async move { + initial_started.notified().await; + peer_tx.send(live).await.unwrap(); + }); + + let result = tokio::time::timeout( + Duration::from_secs(1), + race_peer_fetches( + initial, + &mut peer_rx, + move |peer| { + let initial_started = initial_started_for_fetch.clone(); + let initial_cancelled = initial_cancelled_for_fetch.clone(); + async move { + if peer.addr.port() == 1000 { + let _signal = DropSignal(initial_cancelled); + initial_started.notify_one(); + std::future::pending::<()>().await; + unreachable!() + } + ( + peer, + MetadataFetchOutcome::Fetched(( + "live-winner".to_string(), + 1, + Vec::new(), + 0, + )), + ) + } + }, + &stats, + ), + ) + .await + .unwrap() + .expect("live peer should win"); + sender.await.unwrap(); + + assert_eq!(result.0.addr.port(), 1001); + assert!(initial_cancelled.load(Ordering::Relaxed)); + let snapshot = stats.snapshot(); + assert_eq!(snapshot.metadata_peer_candidates, 2); + assert_eq!(snapshot.metadata_peer_canceled, 1); + } + + #[test] + fn full_queue_evicts_oldest_for_newer_hash() { + let start = Instant::now(); + let mut queue = PendingHashQueue::new(2, Duration::from_secs(60)); + queue.push(event("old", 1000, start), start); + queue.push( + event("middle", 1001, start + Duration::from_secs(1)), + start + Duration::from_secs(1), + ); + let result = queue.push( + event("new", 1002, start + Duration::from_secs(2)), + start + Duration::from_secs(2), + ); + + assert_eq!(result, QueuePushKind::EvictedOldest); + assert!(!queue.contains("old")); + assert!(queue.contains("middle")); + assert!(queue.contains("new")); + } + + #[test] + fn queue_expires_stale_hashes_and_pops_newest_first() { + let start = Instant::now(); + let mut queue = PendingHashQueue::new(4, Duration::from_secs(10)); + queue.push(event("older", 1000, start), start); + queue.push( + event("newer", 1001, start + Duration::from_secs(1)), + start + Duration::from_secs(1), + ); + + let in_flight = HashMap::new(); + assert!( + queue + .pop_newest_ready(start + Duration::from_millis(24), &in_flight) + .is_none() + ); + let newest = queue + .pop_newest_ready(start + Duration::from_secs(2), &in_flight) + .unwrap(); + assert_eq!(newest.info_hash, "newer"); + assert_eq!(queue.expire(start + Duration::from_secs(11)), 1); + assert!(queue.is_empty()); + } + + #[tokio::test] + async fn gate_rejection_does_not_emit_completion() { + let (hash_tx, hash_rx) = mpsc::channel(4); + let gate_calls = Arc::new(AtomicUsize::new(0)); + let gate_calls_for_callback = gate_calls.clone(); + let gate = fetch_gate(move |_| { + gate_calls_for_callback.fetch_add(1, Ordering::Relaxed); + async { false } + }); + let (completion_tx, mut completion_rx) = mpsc::unbounded_channel(); + let completion = completion_callback(move |result| { + let _ = completion_tx.send(result); + }); + let shutdown = CancellationToken::new(); + let scheduler = MetadataScheduler::new( + hash_rx, + Arc::new(RbitFetcher::new(1)), + MetadataSchedulerLimits { + queue_size: 4, + concurrency: 1, + }, + MetadataSchedulerCallbacks { + torrent: empty_torrent_callback(), + fetch_gate: gate, + completion, + }, + Arc::new(AtomicUsize::new(0)), + shutdown.clone(), + ); + let scheduler_task = tokio::spawn(scheduler.run()); + hash_tx + .send(event("bad", 1000, Instant::now())) + .await + .unwrap(); + + tokio::time::timeout(Duration::from_secs(1), async { + while gate_calls.load(Ordering::Relaxed) == 0 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + assert!( + tokio::time::timeout(Duration::from_millis(50), completion_rx.recv()) + .await + .is_err() + ); + + shutdown.cancel(); + scheduler_task.await.unwrap(); + } + + #[tokio::test] + async fn admitted_failure_emits_one_completion() { + let (hash_tx, hash_rx) = mpsc::channel(4); + let gate = fetch_gate(|_| async { true }); + let (completion_tx, mut completion_rx) = mpsc::unbounded_channel(); + let completion = completion_callback(move |result| { + let _ = completion_tx.send(result); + }); + let shutdown = CancellationToken::new(); + let scheduler = MetadataScheduler::new( + hash_rx, + Arc::new(RbitFetcher::new(1)), + MetadataSchedulerLimits { + queue_size: 4, + concurrency: 1, + }, + MetadataSchedulerCallbacks { + torrent: empty_torrent_callback(), + fetch_gate: gate, + completion, + }, + Arc::new(AtomicUsize::new(0)), + shutdown.clone(), + ); + let scheduler_task = tokio::spawn(scheduler.run()); + hash_tx + .send(event("bad", 1000, Instant::now())) + .await + .unwrap(); + + let result = tokio::time::timeout(Duration::from_secs(1), completion_rx.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!(result.info_hash, "bad"); + assert_eq!(result.status, MetadataFetchCompletionStatus::FetchFailed); + assert_eq!(result.attempts, 0); + assert!( + tokio::time::timeout(Duration::from_millis(50), completion_rx.recv()) + .await + .is_err() + ); + + shutdown.cancel(); + scheduler_task.await.unwrap(); + } +} diff --git a/dht-crawler/src/server.rs b/dht-crawler/src/server.rs new file mode 100644 index 0000000..4ed22ac --- /dev/null +++ b/dht-crawler/src/server.rs @@ -0,0 +1,1060 @@ +// 负责组合 DHT 网络组件回调生命周期和公开运行接口 + +use crate::addr::is_valid_node_addr; +use crate::budget::{RateBucket, SharedRateBudget}; +use crate::crawl_config::ResolvedCrawlConfig; +use crate::crawl_engine::CrawlEngine; +use crate::error::Result; +use crate::krpc::encode_response; +use crate::metadata::RbitFetcher; +use crate::node_id::{neighbor_node_id, random_node_id, transaction_id_from_bytes}; +use crate::peer_lookup::{ + PeerLookupHandle, PeerLookupRuntime, is_peer_lookup_tid, spawn_peer_lookup, +}; +use crate::protocol::{DhtArgs, DhtMessage}; +use crate::runtime_stats::{DhtRuntimeLimits, DhtRuntimeStats}; +use crate::sample_infohashes::{ + SampleHashAdmissionCallback, SampleInfohashesHandle, SampleInfohashesRuntime, + is_sample_infohashes_tid, spawn_sample_infohashes, +}; +use crate::scheduler::{ + MetadataCompletionCallback, MetadataFetchCallback, MetadataScheduler, + MetadataSchedulerCallbacks, MetadataSchedulerLimits, MetadataSchedulerRuntime, + TorrentAckCallback, +}; +use crate::types::{DHTOptions, MetadataFetchCompletion, NetMode, NodeTuple, TorrentInfo}; +use crate::udp_buffer::UdpBufferPool; +use crate::udp_ingress::{WorkerHandle, spawn_udp_listener}; +use ahash::AHashMap; +use arc_swap::ArcSwapOption; +use bytes::BytesMut; +#[cfg(feature = "metrics")] +use metrics::counter; +use rand::RngExt; +use socket2::{Domain, Protocol, Socket, Type}; +use std::collections::{HashMap, VecDeque}; +use std::future::Future; +use std::hash::{Hash, Hasher}; +use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::{Duration, Instant}; +use tokio::net::UdpSocket; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; + +type FilterCallback = Box bool + Send + Sync + 'static>; +type ErrorCallback = Box; + +struct QueryResponse<'a> { + transaction_id: &'a [u8], + remote_addr: SocketAddr, + local_addr: SocketAddr, + query_type: &'a str, + sender_id: Option<&'a [u8]>, + target_id: Option<&'a [u8]>, +} + +#[derive(Clone, Copy)] +struct SourceResponseWindow { + started_at: Instant, + last_seen: Instant, + count: u32, +} + +struct WorkerResponseLimiter { + regular_packets: RateBucket, + regular_bytes: RateBucket, + priority_packets: RateBucket, + priority_bytes: RateBucket, + per_source_rate: u32, + sources: AHashMap, + source_expiry: VecDeque<(Instant, SocketAddr)>, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ResponsePermit { + Regular, + PriorityReserve, + Rejected, +} + +fn reserve_quota(total: u32) -> u32 { + if total == 0 { + 0 + } else { + total.div_ceil(10).min(total) + } +} + +impl WorkerResponseLimiter { + fn new(packet_rate: u32, byte_rate: u64, per_source_rate: u32, now: Instant) -> Self { + let byte_rate = byte_rate.min(u32::MAX as u64) as u32; + let priority_packet_rate = reserve_quota(packet_rate); + let priority_byte_rate = reserve_quota(byte_rate); + let packet_rate = packet_rate.saturating_sub(priority_packet_rate); + let byte_rate = byte_rate.saturating_sub(priority_byte_rate); + Self { + regular_packets: RateBucket::per_second( + packet_rate, + packet_rate.div_ceil(5).max(1), + true, + now, + ), + regular_bytes: RateBucket::per_second( + byte_rate, + byte_rate.div_ceil(5).max(512), + true, + now, + ), + priority_packets: RateBucket::per_second( + priority_packet_rate, + priority_packet_rate.div_ceil(5).max(1), + true, + now, + ), + priority_bytes: RateBucket::per_second( + priority_byte_rate, + priority_byte_rate.div_ceil(5).max(512), + true, + now, + ), + per_source_rate, + sources: AHashMap::new(), + source_expiry: VecDeque::new(), + } + } + + fn acquire( + &mut self, + addr: SocketAddr, + encoded_len: usize, + is_priority: bool, + now: Instant, + ) -> ResponsePermit { + self.expire_sources(now); + if Self::take_budget( + &mut self.regular_packets, + &mut self.regular_bytes, + encoded_len, + now, + ) { + if self.acquire_source_slot(addr, now) { + return ResponsePermit::Regular; + } + Self::refund_budget( + &mut self.regular_packets, + &mut self.regular_bytes, + encoded_len, + ); + return ResponsePermit::Rejected; + } + + if is_priority + && Self::take_budget( + &mut self.priority_packets, + &mut self.priority_bytes, + encoded_len, + now, + ) + { + if self.acquire_source_slot(addr, now) { + return ResponsePermit::PriorityReserve; + } + Self::refund_budget( + &mut self.priority_packets, + &mut self.priority_bytes, + encoded_len, + ); + } + + ResponsePermit::Rejected + } + + fn take_budget( + packets: &mut RateBucket, + bytes: &mut RateBucket, + encoded_len: usize, + now: Instant, + ) -> bool { + if !packets.try_take_one(now) { + return false; + } + if !bytes.try_take_exact(encoded_len, now) { + packets.refund_one(); + return false; + } + true + } + + fn refund_budget(packets: &mut RateBucket, bytes: &mut RateBucket, encoded_len: usize) { + packets.refund_one(); + bytes.refund(encoded_len); + } + + fn acquire_source_slot(&mut self, addr: SocketAddr, now: Instant) -> bool { + if self.per_source_rate == 0 { + return false; + } + + let entry = self.sources.entry(addr).or_insert(SourceResponseWindow { + started_at: now, + last_seen: now, + count: 0, + }); + if now + .checked_duration_since(entry.started_at) + .unwrap_or_default() + >= Duration::from_secs(1) + { + entry.started_at = now; + entry.count = 0; + } + if entry.count >= self.per_source_rate { + return false; + } + entry.count += 1; + entry.last_seen = now; + self.source_expiry + .push_back((now + Duration::from_secs(60), addr)); + true + } + + fn expire_sources(&mut self, now: Instant) { + while let Some((deadline, addr)) = self.source_expiry.front().copied() { + if deadline > now { + break; + } + self.source_expiry.pop_front(); + if self.sources.get(&addr).is_some_and(|entry| { + now.checked_duration_since(entry.last_seen) + .unwrap_or_default() + >= Duration::from_secs(60) + }) { + self.sources.remove(&addr); + } + } + } +} + +#[derive(Debug, Clone)] +/// InfoHash and announcing Peer submitted to the Metadata scheduler. +pub struct HashDiscovered { + /// Lowercase hexadecimal InfoHash. + pub info_hash: String, + /// Peer endpoint derived from announce `port`/`implied_port`. + pub peer_addr: SocketAddr, + /// Monotonic discovery time used for freshness and queue ordering. + pub discovered_at: std::time::Instant, +} + +#[derive(Clone)] +/// Cloneable BEP-5 server handle and primary crate entry point. +pub struct DHTServer { + options: DHTOptions, + crawl_config: ResolvedCrawlConfig, + node_id: [u8; 20], + sockets_by_bind_addr: Arc>>, + token_secret: [u8; 10], + torrent_callback: Arc>, + hash_filter: Arc>, + on_metadata_fetch: Arc>, + sample_hash_admission: Arc>, + metadata_completion_callback: Arc>, + error_callback: Arc>, + crawl_engine: Arc, + peer_lookup: PeerLookupHandle, + metadata_fetcher: Arc, + sample_infohashes: SampleInfohashesHandle, + hash_events_tx: mpsc::Sender, + metadata_queue_len: Arc, + max_metadata_queue_size: usize, + runtime_stats: DhtRuntimeStats, + shutdown: CancellationToken, +} + +fn create_udp_sock(domain: Domain, ty: Type, addr: SocketAddr) -> std::io::Result { + let sock = Socket::new(domain, ty, Some(Protocol::UDP))?; + #[cfg(not(windows))] + { + sock.set_reuse_port(true)?; + if addr.is_ipv6() { + sock.set_only_v6(true)?; + } + } + let _ = sock.set_reuse_address(true); + sock.set_nonblocking(true)?; + let _ = sock.set_recv_buffer_size(32 * 1024 * 1024); + let _ = sock.set_send_buffer_size(8 * 1024 * 1024); + sock.bind(&addr.into())?; + UdpSocket::from_std(sock.into()) +} + +fn split_u32_quota(total: u32, workers: usize, worker: usize) -> u32 { + let workers = workers.max(1) as u32; + total / workers + u32::from((worker as u32) < total % workers) +} + +fn split_u64_quota(total: u64, workers: usize, worker: usize) -> u64 { + let workers = workers.max(1) as u64; + total / workers + u64::from((worker as u64) < total % workers) +} + +impl DHTServer { + /// Validates options, binds configured UDP sockets and constructs bounded pipelines. + /// + /// Background Metadata scheduling begins during construction. Active crawling and UDP receive + /// loops begin when [`Self::start`] is awaited. + pub async fn new(options: DHTOptions) -> Result { + let crawl_config = ResolvedCrawlConfig::from_options(&options.crawl); + let runtime_stats = DhtRuntimeStats::with_limits(DhtRuntimeLimits { + metadata_queue: options.metadata.max_queue_size.max(1), + node_pool: crawl_config.pool_capacity, + node_pool_low_watermark: crawl_config.low_watermark, + find_node_in_flight: crawl_config.max_in_flight, + initial_find_node_rate: crawl_config.max_find_node_rate_per_sec, + hash_ingress_queue: options.hash_queue_capacity, + crawl_priority_queue: crawl_config.priority_event_channel_capacity, + crawl_discovery_queue: crawl_config.discovery_event_channel_capacity, + }); + const ANY_V4_ADDR: SocketAddr = + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 8080); + const ANY_V6_ADDR: SocketAddr = + SocketAddr::new(IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 0)), 8080); + + let mut sockets_by_bind_addr = HashMap::new(); + match options.netmode { + NetMode::Ipv4Only => { + let mut addr = ANY_V4_ADDR; + addr.set_port(options.port); + let sock = create_udp_sock(Domain::IPV4, Type::DGRAM, addr)?; + sockets_by_bind_addr.insert(addr, Arc::new(sock)); + } + NetMode::Ipv6Only => { + let mut addr = ANY_V6_ADDR; + addr.set_port(options.port); + let sock = create_udp_sock(Domain::IPV6, Type::DGRAM, addr)?; + sockets_by_bind_addr.insert(addr, Arc::new(sock)); + } + NetMode::DualStack => { + let mut addr_v4 = ANY_V4_ADDR; + addr_v4.set_port(options.port); + let sock_v4 = create_udp_sock(Domain::IPV4, Type::DGRAM, addr_v4)?; + sockets_by_bind_addr.insert(addr_v4, Arc::new(sock_v4)); + + let mut addr_v6 = ANY_V6_ADDR; + addr_v6.set_port(options.port); + let sock_v6 = create_udp_sock(Domain::IPV6, Type::DGRAM, addr_v6)?; + sockets_by_bind_addr.insert(addr_v6, Arc::new(sock_v6)); + } + } + + let node_id = random_node_id(); + let mut token_secret = [0u8; 10]; + rand::rng().fill(&mut token_secret); + + let (hash_events_tx, hash_rx) = + 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, + node_id, + &sockets_by_bind_addr, + crawl_engine.snapshot.clone(), + hash_events_tx.clone(), + PeerLookupRuntime { + options: options.peer_lookup.clone(), + stats: runtime_stats.clone(), + outbound_query_budget: outbound_query_budget.clone(), + shutdown: shutdown.clone(), + }, + ); + let sample_infohashes = spawn_sample_infohashes( + options.netmode, + node_id, + &sockets_by_bind_addr, + crawl_engine.snapshot.clone(), + crawl_engine.clone(), + peer_lookup.clone(), + SampleInfohashesRuntime { + options: options.sample_infohashes.clone(), + stats: runtime_stats.clone(), + outbound_query_budget, + hash_admission: sample_hash_admission.clone(), + shutdown: shutdown.clone(), + }, + ); + + let scheduler = MetadataScheduler::new_with_runtime( + hash_rx, + fetcher.clone(), + MetadataSchedulerLimits { + queue_size: options.metadata.max_queue_size, + concurrency: options.metadata.max_worker_count, + }, + MetadataSchedulerCallbacks { + torrent: torrent_callback.clone(), + fetch_gate: on_metadata_fetch.clone(), + completion: metadata_completion_callback.clone(), + }, + metadata_queue_len.clone(), + shutdown.clone(), + MetadataSchedulerRuntime { + stats: runtime_stats.clone(), + peer_lookup_tx: Some(peer_lookup.request_sender()), + }, + ); + tokio::spawn(scheduler.run()); + + let max_metadata_queue_size = options.metadata.max_queue_size; + Ok(Self { + options, + crawl_config: crawl_config.clone(), + node_id, + sockets_by_bind_addr: Arc::new(sockets_by_bind_addr), + token_secret, + 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, + peer_lookup, + metadata_fetcher: fetcher, + sample_infohashes, + hash_events_tx, + metadata_queue_len, + max_metadata_queue_size, + runtime_stats, + shutdown, + }) + } + + /// Registers the asynchronous admission gate invoked before the first real Peer attempt. + /// + /// Returning `false` rejects the job without downloading Metadata and without emitting a + /// [`MetadataFetchCompletion`]. If no gate is registered, jobs are admitted. + pub fn on_metadata_fetch(&self, callback: F) + where + F: Fn(String) -> Fut + Send + Sync + 'static, + Fut: Future + Send + 'static, + { + let callback: Arc = + Arc::new(Box::new(move |hash| Box::pin(callback(hash)))); + 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. + pub fn on_torrent(&self, callback: F) + where + F: Fn(TorrentInfo) + Send + Sync + 'static, + { + let callback: Arc = Arc::new(Box::new(move |torrent| { + callback(torrent); + true + })); + self.torrent_callback.store(Some(callback)); + } + + /// Registers a torrent callback that acknowledges application delivery. + /// + /// Returning `true` produces [`crate::MetadataFetchCompletionStatus::Accepted`]. Returning + /// `false` produces [`crate::MetadataFetchCompletionStatus::DeliveryRejected`]: the Metadata + /// download was + /// valid, but the application did not accept it. A callback panic is caught and treated as + /// rejected delivery. Registering this callback replaces any previous torrent callback. + pub fn on_torrent_with_ack(&self, callback: F) + where + F: Fn(TorrentInfo) -> bool + Send + Sync + 'static, + { + let callback: Arc = Arc::new(Box::new(callback)); + self.torrent_callback.store(Some(callback)); + } + + /// Registers a callback invoked exactly once when an admitted Metadata job terminates. + /// + /// Gate rejection does not emit a completion. `attempts` counts real Peer network attempts; + /// failure-cache skips are excluded. Registering a new callback replaces the previous one. + pub fn on_metadata_fetch_complete(&self, callback: F) + where + F: Fn(MetadataFetchCompletion) + Send + Sync + 'static, + { + let callback: Arc = Arc::new(Box::new(callback)); + self.metadata_completion_callback.store(Some(callback)); + } + + /// Registers an early synchronous InfoHash filter for valid `announce_peer` queries. + /// + /// Returning `false` prevents the Hash from entering the bounded ingress queue. Registering a + /// new filter replaces the previous one. + pub fn filter(&self, filter: F) + where + F: Fn(&str) -> bool + Send + Sync + 'static, + { + let filter: Arc = Arc::new(Box::new(filter)); + self.hash_filter.store(Some(filter)); + } + + /// Alias for [`Self::filter`]. + pub fn set_filter(&self, filter: F) + where + F: Fn(&str) -> bool + Send + Sync + 'static, + { + self.filter(filter); + } + + /// Registers a runtime error callback, replacing any previous callback. + /// + /// Initialization and `start` errors are still returned through [`Result`]. + pub fn on_error(&self, callback: F) + where + F: Fn(crate::error::DHTError) + Send + Sync + 'static, + { + let callback: Arc = Arc::new(Box::new(callback)); + self.error_callback.store(Some(callback)); + } + + fn emit_error(&self, error: crate::error::DHTError) { + if let Some(callback) = self.error_callback.load_full() { + callback(error); + } + } + + /// Returns the current strict FIFO crawl-pool size. + pub fn get_node_pool_size(&self) -> usize { + self.crawl_engine.node_count.load(Ordering::Relaxed) + } + + /// Returns a cheap cloneable handle to transport-neutral atomic runtime statistics. + pub fn runtime_stats(&self) -> DhtRuntimeStats { + self.runtime_stats.clone() + } + + /// Performs one bounded active `get_peers` lookup and returns its unique Peer addresses. + pub async fn lookup_peers(&self, info_hash: [u8; 20]) -> Result { + tokio::time::timeout(Duration::from_secs(10), self.peer_lookup.lookup(info_hash)) + .await + .map_err(|_| crate::error::DHTError::Other("Peer Lookup 等待超时".to_owned()))? + .map_err(|message| crate::error::DHTError::Other(message.to_owned())) + } + + /// Verifies that a Peer completes the BitTorrent handshake for the requested InfoHash. + pub async fn verify_peer_handshake(&self, info_hash: [u8; 20], peer: SocketAddr) -> bool { + self.metadata_fetcher + .verify_handshake(&info_hash, peer) + .await + } + + /// Starts crawl and UDP background tasks, then waits until [`Self::shutdown`] is called. + /// + /// A shut-down server cannot be restarted. Construct a new [`DHTServer`] for another run. + pub async fn start(&self) -> Result<()> { + if self.shutdown.is_cancelled() { + return Err(crate::error::DHTError::Other( + "DHT server is already shut down".to_string(), + )); + } + + self.crawl_engine.spawn( + self.options.netmode, + self.node_id, + &self.sockets_by_bind_addr, + self.metadata_queue_len.clone(), + self.max_metadata_queue_size, + self.shutdown.clone(), + ); + + let buffer_pool = UdpBufferPool::new(); + let workers = self.spawn_workers(buffer_pool.clone()); + for sock in self.sockets_by_bind_addr.values().cloned() { + spawn_udp_listener( + sock, + workers.clone(), + self.shutdown.clone(), + buffer_pool.clone(), + self.runtime_stats.clone(), + )?; + } + + self.shutdown.cancelled().await; + Ok(()) + } + + /// Cancels UDP, crawl and Metadata tasks. This method is safe to call more than once. + pub fn shutdown(&self) { + self.shutdown.cancel(); + } + + fn spawn_workers(&self, buffer_pool: UdpBufferPool) -> Vec { + let server = self.clone(); + let shutdown = self.shutdown.clone(); + let num_workers = std::thread::available_parallelism() + .map(|n| n.get()) + .unwrap_or(8); + let queue_size = 5_000; + + let mut workers: Vec = Vec::with_capacity(num_workers); + for worker_id in 0..num_workers { + let (tx, mut rx) = mpsc::channel(queue_size); + workers.push(tx); + + let server_clone = server.clone(); + let cancellation_token = shutdown.clone(); + let pool = buffer_pool.clone(); + let packet_rate = split_u32_quota( + self.crawl_config.max_response_rate_per_sec, + num_workers, + worker_id, + ); + let byte_rate = split_u64_quota( + self.crawl_config.max_response_bytes_per_sec, + num_workers, + worker_id, + ); + let per_source_rate = self.crawl_config.max_response_rate_per_source; + + tokio::spawn(async move { + let mut response_limiter = WorkerResponseLimiter::new( + packet_rate, + byte_rate, + per_source_rate, + Instant::now(), + ); + let mut response_buffer = BytesMut::with_capacity(512); + loop { + tokio::select! { + _ = cancellation_token.cancelled() => break, + msg = rx.recv() => { + match msg { + Some((packet, remote_addr, local_addr)) => { + if let Err(e) = server_clone + .handle_message( + packet.payload(), + remote_addr, + local_addr, + &mut response_limiter, + &mut response_buffer, + ) + .await + { + server_clone.emit_error(e); + } + pool.release(packet.buf); + } + None => break, + } + } + } + } + }); + } + workers + } + + async fn handle_message( + &self, + data: &[u8], + remote_addr: SocketAddr, + local_addr: SocketAddr, + response_limiter: &mut WorkerResponseLimiter, + response_buffer: &mut BytesMut, + ) -> Result<()> { + if self.sockets_by_bind_addr.get(&local_addr).is_none() { + return Ok(()); + } + + let msg: DhtMessage = match serde_bencode::from_bytes(data) { + Ok(m) => m, + Err(_) => { + #[cfg(feature = "metrics")] + counter!("dht_messages_parse_error_total").increment(1); + return Ok(()); + } + }; + + #[cfg(feature = "metrics")] + { + let label = match msg.y.as_str() { + "q" => "q", + "r" => "r", + "e" => "e", + _ => "unknown", + }; + counter!("dht_messages_processed_total", "type" => label).increment(1); + } + + match msg.y.as_str() { + "q" => { + if let Some(q_type) = &msg.q { + self.handle_query( + &msg, + q_type.as_bytes(), + remote_addr, + local_addr, + response_limiter, + response_buffer, + ) + .await?; + } + } + "r" => { + if let Some(response) = msg.r + && let Some(tid) = transaction_id_from_bytes(&msg.t) + { + if is_sample_infohashes_tid(&tid) { + self.sample_infohashes + .route_response(remote_addr, tid, response); + } else if is_peer_lookup_tid(&tid) { + self.peer_lookup.route_response(remote_addr, tid, response); + } else { + self.crawl_engine.route_response(remote_addr, tid, response); + } + } + } + _ => {} + } + Ok(()) + } + + async fn handle_query( + &self, + msg: &DhtMessage, + query_type: &[u8], + remote_addr: SocketAddr, + local_addr: SocketAddr, + response_limiter: &mut WorkerResponseLimiter, + response_buffer: &mut BytesMut, + ) -> Result<()> { + let args = match &msg.a { + Some(a) => a, + None => return Ok(()), + }; + + let transaction_id = &msg.t; + let sender_id: Option<&[u8]> = args.id.as_deref().map(|v| v.as_slice()); + let target_id_fallback: Option<&[u8]> = args + .target + .as_deref() + .or(args.info_hash.as_deref()) + .map(|v| v.as_slice()); + let q_str = std::str::from_utf8(query_type).unwrap_or(""); + self.runtime_stats.inbound_query(q_str); + + if let Some(sender_id) = sender_id + && sender_id.len() == 20 + && is_valid_node_addr(&remote_addr) + { + self.crawl_engine.route_discovered(NodeTuple { + id: sender_id + .try_into() + .expect("validated DHT sender id contains 20 bytes"), + addr: remote_addr, + }); + } + + #[cfg(feature = "metrics")] + { + let label = match q_str { + "ping" => "ping", + "find_node" => "find_node", + "get_peers" => "get_peers", + "announce_peer" => "announce_peer", + "vote" => "vote", + _ => "other_or_invalid", + }; + counter!("dht_queries_total", "q" => label).increment(1); + } + + if q_str == "announce_peer" { + self.handle_announce_peer(args, remote_addr).await?; + } + + self.send_response( + QueryResponse { + transaction_id, + remote_addr, + local_addr, + query_type: q_str, + sender_id, + target_id: target_id_fallback, + }, + response_limiter, + response_buffer, + ) + .await?; + Ok(()) + } + + async fn handle_announce_peer(&self, args: &DhtArgs, addr: SocketAddr) -> Result<()> { + if let Some(token) = &args.token { + if !self.validate_token(token, addr) { + self.runtime_stats.announce_invalid_token(); + #[cfg(feature = "metrics")] + counter!("dht_announce_peer_blocked_total", "reason" => "invalid_token") + .increment(1); + return Ok(()); + } + } else { + self.runtime_stats.announce_invalid_token(); + return Ok(()); + } + + if let Some(info_hash) = &args.info_hash { + let info_hash_arr: [u8; 20] = match info_hash.as_ref().try_into() { + Ok(arr) => arr, + Err(_) => { + self.runtime_stats.announce_invalid_token(); + return Ok(()); + } + }; + let hash_hex = hex::encode(info_hash_arr); + + if let Some(filter) = self.hash_filter.load_full() + && !filter(&hash_hex) + { + self.runtime_stats.announce_filtered(); + #[cfg(feature = "metrics")] + counter!("dht_announce_peer_blocked_total", "reason" => "filtered").increment(1); + return Ok(()); + } + + #[cfg(feature = "metrics")] + counter!("dht_info_hashes_discovered_total").increment(1); + + let port = if let Some(implied) = args.implied_port { + if implied != 0 { + addr.port() + } else { + args.port.unwrap_or(0) + } + } else { + args.port.unwrap_or(addr.port()) + }; + + if port > 0 { + self.runtime_stats.announce_accepted(); + self.runtime_stats.hash_received(); + let event = HashDiscovered { + info_hash: hash_hex, + peer_addr: SocketAddr::new(addr.ip(), port), + discovered_at: std::time::Instant::now(), + }; + + let enqueue_result = self.hash_events_tx.try_send(event); + self.runtime_stats.set_hash_ingress_queue_depth( + self.hash_events_tx + .max_capacity() + .saturating_sub(self.hash_events_tx.capacity()), + ); + if enqueue_result.is_err() { + self.runtime_stats.hash_ingress_dropped(); + #[cfg(feature = "metrics")] + counter!("dht_metadata_ingress_dropped_total", "reason" => "queue_full") + .increment(1); + #[cfg(debug_assertions)] + log::debug!("Hash queue is full; dropping hash"); + } + } + } + Ok(()) + } + + async fn send_response( + &self, + response: QueryResponse<'_>, + response_limiter: &mut WorkerResponseLimiter, + response_buffer: &mut BytesMut, + ) -> Result<()> { + let socket = match self.sockets_by_bind_addr.get(&response.local_addr) { + Some(sock) => sock, + None => return Ok(()), + }; + + let reference_id = response.sender_id.or(response.target_id); + let my_id = if let Some(target) = reference_id { + let generated = neighbor_node_id(target, &self.node_id); + <[u8; 20]>::try_from(generated.as_slice()).expect("neighbor id is always 20 bytes") + } else { + self.node_id + }; + let token = self.generate_token(response.remote_addr); + let include_nodes = + response.query_type == "get_peers" || response.query_type == "find_node"; + let requestor_is_ipv6 = response.remote_addr.is_ipv6(); + let nodes = if include_nodes { + let filter_ipv6 = match self.options.netmode { + NetMode::Ipv4Only => Some(false), + NetMode::Ipv6Only => Some(true), + NetMode::DualStack => Some(requestor_is_ipv6), + }; + let snapshot = self.crawl_engine.snapshot.load(); + snapshot.random_nodes(8, filter_ipv6) + } else { + Vec::new() + }; + encode_response( + response_buffer, + response.transaction_id, + &my_id, + &token, + &nodes, + requestor_is_ipv6, + ); + let is_priority = response.query_type == "ping" || response.query_type == "get_peers"; + match response_limiter.acquire( + response.remote_addr, + response_buffer.len(), + is_priority, + Instant::now(), + ) { + ResponsePermit::Regular => self.runtime_stats.response_normal(), + ResponsePermit::PriorityReserve => { + self.runtime_stats.udp_response_priority_reserved(); + #[cfg(feature = "metrics")] + counter!( + "dht_udp_responses_priority_reserved_total", + "query" => if response.query_type == "ping" { "ping" } else { "get_peers" } + ) + .increment(1); + } + ResponsePermit::Rejected => { + self.runtime_stats.udp_response_rate_limited(); + #[cfg(feature = "metrics")] + counter!("dht_udp_responses_dropped_total", "reason" => "rate_limit").increment(1); + return Ok(()); + } + } + match socket.send_to(response_buffer, response.remote_addr).await { + Ok(len) => { + self.runtime_stats.udp_sent(len); + #[cfg(feature = "metrics")] + { + counter!("dht_udp_bytes_sent_total").increment(len as u64); + counter!("dht_udp_packets_sent_total", "type" => "response").increment(1); + } + } + Err(_) => self.runtime_stats.response_send_failed(), + } + Ok(()) + } + + fn generate_token(&self, addr: SocketAddr) -> [u8; 8] { + let mut hasher = ahash::AHasher::default(); + match addr.ip() { + IpAddr::V4(ip) => ip.octets().hash(&mut hasher), + IpAddr::V6(ip) => ip.octets().hash(&mut hasher), + } + self.token_secret.hash(&mut hasher); + hasher.finish().to_le_bytes() + } + + fn validate_token(&self, token: &[u8], addr: SocketAddr) -> bool { + if token.len() != 8 { + return false; + } + let expected = self.generate_token(addr); + token == expected + } +} + +#[cfg(test)] +mod response_limiter_tests { + use super::*; + + #[test] + fn response_limiter_enforces_packet_byte_and_source_limits() { + let start = Instant::now(); + let addr: SocketAddr = "8.8.8.8:6881".parse().unwrap(); + let mut limiter = WorkerResponseLimiter::new(2, 200, 1, start); + assert_eq!( + limiter.acquire(addr, 100, false, start), + ResponsePermit::Regular + ); + assert_eq!( + limiter.acquire(addr, 100, false, start), + ResponsePermit::Rejected + ); + assert_eq!( + limiter.acquire(addr, 100, false, start + Duration::from_secs(1)), + ResponsePermit::Regular + ); + assert_eq!( + limiter.acquire( + "1.1.1.1:6881".parse().unwrap(), + 513, + false, + start + Duration::from_secs(2) + ), + ResponsePermit::Rejected + ); + } + + #[test] + fn ping_and_get_peers_can_use_the_priority_reserve() { + let start = Instant::now(); + let mut limiter = WorkerResponseLimiter::new(10, 1_000, 100, start); + let first: SocketAddr = "8.8.8.8:6881".parse().unwrap(); + let second: SocketAddr = "1.1.1.1:6881".parse().unwrap(); + let priority: SocketAddr = "9.9.9.9:6881".parse().unwrap(); + + assert_eq!( + limiter.acquire(first, 50, false, start), + ResponsePermit::Regular + ); + assert_eq!( + limiter.acquire(second, 50, false, start), + ResponsePermit::Regular + ); + assert_eq!( + limiter.acquire(priority, 50, false, start), + ResponsePermit::Rejected + ); + assert_eq!( + limiter.acquire(priority, 50, true, start), + ResponsePermit::PriorityReserve + ); + assert_eq!( + limiter.acquire("4.4.4.4:6881".parse().unwrap(), 50, true, start), + ResponsePermit::Rejected + ); + } +} diff --git a/dht-crawler/src/types.rs b/dht-crawler/src/types.rs new file mode 100644 index 0000000..004b07c --- /dev/null +++ b/dht-crawler/src/types.rs @@ -0,0 +1,408 @@ +// 负责定义 DHT 配置回调载荷和公开网络数据类型 + +use serde::{Deserialize, Serialize}; +use std::net::SocketAddr; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +/// IP families on which the DHT server listens and crawls. +pub enum NetMode { + /// Bind and crawl IPv4 only. + Ipv4Only, + /// Bind and crawl IPv6 only. + Ipv6Only, + #[default] + /// Bind separate IPv4 and IPv6 sockets. + DualStack, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +/// Validated torrent metadata delivered to the application callback. +pub struct TorrentInfo { + /// Lowercase hexadecimal SHA1 of the bencoded info dictionary. + pub info_hash: String, + /// Magnet URI containing the InfoHash. + pub magnet_link: String, + /// Torrent display name. + pub name: String, + /// Sum of file sizes in bytes. + pub total_size: u64, + /// Files described by the torrent. + pub files: Vec, + /// Torrent piece length in bytes, or zero if absent. + pub piece_length: u64, + /// Peer addresses used to obtain the Metadata. + pub peers: Vec, + /// Completion time as Unix seconds. + pub timestamp: u64, +} + +/// Final outcome of a metadata fetch that passed the admission callback. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum MetadataFetchCompletionStatus { + /// Metadata was fetched and accepted by the torrent callback. + Accepted, + /// All available peer candidates failed. + FetchFailed, + /// Metadata was fetched, but the application did not accept it. + DeliveryRejected, +} + +/// Report emitted exactly once after an admitted metadata fetch finishes. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct MetadataFetchCompletion { + /// InfoHash that reached a terminal state. + pub info_hash: String, + /// Final download/delivery status. + pub status: MetadataFetchCompletionStatus, + /// Real Peer network attempts; failure-cache skips are excluded. + pub attempts: usize, +} + +impl MetadataFetchCompletion { + /// Returns true only for [`MetadataFetchCompletionStatus::Accepted`]. + pub fn is_success(&self) -> bool { + self.status == MetadataFetchCompletionStatus::Accepted + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +/// One file entry from the validated info dictionary. +pub struct FileInfo { + /// Slash-separated relative path. + pub path: String, + /// File size in bytes. + pub size: u64, +} + +#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)] +/// Compact DHT node tuple used by crawl and routing code. +pub struct NodeTuple { + /// Twenty-byte DHT node ID. + pub id: [u8; 20], + /// Public UDP endpoint. + pub addr: SocketAddr, +} + +impl TorrentInfo { + /// Formats [`Self::total_size`] using binary thresholds and a short unit suffix. + pub fn format_size(&self) -> String { + format_bytes(self.total_size) + } +} + +impl FileInfo { + /// Formats [`Self::size`] using binary thresholds and a short unit suffix. + pub fn format_size(&self) -> String { + format_bytes(self.size) + } +} + +fn format_bytes(bytes: u64) -> String { + const UNITS: &[&str] = &["B", "KB", "MB", "GB", "TB"]; + let mut size = bytes as f64; + let mut unit_index = 0; + while size >= 1024.0 && unit_index < UNITS.len() - 1 { + size /= 1024.0; + unit_index += 1; + } + format!("{size:.2} {}", UNITS[unit_index]) +} + +#[derive(Debug, Clone)] +/// Complete server configuration. +pub struct DHTOptions { + /// UDP listen port. + pub port: u16, + /// Enabled IP families. + 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. + pub peer_lookup: PeerLookupOptions, + /// Active BEP-51 InfoHash sampling limits. + pub sample_infohashes: SampleInfohashesOptions, + /// Active crawl, node-pool and scheduler limits. + pub crawl: CrawlOptions, +} + +#[derive(Debug, Clone)] +/// Metadata download and failure-cache limits. +pub struct MetadataOptions { + /// End-to-end timeout for one Peer attempt, in seconds. + pub timeout_secs: u64, + /// Maximum number of deduplicated pending InfoHashes. + 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. + pub peer_failure_ttl_secs: u64, +} + +#[derive(Debug, Clone)] +/// Active get_peers lookup budgets used to discover additional Metadata Peers. +pub struct PeerLookupOptions { + /// Maximum new InfoHash lookups started per second. Zero disables active lookup. + pub max_lookups_per_second: u32, + /// Maximum lookup budget consumed immediately after an idle period. + pub burst: u32, + /// Maximum InfoHash lookups kept active at the same time. + pub max_active_lookups: usize, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +/// Terminal result of one bounded active `get_peers` lookup. +pub struct PeerLookupResult { + /// Unique Peer socket addresses returned by the DHT. + pub peers: Vec, + /// UDP queries sent for this lookup. + pub queries: usize, +} + +#[derive(Debug, Clone)] +/// Active BEP-51 `sample_infohashes` discovery limits. +pub struct SampleInfohashesOptions { + /// Maximum BEP-51 queries started per second. Zero disables sampling. + pub max_queries_per_second: u32, + /// Maximum sampling budget consumed immediately after an idle period. + pub burst: u32, + /// Maximum outstanding BEP-51 requests. + pub max_in_flight: usize, + /// Per-request timeout in milliseconds. + pub request_timeout_millis: u64, + /// Retry delay for timeouts or nodes that do not return samples, in seconds. + pub unsupported_backoff_secs: u64, + /// Maximum sampled InfoHashes retained for bounded in-memory deduplication. + pub dedup_capacity: usize, +} + +#[derive(Debug, Clone, Default)] +/// Active crawl configuration grouped by responsibility. +pub struct CrawlOptions { + /// FIFO node-pool and responsive-ring limits. + pub pool: PoolOptions, + /// Query, replacement and response budgets. + pub rate_limit: RateLimitOptions, + /// Bootstrap sources and retry policy. + pub bootstrap: BootstrapOptions, + /// Target-generation policy. + pub target: TargetOptions, + /// Internal bounded-channel and snapshot limits. + pub scheduler: SchedulerOptions, +} + +#[derive(Debug, Clone)] +/// Independent crawl and UDP-response budgets. +pub struct RateLimitOptions { + /// Maximum active find_node queries scheduled per second. + pub max_find_node_rate_per_sec: u32, + /// Maximum query budget consumed in one scheduler tick. + pub burst: u32, + /// Maximum total pending find_node transactions. + pub max_in_flight: usize, + /// Pending find_node timeout in seconds. + pub request_timeout_secs: u64, + /// Maximum never-before-probed destinations per minute. + pub max_new_destinations_per_minute: u32, + /// Maximum outbound DHT response packets per second. + pub max_response_rate_per_sec: u32, + /// Maximum encoded outbound DHT response bytes per second. + pub max_response_bytes_per_sec: u64, + /// Maximum response packets per source address per second. + pub max_response_rate_per_source: u32, + /// Remaining query-rate percentage when Metadata pressure reaches 95%. + pub metadata_pressure_floor_percent: u8, + /// Maximum FIFO replacements per minute after the pool has warmed. + pub max_replacements_per_minute: u32, + /// Maximum pending find_node transactions per IP subnet. + pub max_in_flight_per_subnet: usize, +} + +#[derive(Debug, Clone)] +/// FIFO crawl-pool and responsive-node reservoir limits. +pub struct PoolOptions { + /// Maximum queued crawl nodes. + pub capacity: usize, + /// How long a probed endpoint is blocked from readmission, in seconds. + pub recent_probe_ttl_secs: u64, + /// Maximum nodes retained for replies and revisit traffic. + pub responsive_capacity: usize, + /// Responsive-node lifetime in seconds. + pub responsive_ttl_secs: u64, + /// Pool size below which bootstrap is considered. + pub low_watermark: usize, +} + +#[derive(Debug, Clone)] +/// Bootstrap hostnames and retry timing. +pub struct BootstrapOptions { + /// Host:port sources resolved when bootstrap is needed. + pub nodes: Vec, + /// Minimum interval between bootstrap rounds, in seconds. + pub interval_secs: u64, + /// Maximum resolved endpoints selected in one round. + pub max_nodes_per_round: usize, + /// Initial failed-source backoff in seconds. + pub source_backoff_base_secs: u64, + /// Maximum failed-source backoff in seconds. + pub source_backoff_max_secs: u64, +} + +#[derive(Debug, Clone)] +/// Distribution used to generate find_node targets and sender IDs. +pub struct TargetOptions { + /// Percentage of targets that are fully random. + pub random_walk_percent: u8, + /// Percentage of targets chosen from sparse routing buckets. + pub sparse_bucket_percent: u8, + /// Whether outbound sender IDs borrow the target's prefix. + pub neighbor_sender_id: bool, +} + +#[derive(Debug, Clone)] +/// Capacities and batch limits for the crawl actor. +pub struct SchedulerOptions { + /// Capacity for response/bootstrap priority events. + pub priority_event_channel_capacity: usize, + /// Capacity for newly discovered node events. + pub discovery_event_channel_capacity: usize, + /// Maximum events drained per actor iteration. + pub event_batch_limit: usize, + /// Maximum discovery nodes drained per actor iteration. + pub node_batch_limit: usize, + /// Maximum responsive nodes published in the lock-free snapshot. + pub routing_snapshot_size: usize, + /// Snapshot publication interval in milliseconds. + pub snapshot_refresh_millis: u64, +} + +impl Default for DHTOptions { + fn default() -> Self { + Self { + 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(), + crawl: CrawlOptions::default(), + } + } +} + +impl Default for MetadataOptions { + fn default() -> Self { + Self { + timeout_secs: 4, + max_queue_size: 10_000, + max_worker_count: 8, + max_connects_per_second: 2, + peer_failure_cache_capacity: 200_000, + peer_failure_ttl_secs: 60, + } + } +} + +impl Default for PeerLookupOptions { + fn default() -> Self { + Self { + max_lookups_per_second: 1, + burst: 1, + max_active_lookups: 4, + } + } +} + +impl Default for SampleInfohashesOptions { + fn default() -> Self { + Self { + 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, + } + } +} + +impl Default for RateLimitOptions { + fn default() -> Self { + Self { + max_find_node_rate_per_sec: 6, + burst: 2, + max_in_flight: 12, + request_timeout_secs: 2, + 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, + metadata_pressure_floor_percent: 25, + max_replacements_per_minute: 25_000, + max_in_flight_per_subnet: 8, + } + } +} + +impl Default for PoolOptions { + fn default() -> Self { + Self { + capacity: 100_000, + recent_probe_ttl_secs: 600, + responsive_capacity: 16_384, + responsive_ttl_secs: 900, + low_watermark: 10_000, + } + } +} + +impl Default for BootstrapOptions { + fn default() -> Self { + Self { + nodes: vec![ + "router.bittorrent.com:6881".to_string(), + "dht.transmissionbt.com:6881".to_string(), + "router.utorrent.com:6881".to_string(), + "dht.aelitis.com:6881".to_string(), + ], + interval_secs: 30, + max_nodes_per_round: 16, + source_backoff_base_secs: 60, + source_backoff_max_secs: 3_600, + } + } +} + +impl Default for TargetOptions { + fn default() -> Self { + Self { + random_walk_percent: 70, + sparse_bucket_percent: 30, + neighbor_sender_id: true, + } + } +} + +impl Default for SchedulerOptions { + fn default() -> Self { + Self { + priority_event_channel_capacity: 8_192, + discovery_event_channel_capacity: 16_384, + event_batch_limit: 256, + node_batch_limit: 4_096, + routing_snapshot_size: 4_096, + snapshot_refresh_millis: 1_000, + } + } +} diff --git a/dht-crawler/src/udp_buffer.rs b/dht-crawler/src/udp_buffer.rs new file mode 100644 index 0000000..d11d339 --- /dev/null +++ b/dht-crawler/src/udp_buffer.rs @@ -0,0 +1,79 @@ +//! UDP 收包缓冲区池：避免每包 `to_owned()` 拷贝。 +//! +//! 单线程 listener 从池中取出固定大小缓冲区，`recv_from` 直接写入； +//! 通过 channel 将缓冲区所有权交给 worker，处理完毕后归还池中复用。 + +use crossbeam_queue::ArrayQueue; +use std::sync::Arc; + +/// 与 `process_udp_packet` 中丢弃阈值一致 +pub const MAX_DHT_UDP_PACKET: usize = 8192; + +/// 预分配缓冲区数量（约等于高峰在途包数） +const INITIAL_POOL_SIZE: usize = 512; +/// 池上限，防止极端背压下无限增长 +const MAX_POOL_SIZE: usize = 4096; + +/// 在途 UDP 包：固定容量缓冲区 + 有效长度 +pub struct UdpPacket { + pub buf: Box<[u8]>, + pub len: usize, +} + +impl UdpPacket { + #[inline] + pub fn payload(&self) -> &[u8] { + &self.buf[..self.len] + } +} + +/// 固定 8KiB 缓冲区的对象池（`recv_from` 零拷贝移交 worker） +#[derive(Clone)] +pub struct UdpBufferPool { + inner: Arc, +} + +struct PoolInner { + free: ArrayQueue>, + buf_capacity: usize, +} + +impl UdpBufferPool { + pub fn new() -> Self { + let free = ArrayQueue::new(MAX_POOL_SIZE); + for _ in 0..INITIAL_POOL_SIZE { + let _ = free.push(alloc_buffer(MAX_DHT_UDP_PACKET)); + } + Self { + inner: Arc::new(PoolInner { + free, + buf_capacity: MAX_DHT_UDP_PACKET, + }), + } + } + + /// 取一块缓冲区；池空时分配新块（背压或突发流量） + pub fn acquire(&self) -> Box<[u8]> { + self.inner + .free + .pop() + .unwrap_or_else(|| alloc_buffer(self.inner.buf_capacity)) + } + + /// 归还缓冲区；池满时直接丢弃，由 GC 回收 + pub fn release(&self, buf: Box<[u8]>) { + if buf.len() != self.inner.buf_capacity { + return; + } + let _ = self.inner.free.push(buf); + } + + pub fn buf_capacity(&self) -> usize { + self.inner.buf_capacity + } +} + +fn alloc_buffer(capacity: usize) -> Box<[u8]> { + let v = vec![0; capacity]; + v.into_boxed_slice() +} diff --git a/dht-crawler/src/udp_ingress.rs b/dht-crawler/src/udp_ingress.rs new file mode 100644 index 0000000..bd6e8fb --- /dev/null +++ b/dht-crawler/src/udp_ingress.rs @@ -0,0 +1,345 @@ +use crate::error::DHTError; +use crate::runtime_stats::DhtRuntimeStats; +use crate::udp_buffer::{MAX_DHT_UDP_PACKET, UdpBufferPool, UdpPacket}; +#[cfg(feature = "metrics")] +use metrics::counter; +use std::hash::{Hash, Hasher}; +use std::net::SocketAddr; +use std::sync::Arc; +use std::time::Duration; +use tokio::net::UdpSocket; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; + +pub(crate) type WorkerHandle = mpsc::Sender<(UdpPacket, SocketAddr, SocketAddr)>; + +pub(crate) fn spawn_udp_listener( + socket: Arc, + mut workers: Vec, + shutdown: CancellationToken, + buffer_pool: UdpBufferPool, + runtime_stats: DhtRuntimeStats, +) -> crate::error::Result<()> { + let local_addr = socket + .local_addr() + .map_err(|e| DHTError::Init(format!("socket local addr failed: {e}")))?; + if workers.is_empty() { + return Err(DHTError::Init( + "spawn_udp_listener: no worker provided".to_string(), + )); + } + tokio::spawn(async move { + loop { + let mut buf = buffer_pool.acquire(); + let recv_buf = &mut buf[..buffer_pool.buf_capacity()]; + + tokio::select! { + _ = shutdown.cancelled() => { + buffer_pool.release(buf); + break; + } + result = socket.recv_from(recv_buf) => { + match result { + Ok((size, origin_addr)) => { + if let Err(ProcessUdpPacketError::NoLiveWorkers) = + process_udp_packet(buf, size, origin_addr, local_addr, &buffer_pool, &runtime_stats, &mut workers) + { + log::warn!("Socket {socket:?} is closing because no worker can process packets."); + break + } + } + Err(_) => { + buffer_pool.release(buf); + tokio::select! { + _ = shutdown.cancelled() => break, + _ = tokio::time::sleep(Duration::from_millis(1)) => {}, + } + } + } + } + } + } + }); + Ok(()) +} + +enum ProcessUdpPacketError { + PacketTooLarge, + InvalidPacket, + ChokedWorkers, + NoLiveWorkers, +} + +fn process_udp_packet( + buf: Box<[u8]>, + size: usize, + origin_addr: SocketAddr, + local_addr: SocketAddr, + buffer_pool: &UdpBufferPool, + runtime_stats: &DhtRuntimeStats, + workers: &mut Vec, +) -> std::result::Result<(), ProcessUdpPacketError> { + runtime_stats.udp_received(); + runtime_stats.udp_received_bytes(size); + #[cfg(feature = "metrics")] + counter!("dht_udp_bytes_received_total").increment(size as u64); + + if size > MAX_DHT_UDP_PACKET { + runtime_stats.udp_invalid(); + #[cfg(feature = "metrics")] + counter!("dht_udp_packets_received_total", "status" => "dropped_size").increment(1); + buffer_pool.release(buf); + return Err(ProcessUdpPacketError::PacketTooLarge); + } + + if size == 0 || buf[0] != b'd' { + runtime_stats.udp_invalid(); + #[cfg(feature = "metrics")] + counter!("dht_udp_packets_received_total", "status" => "dropped_magic").increment(1); + buffer_pool.release(buf); + return Err(ProcessUdpPacketError::InvalidPacket); + } + + let mut packet = UdpPacket { buf, len: size }; + let mut hasher = ahash::AHasher::default(); + origin_addr.hash(&mut hasher); + let origin_hash = hasher.finish() as usize; + + 'select_worker: loop { + if workers.is_empty() { + buffer_pool.release(packet.buf); + return Err(ProcessUdpPacketError::NoLiveWorkers); + } + + let worker_count = workers.len(); + let preferred_index = origin_hash % worker_count; + for offset in 0..worker_count { + let worker_index = (preferred_index + offset) % worker_count; + match workers[worker_index].try_send((packet, origin_addr, local_addr)) { + Ok(_) => { + #[cfg(feature = "metrics")] + counter!("dht_udp_packets_received_total", "status" => "ok").increment(1); + return Ok(()); + } + Err(mpsc::error::TrySendError::Full((p, _, _))) => { + packet = p; + } + Err(mpsc::error::TrySendError::Closed((p, _, _))) => { + packet = p; + log::warn!("UDP worker dropped."); + workers.swap_remove(worker_index); + continue 'select_worker; + } + } + } + + #[cfg(feature = "metrics")] + counter!("dht_udp_packets_received_total", "status" => "queue_full").increment(1); + runtime_stats.udp_queue_full(); + buffer_pool.release(packet.buf); + return Err(ProcessUdpPacketError::ChokedWorkers); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn addresses() -> (SocketAddr, SocketAddr) { + ( + "8.8.8.8:6881".parse().unwrap(), + "0.0.0.0:12313".parse().unwrap(), + ) + } + + fn buffer(pool: &UdpBufferPool, first: u8) -> Box<[u8]> { + let mut buf = pool.acquire(); + buf[0] = first; + buf + } + + fn packet(pool: &UdpBufferPool, first: u8) -> UdpPacket { + UdpPacket { + buf: buffer(pool, first), + len: 1, + } + } + + fn preferred_index(origin_addr: SocketAddr, worker_count: usize) -> usize { + let mut hasher = ahash::AHasher::default(); + origin_addr.hash(&mut hasher); + (hasher.finish() as usize) % worker_count + } + + #[test] + fn available_preferred_worker_is_used_first() { + let pool = UdpBufferPool::new(); + let stats = DhtRuntimeStats::default(); + let (origin_addr, local_addr) = addresses(); + let (tx0, mut rx0) = mpsc::channel(1); + let (tx1, mut rx1) = mpsc::channel(1); + let mut workers = vec![tx0, tx1]; + let preferred = preferred_index(origin_addr, workers.len()); + + assert!( + process_udp_packet( + buffer(&pool, b'd'), + 1, + origin_addr, + local_addr, + &pool, + &stats, + &mut workers, + ) + .is_ok() + ); + + let (preferred_rx, fallback_rx) = if preferred == 0 { + (&mut rx0, &mut rx1) + } else { + (&mut rx1, &mut rx0) + }; + let (forwarded, _, _) = preferred_rx.try_recv().unwrap(); + assert_eq!(forwarded.payload(), b"d"); + assert!(matches!( + fallback_rx.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + )); + pool.release(forwarded.buf); + } + + #[test] + fn full_preferred_worker_falls_back_to_available_worker() { + let pool = UdpBufferPool::new(); + let stats = DhtRuntimeStats::default(); + let (origin_addr, local_addr) = addresses(); + let (tx0, mut rx0) = mpsc::channel(1); + let (tx1, mut rx1) = mpsc::channel(1); + let mut workers = vec![tx0, tx1]; + let preferred = preferred_index(origin_addr, workers.len()); + + workers[preferred] + .try_send((packet(&pool, b'x'), origin_addr, local_addr)) + .unwrap(); + + assert!( + process_udp_packet( + buffer(&pool, b'd'), + 1, + origin_addr, + local_addr, + &pool, + &stats, + &mut workers, + ) + .is_ok() + ); + + let (preferred_rx, fallback_rx) = if preferred == 0 { + (&mut rx0, &mut rx1) + } else { + (&mut rx1, &mut rx0) + }; + let (queued, _, _) = preferred_rx.try_recv().unwrap(); + let (forwarded, _, _) = fallback_rx.try_recv().unwrap(); + assert_eq!(queued.payload(), b"x"); + assert_eq!(forwarded.payload(), b"d"); + pool.release(queued.buf); + pool.release(forwarded.buf); + } + + #[test] + fn closed_preferred_worker_is_removed_before_fallback() { + let pool = UdpBufferPool::new(); + let stats = DhtRuntimeStats::default(); + let (origin_addr, local_addr) = addresses(); + let (closed_tx, closed_rx) = mpsc::channel(1); + drop(closed_rx); + let (open_tx, mut open_rx) = mpsc::channel(1); + let preferred = preferred_index(origin_addr, 2); + let mut workers = if preferred == 0 { + vec![closed_tx, open_tx] + } else { + vec![open_tx, closed_tx] + }; + + assert!( + process_udp_packet( + buffer(&pool, b'd'), + 1, + origin_addr, + local_addr, + &pool, + &stats, + &mut workers, + ) + .is_ok() + ); + + assert_eq!(workers.len(), 1); + let (forwarded, _, _) = open_rx.try_recv().unwrap(); + assert_eq!(forwarded.payload(), b"d"); + pool.release(forwarded.buf); + } + + #[test] + fn packet_is_dropped_only_after_all_live_workers_are_full() { + let pool = UdpBufferPool::new(); + let stats = DhtRuntimeStats::default(); + let (origin_addr, local_addr) = addresses(); + let (tx0, mut rx0) = mpsc::channel(1); + let (tx1, mut rx1) = mpsc::channel(1); + let mut workers = vec![tx0, tx1]; + for worker in &workers { + worker + .try_send((packet(&pool, b'x'), origin_addr, local_addr)) + .unwrap(); + } + + let result = process_udp_packet( + buffer(&pool, b'd'), + 1, + origin_addr, + local_addr, + &pool, + &stats, + &mut workers, + ); + + assert!(matches!(result, Err(ProcessUdpPacketError::ChokedWorkers))); + let snapshot = stats.snapshot(); + assert_eq!(snapshot.udp_received, 1); + assert_eq!(snapshot.udp_queue_full, 1); + assert_eq!(snapshot.udp_invalid, 0); + assert_eq!(workers.len(), 2); + for receiver in [&mut rx0, &mut rx1] { + let (queued, _, _) = receiver.try_recv().unwrap(); + assert_eq!(queued.payload(), b"x"); + pool.release(queued.buf); + } + } + + #[test] + fn invalid_packet_updates_runtime_stats() { + let pool = UdpBufferPool::new(); + let stats = DhtRuntimeStats::default(); + let (origin_addr, local_addr) = addresses(); + let mut workers = Vec::new(); + + let result = process_udp_packet( + buffer(&pool, b'x'), + 1, + origin_addr, + local_addr, + &pool, + &stats, + &mut workers, + ); + + assert!(matches!(result, Err(ProcessUdpPacketError::InvalidPacket))); + let snapshot = stats.snapshot(); + assert_eq!(snapshot.udp_received, 1); + assert_eq!(snapshot.udp_invalid, 1); + assert_eq!(snapshot.udp_queue_full, 0); + } +} diff --git a/dht-search.example.toml b/dht-search.example.toml new file mode 100644 index 0000000..8d31df9 --- /dev/null +++ b/dht-search.example.toml @@ -0,0 +1,35 @@ +# 定义 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" + +[verification] +enabled = true +queue_capacity = 10000 +max_active = 2 +max_peer_attempts = 3 +lease_secs = 60 +poll_interval_millis = 250 diff --git a/dht-search/Cargo.toml b/dht-search/Cargo.toml new file mode 100644 index 0000000..50d8ab7 --- /dev/null +++ b/dht-search/Cargo.toml @@ -0,0 +1,37 @@ +# 定义 DHT 元数据搜索应用的独立依赖和构建入口 + +[package] +name = "dht-search" +version = "0.1.0" +edition.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +publish = false + +[features] +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"] } +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" +tower = { version = "0.5", features = ["util"] } diff --git a/dht-search/README.md b/dht-search/README.md new file mode 100644 index 0000000..a63d9cc --- /dev/null +++ b/dht-search/README.md @@ -0,0 +1,136 @@ +# 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 连接数量 | + +Metadata 下载和可用性握手共用 `metadata_connects_per_second` 预算不会各自叠加 + +### 按需可用性验证 + +搜索结果和详情访问只会把已过冷却期的种子异步加入持久化验证队列 HTTP 响应不会等待 DHT 或 Peer 网络 + +| 配置项 | 默认值 | 作用 | +|---|---:|---| +| `verification.enabled` | `true` | 是否启用按需可用性验证 | +| `verification.queue_capacity` | `10000` | 持久化验证队列容量 | +| `verification.max_active` | `2` | 同时验证的种子数量 | +| `verification.max_peer_attempts` | `3` | 每个种子最多握手的 Peer 数量 | +| `verification.lease_secs` | `60` | 异常退出后验证任务重新可领取的租约时间 | +| `verification.poll_interval_millis` | `250` | 持久化队列轮询间隔 | + +详情访问使用高优先级 搜索结果使用普通优先级 队列满时高优先级可以替换最旧普通任务 + +可用性分为 `unknown` `active` 和 `possibly_stale` 一次或多次验证失败只表示当前可能没有可连接 Peer 不会删除种子 + +新抓取记录会把成功下载 Metadata 的来源 Peer 视为一次有效验证 旧记录按需复查时会同时使用新 DHT 结果和已保存的成功来源 Peer + +热度是近期 DHT 发现强度最近出现时间和可连接 Peer 数的综合活跃度分数 不代表全球下载量 + +桌面网络不要在不了解路由器 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 索引执行不会把全部记录加载到内存过滤 + +搜索响应包含 `heat` 和 `availability` 摘要 详情响应包含完整验证时间 Peer 数和连续失败次数 + +## 数据恢复 + +RocksDB 是权威数据源而 Tantivy 是可重建索引 + +当 Tantivy 目录不存在时应用会把全部 RocksDB 记录重新标记为待索引并自动完成全量重建 + +RocksDB v1 会在启动时通过可恢复迁移升级到 v2 旧 Tantivy schema 会先备份再从 RocksDB 重建 重建提交完成后清理备份 + +正常退出会先停止 DHT 再排空持久化队列提交剩余索引最后关闭 HTTP 服务 diff --git a/dht-search/src/api/handlers.rs b/dht-search/src/api/handlers.rs new file mode 100644 index 0000000..e5f8dd8 --- /dev/null +++ b/dht-search/src/api/handlers.rs @@ -0,0 +1,174 @@ +// 负责处理搜索详情统计和健康检查请求 + +use std::str::FromStr; + +use axum::{ + Json, + extract::{Path, Query, State}, + http::StatusCode, + response::{IntoResponse, Response}, +}; +use dht_search::{ + domain::InfoHash, + search::{SearchOptions, SearchPage}, + storage::VerificationPriority, +}; + +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(); + let verification = state + .verification + .as_ref() + .map(|ingress| ingress.stats().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(), + verification_queue: verification.map_or(0, |stats| stats.queue_depth), + verification_accepted: verification.map_or(0, |stats| stats.accepted), + verification_deduplicated: verification.map_or(0, |stats| stats.deduplicated), + verification_rejected_full: verification.map_or(0, |stats| stats.rejected_full), + verification_started: verification.map_or(0, |stats| stats.started), + verification_succeeded: verification.map_or(0, |stats| stats.succeeded), + verification_failed: verification.map_or(0, |stats| stats.failed), + verification_peers_discovered: verification.map_or(0, |stats| stats.peers_discovered), + verification_handshakes_succeeded: verification + .map_or(0, |stats| stats.handshakes_succeeded), + }) +} + +pub(crate) async fn search( + State(state): State, + Query(request): Query, +) -> Result, ApiError> { + let verification = state.verification.clone(); + 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()))?; + if let Some(verification) = &verification { + let hashes = page + .hits + .iter() + .filter_map(|hit| InfoHash::from_str(&hit.info_hash).ok()) + .collect(); + verification + .enqueue(hashes, VerificationPriority::Normal) + .await; + } + Ok(Json(page)) +} + +pub(crate) async fn torrent( + State(state): State, + Path(info_hash): Path, +) -> Result, ApiError> { + let verification = state.verification.clone(); + 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"))?; + if let Some(verification) = &verification { + verification + .enqueue(vec![info_hash], VerificationPriority::High) + .await; + } + 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 new file mode 100644 index 0000000..a7e30bb --- /dev/null +++ b/dht-search/src/api/mod.rs @@ -0,0 +1,147 @@ +// 负责组合 HTTP 路由和共享接口状态但不直接访问数据库实现 + +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, verification::VerificationIngress}; + +#[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) verification: Option, +} + +pub(crate) async fn serve( + listen: SocketAddr, + state: ApiState, + cancel: CancellationToken, +) -> std::io::Result<()> { + let router = router(state); + let listener = tokio::net::TcpListener::bind(listen).await?; + tracing::info!(%listen, "HTTP 服务启动"); + axum::serve(listener, router) + .with_graceful_shutdown(cancel.cancelled_owned()) + .await +} + +fn router(state: ApiState) -> 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) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use axum::{ + body::{Body, to_bytes}, + http::{Request, StatusCode}, + }; + use dht_crawler::{DhtRuntimeStats, FileInfo, TorrentInfo}; + use dht_search::{ + domain::{InfoHash, TorrentRecord}, + search::SearchEngine, + storage::{RocksTorrentRepository, TorrentRepository}, + }; + use tempfile::TempDir; + use tower::ServiceExt; + + use crate::{crawler::pipeline::PersistencePipeline, verification::VerificationIngress}; + + use super::*; + + #[tokio::test] + async fn search_and_detail_return_user_fields_and_enqueue_verification() { + let directory = TempDir::new().unwrap(); + let repository = + Arc::new(RocksTorrentRepository::open(directory.path().join("rocksdb")).unwrap()); + let mut record = TorrentRecord::try_from(TorrentInfo { + info_hash: "0101010101010101010101010101010101010101".into(), + magnet_link: String::new(), + name: "Example Movie".into(), + total_size: 42, + files: vec![FileInfo { + path: "movie.mkv".into(), + size: 42, + }], + piece_length: 16_384, + peers: vec!["127.0.0.1:6881".into()], + timestamp: 10, + }) + .unwrap(); + record.availability = dht_search::domain::Availability::default(); + repository.upsert(record.clone()).unwrap(); + let search = SearchEngine::open(directory.path().join("tantivy")).unwrap(); + search.index_records(std::slice::from_ref(&record)).unwrap(); + let repository_trait: Arc = repository.clone(); + let persistence = PersistencePipeline::start(repository_trait.clone(), 4); + let verification = VerificationIngress::for_test(repository.clone(), 10); + let app = router(ApiState { + repository: repository_trait, + search, + dht_stats: DhtRuntimeStats::default(), + persistence: persistence.ingress.clone(), + verification: Some(verification), + }); + + let response = app + .clone() + .oneshot( + Request::builder() + .uri("/search?q=Example") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let json: serde_json::Value = + serde_json::from_slice(&to_bytes(response.into_body(), usize::MAX).await.unwrap()) + .unwrap(); + assert_eq!(json["hits"][0]["name"], "Example Movie"); + assert!(json["hits"][0]["heat"]["score"].is_number()); + assert_eq!(json["hits"][0]["availability"]["status"], "unknown"); + assert_eq!(repository.verification_queue_len().unwrap(), 1); + + let response = app + .oneshot( + Request::builder() + .uri(format!("/torrents/{}", InfoHash::from_bytes([1; 20]))) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let json: serde_json::Value = + serde_json::from_slice(&to_bytes(response.into_body(), usize::MAX).await.unwrap()) + .unwrap(); + assert_eq!(json["name"], "Example Movie"); + assert!( + json["magnet_link"] + .as_str() + .unwrap() + .starts_with("magnet:?xt=") + ); + assert_eq!(json["files"][0]["path"], "movie.mkv"); + assert_eq!(repository.verification_queue_len().unwrap(), 1); + + persistence.close_and_join().await.unwrap(); + } +} diff --git a/dht-search/src/api/request.rs b/dht-search/src/api/request.rs new file mode 100644 index 0000000..aa78dd9 --- /dev/null +++ b/dht-search/src/api/request.rs @@ -0,0 +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 new file mode 100644 index 0000000..4df2355 --- /dev/null +++ b/dht-search/src/api/response.rs @@ -0,0 +1,91 @@ +// 负责定义稳定的 HTTP 响应模型和领域对象转换边界 + +use dht_search::domain::{Availability, Heat, TorrentFile, TorrentRecord}; +use serde::Serialize; +use std::time::{SystemTime, UNIX_EPOCH}; + +#[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, + pub(crate) verification_queue: u64, + pub(crate) verification_accepted: u64, + pub(crate) verification_deduplicated: u64, + pub(crate) verification_rejected_full: u64, + pub(crate) verification_started: u64, + pub(crate) verification_succeeded: u64, + pub(crate) verification_failed: u64, + pub(crate) verification_peers_discovered: u64, + pub(crate) verification_handshakes_succeeded: 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, + pub(crate) heat: Heat, + pub(crate) availability: Availability, +} + +impl From for TorrentResponse { + fn from(record: TorrentRecord) -> Self { + let info_hash = record.info_hash.to_string(); + let heat = record.heat(unix_timestamp()); + 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, + heat, + availability: record.availability, + } + } +} + +fn unix_timestamp() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() +} diff --git a/dht-search/src/app.rs b/dht-search/src/app.rs new file mode 100644 index 0000000..028a2a4 --- /dev/null +++ b/dht-search/src/app.rs @@ -0,0 +1,361 @@ +// 负责连接采集存储索引和接口层并定义应用级启动顺序 + +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, verification, +}; + +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 任务完成" + ) + }); + + let verification_cancel = CancellationToken::new(); + let (verification_fatal_tx, mut verification_fatal) = tokio::sync::oneshot::channel(); + let mut verification_fatal_guard = None; + let (verification_ingress, verification_task) = if config.verification.enabled { + let (ingress, worker) = verification::start( + repository.clone(), + server.clone(), + config.verification.clone(), + verification_cancel.clone(), + ); + let cancel = verification_cancel.clone(); + let task = tokio::spawn(async move { + let result = match worker.await { + Ok(result) => result, + Err(error) => Err(error.to_string()), + }; + if !cancel.is_cancelled() { + let message = result + .as_ref() + .err() + .cloned() + .unwrap_or_else(|| "可用性验证 worker 意外停止".to_owned()); + let _ = verification_fatal_tx.send(message); + } + result + }); + (Some(ingress), Some(task)) + } else { + verification_fatal_guard = Some(verification_fatal_tx); + (None, None) + }; + + 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(), + verification: verification_ingress, + }, + 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)) + } + fatal = &mut verification_fatal => { + let message = fatal.unwrap_or_else(|_| "可用性验证 worker 意外停止".to_owned()); + Err(AppError::VerificationWorker(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}"))), + } + } + }; + + verification_cancel.cancel(); + drop(verification_fatal_guard); + if let Some(task) = verification_task { + match task.await { + Ok(Ok(())) => {} + Ok(Err(error)) => return Err(AppError::VerificationWorker(error)), + Err(error) => return Err(AppError::VerificationWorker(error.to_string())), + } + } + 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(_) => search.cleanup_rebuild_backup().map_err(|error| error.to_string())?, + 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; + } + } + search + .cleanup_rebuild_backup() + .map_err(|error| error.to_string())?; + 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 new file mode 100644 index 0000000..fdbdf74 --- /dev/null +++ b/dht-search/src/config.rs @@ -0,0 +1,341 @@ +// 负责加载校验和提供应用配置但不执行任何业务逻辑 + +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, + pub(crate) verification: VerificationConfig, +} + +#[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, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub(crate) struct VerificationConfig { + pub(crate) enabled: bool, + pub(crate) queue_capacity: usize, + pub(crate) max_active: usize, + pub(crate) max_peer_attempts: usize, + pub(crate) lease_secs: u64, + pub(crate) poll_interval_millis: u64, +} + +#[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.verification.queue_capacity == 0 + || self.verification.max_active == 0 + || self.verification.max_peer_attempts == 0 + || self.verification.lease_secs == 0 + || self.verification.poll_interval_millis == 0 + { + return Err(AppError::Config( + "验证队列容量并发尝试数租约和轮询间隔必须大于零".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(), + verification: VerificationConfig::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 Default for VerificationConfig { + fn default() -> Self { + Self { + enabled: true, + queue_capacity: 10_000, + max_active: 2, + max_peer_attempts: 3, + lease_secs: 60, + poll_interval_millis: 250, + } + } +} + +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 new file mode 100644 index 0000000..5611e51 --- /dev/null +++ b/dht-search/src/crawler/mod.rs @@ -0,0 +1,4 @@ +// 负责组合 DHT 发现 Metadata 下载和持久化提交管线 + +pub(crate) mod pipeline; +mod worker; diff --git a/dht-search/src/crawler/pipeline.rs b/dht-search/src/crawler/pipeline.rs new file mode 100644 index 0000000..ade932e --- /dev/null +++ b/dht-search/src/crawler/pipeline.rs @@ -0,0 +1,273 @@ +// 负责定义采集阶段之间的有界队列背压和任务流转规则 + +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, VerificationResult}, + storage::{ + UpsertOutcome, VerificationEnqueueOutcome, VerificationPriority, VerificationRequest, + }, + }; + + 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 enqueue_verification( + &self, + _: &[InfoHash], + _: VerificationPriority, + _: u64, + _: usize, + ) -> Result { + Ok(VerificationEnqueueOutcome::default()) + } + + fn claim_verification( + &self, + _: u64, + _: u64, + ) -> Result, StorageError> { + Ok(None) + } + + fn finish_verification( + &self, + _: InfoHash, + _: VerificationResult, + ) -> Result<(), StorageError> { + Ok(()) + } + + fn verification_queue_len(&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/crawler/worker.rs b/dht-search/src/crawler/worker.rs new file mode 100644 index 0000000..54ca193 --- /dev/null +++ b/dht-search/src/crawler/worker.rs @@ -0,0 +1 @@ +// 负责消费 infohash 并执行受资源限制的 Metadata 下载任务 diff --git a/dht-search/src/domain/fingerprint.rs b/dht-search/src/domain/fingerprint.rs new file mode 100644 index 0000000..f66d015 --- /dev/null +++ b/dht-search/src/domain/fingerprint.rs @@ -0,0 +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 new file mode 100644 index 0000000..d43e291 --- /dev/null +++ b/dht-search/src/domain/mod.rs @@ -0,0 +1,12 @@ +// 负责导出不依赖存储搜索和传输实现的核心领域模型 + +mod fingerprint; +mod torrent; + +pub use fingerprint::content_key; +#[cfg(test)] +pub(crate) use torrent::test_record; +pub use torrent::{ + Availability, AvailabilityStatus, Heat, HeatLevel, IndexState, InfoHash, RECORD_SCHEMA_VERSION, + TorrentFile, TorrentRecord, TorrentRecordError, VerificationResult, +}; diff --git a/dht-search/src/domain/torrent.rs b/dht-search/src/domain/torrent.rs new file mode 100644 index 0000000..e74f119 --- /dev/null +++ b/dht-search/src/domain/torrent.rs @@ -0,0 +1,417 @@ +// 负责定义种子元数据文件条目发现状态和索引状态模型 + +use dht_crawler::TorrentInfo; +use serde::{Deserialize, Serialize}; +use std::{fmt, str::FromStr}; + +use super::fingerprint::content_key; + +pub const RECORD_SCHEMA_VERSION: u16 = 2; +const MAX_STORED_PEERS: usize = 32; +const ACTIVITY_SCALE: u64 = 1_000; +const ACTIVITY_HALF_LIFE_SECS: f64 = 86_400.0; + +#[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, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] +#[serde(rename_all = "snake_case")] +pub enum AvailabilityStatus { + #[default] + Unknown, + Active, + PossiblyStale, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)] +pub struct Availability { + pub status: AvailabilityStatus, + pub last_verified_at: Option, + pub last_success_at: Option, + pub discovered_peers: u32, + pub reachable_peers: u32, + pub consecutive_failures: u32, + pub next_check_at: u64, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub struct VerificationResult { + pub verified_at: u64, + pub discovered_peers: u32, + pub reachable_peers: u32, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum HeatLevel { + Hot, + Active, + Normal, + Cold, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +pub struct Heat { + pub score: u8, + pub level: HeatLevel, +} + +impl Heat { + pub fn from_score(score: u8) -> Self { + let level = match score { + 75..=100 => HeatLevel::Hot, + 50..=74 => HeatLevel::Active, + 25..=49 => HeatLevel::Normal, + _ => HeatLevel::Cold, + }; + Self { score, level } + } +} + +#[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, + #[serde(default)] + pub availability: Availability, + #[serde(default = "default_activity_score")] + pub activity_score_millis: u64, + #[serde(default)] + pub activity_updated_at: u64, +} + +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); + self.update_activity(timestamp); + 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()); + } + } + } + + pub fn normalize_schema(&mut self) { + if self.activity_updated_at == 0 { + self.activity_updated_at = self.last_seen; + } + self.schema_version = RECORD_SCHEMA_VERSION; + } + + pub fn apply_verification(&mut self, result: VerificationResult) { + self.availability.last_verified_at = Some(result.verified_at); + self.availability.discovered_peers = result.discovered_peers; + self.availability.reachable_peers = result.reachable_peers; + if result.reachable_peers > 0 { + self.availability.status = AvailabilityStatus::Active; + self.availability.last_success_at = Some(result.verified_at); + self.availability.consecutive_failures = 0; + self.availability.next_check_at = result.verified_at.saturating_add(86_400); + } else { + self.availability.status = AvailabilityStatus::PossiblyStale; + self.availability.consecutive_failures = + self.availability.consecutive_failures.saturating_add(1); + self.availability.next_check_at = result + .verified_at + .saturating_add(failure_retry_secs(self.availability.consecutive_failures)); + } + self.index_state = IndexState::Pending; + } + + pub fn heat(&self, now: u64) -> Heat { + let activity = decayed_activity(self.activity_score_millis, self.activity_updated_at, now) + as f64 + / ACTIVITY_SCALE as f64; + let discovery = (activity.ln_1p() / 101_f64.ln()).clamp(0.0, 1.0); + let age = now.saturating_sub(self.last_seen) as f64; + let freshness = 2_f64.powf(-age / (7.0 * 86_400.0)); + let availability = match self.availability.status { + AvailabilityStatus::Unknown => 0.25, + AvailabilityStatus::PossiblyStale => 0.0, + AvailabilityStatus::Active => { + ((f64::from(self.availability.reachable_peers) + 1.0).ln() / 4_f64.ln()) + .clamp(0.0, 1.0) + } + }; + let score = (100.0 * (0.60 * discovery + 0.25 * freshness + 0.15 * availability)) + .round() + .clamp(0.0, 100.0) as u8; + Heat::from_score(score) + } + + fn update_activity(&mut self, timestamp: u64) { + self.activity_score_millis = decayed_activity( + self.activity_score_millis, + self.activity_updated_at, + timestamp, + ) + .saturating_add(ACTIVITY_SCALE); + self.activity_updated_at = self.activity_updated_at.max(timestamp); + } +} + +fn default_activity_score() -> u64 { + ACTIVITY_SCALE +} + +fn decayed_activity(value: u64, updated_at: u64, now: u64) -> u64 { + if updated_at == 0 || now <= updated_at { + return value; + } + let elapsed = now - updated_at; + (value as f64 * 2_f64.powf(-(elapsed as f64) / ACTIVITY_HALF_LIFE_SECS)).round() as u64 +} + +fn failure_retry_secs(failures: u32) -> u64 { + match failures { + 0 | 1 => 3_600, + 2 => 6 * 3_600, + 3 => 24 * 3_600, + 4 => 72 * 3_600, + _ => 7 * 24 * 3_600, + } +} + +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); + let reachable_peers = source_peers.len().min(u32::MAX as usize) as u32; + let availability = if reachable_peers > 0 { + Availability { + status: AvailabilityStatus::Active, + last_verified_at: Some(info.timestamp), + last_success_at: Some(info.timestamp), + discovered_peers: reachable_peers, + reachable_peers, + consecutive_failures: 0, + next_check_at: info.timestamp.saturating_add(86_400), + } + } else { + Availability::default() + }; + + Ok(Self { + 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, + availability, + activity_score_millis: ACTIVITY_SCALE, + activity_updated_at: info.timestamp, + }) + } +} + +#[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, + availability: Availability::default(), + activity_score_millis: ACTIVITY_SCALE, + activity_updated_at: timestamp, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use dht_crawler::FileInfo; + + #[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); + } + + #[test] + fn activity_decays_and_recent_observation_increases_heat() { + let mut record = test_record(1, 10); + let old_heat = record.heat(10 + 30 * 86_400).score; + record.observe_again(10 + 30 * 86_400, &[]); + let new_heat = record.heat(10 + 30 * 86_400).score; + assert!(new_heat > old_heat); + } + + #[test] + fn verification_success_and_failures_update_status_and_backoff() { + let mut record = test_record(1, 10); + record.apply_verification(VerificationResult { + verified_at: 100, + discovered_peers: 3, + reachable_peers: 2, + }); + assert_eq!(record.availability.status, AvailabilityStatus::Active); + assert_eq!(record.availability.next_check_at, 86_500); + + record.apply_verification(VerificationResult { + verified_at: 200, + discovered_peers: 0, + reachable_peers: 0, + }); + assert_eq!( + record.availability.status, + AvailabilityStatus::PossiblyStale + ); + assert_eq!(record.availability.consecutive_failures, 1); + assert_eq!(record.availability.next_check_at, 3_800); + } + + #[test] + fn freshly_downloaded_metadata_is_immediately_active() { + let record = TorrentRecord::try_from(TorrentInfo { + info_hash: "0101010101010101010101010101010101010101".into(), + magnet_link: String::new(), + name: "Example".into(), + total_size: 42, + files: vec![FileInfo { + path: "example.bin".into(), + size: 42, + }], + piece_length: 16_384, + peers: vec!["127.0.0.1:6881".into()], + timestamp: 100, + }) + .unwrap(); + assert_eq!(record.availability.status, AvailabilityStatus::Active); + assert_eq!(record.availability.reachable_peers, 1); + assert_eq!(record.availability.last_verified_at, Some(100)); + assert_eq!(record.availability.next_check_at, 86_500); + } +} diff --git a/dht-search/src/error.rs b/dht-search/src/error.rs new file mode 100644 index 0000000..0a20e5a --- /dev/null +++ b/dht-search/src/error.rs @@ -0,0 +1,25 @@ +// 负责定义应用层统一错误类型和跨模块错误转换边界 + +#[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), + #[error("可用性验证 worker 失败: {0}")] + VerificationWorker(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 new file mode 100644 index 0000000..e898860 --- /dev/null +++ b/dht-search/src/main.rs @@ -0,0 +1,25 @@ +// 负责组装应用依赖启动运行时并协调服务生命周期 + +use clap::Parser; + +mod api; +mod app; +mod config; +mod crawler; +mod error; +mod shutdown; +mod telemetry; +mod verification; + +#[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 new file mode 100644 index 0000000..e7eca98 --- /dev/null +++ b/dht-search/src/search/indexer.rs @@ -0,0 +1,485 @@ +// 负责批量写入删除提交和从权威存储重建 Tantivy 索引 + +use std::{ + collections::BTreeSet, + ops::Bound, + path::{Path, PathBuf}, + sync::{Arc, Mutex}, + time::{SystemTime, UNIX_EPOCH}, +}; + +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::{AvailabilityStatus, Heat, TorrentRecord}; +use crate::storage::TorrentRepository; + +use super::{ + IndexingError, SearchError, + query::{AvailabilitySummary, 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, + backup_path: Mutex>, +} + +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> { + let path = path.as_ref().to_path_buf(); + std::fs::create_dir_all(&path) + .map_err(|error| SearchError::Directory(error.to_string()))?; + let mut directory = MmapDirectory::open(&path) + .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 mut created = !exists; + let mut backup_path = None; + let index = if exists { + let index = Index::open(directory)?; + if index.schema() != expected_schema { + drop(index); + let backup = backup_path_for(&path); + std::fs::rename(&path, &backup) + .map_err(|error| SearchError::Directory(error.to_string()))?; + std::fs::create_dir_all(&path) + .map_err(|error| SearchError::Directory(error.to_string()))?; + directory = MmapDirectory::open(&path) + .map_err(|error| SearchError::Directory(error.to_string()))?; + backup_path = Some(backup); + created = true; + Index::open_or_create(directory, expected_schema)? + } else { + 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, + backup_path: Mutex::new(backup_path), + }), + }, + created, + )) + } + + pub fn cleanup_rebuild_backup(&self) -> Result<(), SearchError> { + let mut backup = self + .inner + .backup_path + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(path) = backup.take() + && path.exists() + { + std::fs::remove_dir_all(path) + .map_err(|error| SearchError::Directory(error.to_string()))?; + } + Ok(()) + } + + 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, + heat: Heat::from_score( + number(&document, fields.heat_score, "heat_score")?.min(100) as u8, + ), + availability: AvailabilitySummary { + status: availability_status(number( + &document, + fields.availability_status, + "availability_status", + )?), + last_verified_at: match number( + &document, + fields.last_verified_at, + "last_verified_at", + )? { + 0 => None, + value => Some(value), + }, + reachable_peers: number(&document, fields.reachable_peers, "reachable_peers")? + .min(u64::from(u32::MAX)) as u32, + }, + }); + } + 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.add_u64( + fields.availability_status, + match record.availability.status { + AvailabilityStatus::Unknown => 0, + AvailabilityStatus::Active => 1, + AvailabilityStatus::PossiblyStale => 2, + }, + ); + document.add_u64( + fields.reachable_peers, + u64::from(record.availability.reachable_peers), + ); + document.add_u64( + fields.last_verified_at, + record.availability.last_verified_at.unwrap_or(0), + ); + document.add_u64( + fields.heat_score, + u64::from(record.heat(unix_timestamp()).score), + ); + document +} + +fn availability_status(value: u64) -> AvailabilityStatus { + match value { + 1 => AvailabilityStatus::Active, + 2 => AvailabilityStatus::PossiblyStale, + _ => AvailabilityStatus::Unknown, + } +} + +fn unix_timestamp() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() +} + +fn backup_path_for(path: &Path) -> PathBuf { + let name = path + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or("tantivy"); + let base = format!("{name}.backup-{}", unix_timestamp()); + let mut candidate = path.with_file_name(&base); + let mut suffix = 1_u32; + while candidate.exists() { + candidate = path.with_file_name(format!("{base}-{suffix}")); + suffix = suffix.saturating_add(1); + } + candidate +} + +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, RECORD_SCHEMA_VERSION, TorrentFile, TorrentRecord}; + + use super::*; + + fn record() -> TorrentRecord { + TorrentRecord { + schema_version: RECORD_SCHEMA_VERSION, + 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, + availability: crate::domain::Availability::default(), + activity_score_millis: 1_000, + activity_updated_at: 20, + } + } + + #[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); + } + + #[test] + fn incompatible_schema_is_backed_up_and_recreated() { + let directory = TempDir::new().unwrap(); + let index_path = directory.path().join("tantivy"); + std::fs::create_dir(&index_path).unwrap(); + let mut old_schema = tantivy::schema::Schema::builder(); + old_schema.add_text_field("old_name", tantivy::schema::TEXT); + Index::create_in_dir(&index_path, old_schema.build()).unwrap(); + + let (engine, created) = SearchEngine::open_with_status(&index_path).unwrap(); + assert!(created); + let backups: Vec<_> = std::fs::read_dir(directory.path()) + .unwrap() + .filter_map(Result::ok) + .filter(|entry| { + entry + .file_name() + .to_string_lossy() + .starts_with("tantivy.backup-") + }) + .collect(); + assert_eq!(backups.len(), 1); + + engine.cleanup_rebuild_backup().unwrap(); + assert!(!backups[0].path().exists()); + assert!(index_path.exists()); + } +} diff --git a/dht-search/src/search/mod.rs b/dht-search/src/search/mod.rs new file mode 100644 index 0000000..2ee4b43 --- /dev/null +++ b/dht-search/src/search/mod.rs @@ -0,0 +1,30 @@ +// 负责暴露全文搜索抽象并隐藏 Tantivy 的具体实现细节 + +mod indexer; +mod query; +mod schema; + +pub use indexer::SearchEngine; +pub use query::{AvailabilitySummary, 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 new file mode 100644 index 0000000..00e6cfe --- /dev/null +++ b/dht-search/src/search/query.rs @@ -0,0 +1,45 @@ +// 负责构建全文查询过滤排序分页和内容聚合逻辑 + +use serde::Serialize; + +use crate::domain::{AvailabilityStatus, Heat}; + +#[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, + pub heat: Heat, + pub availability: AvailabilitySummary, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct AvailabilitySummary { + pub status: AvailabilityStatus, + pub last_verified_at: Option, + pub reachable_peers: u32, +} + +#[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 new file mode 100644 index 0000000..3dd4f50 --- /dev/null +++ b/dht-search/src/search/schema.rs @@ -0,0 +1,59 @@ +// 负责定义 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) availability_status: Field, + pub(crate) reachable_peers: Field, + pub(crate) last_verified_at: Field, + pub(crate) heat_score: 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 availability_status = builder.add_u64_field("availability_status", FAST | STORED); + let reachable_peers = builder.add_u64_field("reachable_peers", FAST | STORED); + let last_verified_at = builder.add_u64_field("last_verified_at", FAST | STORED); + let heat_score = builder.add_u64_field("heat_score", FAST | 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, + availability_status, + reachable_peers, + last_verified_at, + heat_score, + }, + ) +} diff --git a/dht-search/src/shutdown.rs b/dht-search/src/shutdown.rs new file mode 100644 index 0000000..6f79fad --- /dev/null +++ b/dht-search/src/shutdown.rs @@ -0,0 +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 new file mode 100644 index 0000000..86415b1 --- /dev/null +++ b/dht-search/src/storage/keys.rs @@ -0,0 +1,165 @@ +// 负责定义稳定的 RocksDB 键空间编码和版本边界 + +use crate::domain::InfoHash; + +pub(crate) const DATABASE_SCHEMA_VERSION: u32 = 2; +pub(crate) const SCHEMA_VERSION_KEY: &[u8] = b"\x00schema-version"; +pub(crate) const MIGRATION_KEY: &[u8] = b"\x00migration-v1-v2"; +pub(crate) const VERIFICATION_QUEUE_COUNT_KEY: &[u8] = b"\x00verification-queue-count"; +const TORRENT_PREFIX: u8 = b't'; +const CONTENT_PREFIX: u8 = b'c'; +const PENDING_INDEX_PREFIX: u8 = b'p'; +const VERIFICATION_HIGH_PREFIX: u8 = b'h'; +const VERIFICATION_NORMAL_PREFIX: u8 = b'n'; +const VERIFICATION_LEASE_PREFIX: u8 = b'l'; +const VERIFICATION_LOCATOR_PREFIX: u8 = b'v'; + +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 verification_task_key( + high_priority: bool, + requested_at: u64, + info_hash: InfoHash, +) -> [u8; 1 + 8 + InfoHash::BYTE_LEN] { + timed_info_hash_key( + if high_priority { + VERIFICATION_HIGH_PREFIX + } else { + VERIFICATION_NORMAL_PREFIX + }, + requested_at, + info_hash, + ) +} + +pub(crate) fn verification_task_prefix(high_priority: bool) -> [u8; 1] { + [if high_priority { + VERIFICATION_HIGH_PREFIX + } else { + VERIFICATION_NORMAL_PREFIX + }] +} + +pub(crate) fn verification_lease_key( + lease_until: u64, + info_hash: InfoHash, +) -> [u8; 1 + 8 + InfoHash::BYTE_LEN] { + timed_info_hash_key(VERIFICATION_LEASE_PREFIX, lease_until, info_hash) +} + +pub(crate) fn verification_lease_prefix() -> [u8; 1] { + [VERIFICATION_LEASE_PREFIX] +} + +pub(crate) fn verification_locator_key(info_hash: InfoHash) -> [u8; 1 + InfoHash::BYTE_LEN] { + prefixed_info_hash(VERIFICATION_LOCATOR_PREFIX, info_hash) +} + +pub(crate) fn decode_timed_info_hash(key: &[u8], prefix: u8) -> Option<(u64, InfoHash)> { + if key.len() != 1 + 8 + InfoHash::BYTE_LEN || key.first().copied() != Some(prefix) { + return None; + } + let timestamp = u64::from_be_bytes(key[1..9].try_into().ok()?); + let bytes = key[9..].try_into().ok()?; + Some((timestamp, InfoHash::from_bytes(bytes))) +} + +pub(crate) fn decode_verification_task(key: &[u8]) -> Option<(bool, u64, InfoHash)> { + match key.first().copied()? { + VERIFICATION_HIGH_PREFIX => { + decode_timed_info_hash(key, VERIFICATION_HIGH_PREFIX).map(|(at, hash)| (true, at, hash)) + } + VERIFICATION_NORMAL_PREFIX => decode_timed_info_hash(key, VERIFICATION_NORMAL_PREFIX) + .map(|(at, hash)| (false, at, hash)), + _ => None, + } +} + +pub(crate) fn decode_verification_lease(key: &[u8]) -> Option<(u64, InfoHash)> { + decode_timed_info_hash(key, VERIFICATION_LEASE_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 +} + +fn timed_info_hash_key( + prefix: u8, + timestamp: u64, + info_hash: InfoHash, +) -> [u8; 1 + 8 + InfoHash::BYTE_LEN] { + let mut key = [0_u8; 1 + 8 + InfoHash::BYTE_LEN]; + key[0] = prefix; + key[1..9].copy_from_slice(×tamp.to_be_bytes()); + key[9..].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 new file mode 100644 index 0000000..b45e0b9 --- /dev/null +++ b/dht-search/src/storage/mod.rs @@ -0,0 +1,13 @@ +// 负责暴露持久化抽象并隐藏 RocksDB 的具体实现细节 + +mod keys; +mod repository; +#[cfg(feature = "rocksdb-storage")] +mod rocks; + +pub use repository::{ + StorageError, TorrentRepository, UpsertOutcome, VerificationEnqueueOutcome, + VerificationPriority, VerificationRequest, +}; +#[cfg(feature = "rocksdb-storage")] +pub use rocks::RocksTorrentRepository; diff --git a/dht-search/src/storage/repository.rs b/dht-search/src/storage/repository.rs new file mode 100644 index 0000000..ce8a72e --- /dev/null +++ b/dht-search/src/storage/repository.rs @@ -0,0 +1,108 @@ +// 负责定义元数据去重状态恢复和索引任务所需的存储接口 + +use crate::domain::{InfoHash, TorrentRecord, VerificationResult}; + +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; + + fn enqueue_verification( + &self, + info_hashes: &[InfoHash], + priority: VerificationPriority, + requested_at: u64, + capacity: usize, + ) -> Result; + + fn claim_verification( + &self, + now: u64, + lease_secs: u64, + ) -> Result, StorageError>; + + fn finish_verification( + &self, + info_hash: InfoHash, + result: VerificationResult, + ) -> Result<(), StorageError>; + + fn verification_queue_len(&self) -> Result; +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum VerificationPriority { + Normal, + High, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct VerificationRequest { + pub info_hash: InfoHash, + pub priority: VerificationPriority, +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct VerificationEnqueueOutcome { + pub accepted: usize, + pub deduplicated: usize, + pub rejected_full: usize, +} + +#[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), + #[error("验证队列计数数据损坏")] + CorruptVerificationQueueCount, +} diff --git a/dht-search/src/storage/rocks.rs b/dht-search/src/storage/rocks.rs new file mode 100644 index 0000000..20405ca --- /dev/null +++ b/dht-search/src/storage/rocks.rs @@ -0,0 +1,799 @@ +// 负责实现 RocksDB 打开配置批量写入精确查询和关闭流程 + +use std::{path::Path, sync::Mutex}; + +use rocksdb::{ + BlockBasedOptions, Cache, DB, DBCompressionType, Direction, IteratorMode, Options, + SliceTransform, WriteBatch, +}; + +use crate::domain::{IndexState, InfoHash, TorrentRecord, VerificationResult}; + +use super::{ + keys::{ + DATABASE_SCHEMA_VERSION, MIGRATION_KEY, SCHEMA_VERSION_KEY, VERIFICATION_QUEUE_COUNT_KEY, + content_member_key, content_members_prefix, decode_content_member_info_hash, + decode_pending_info_hash, decode_verification_lease, decode_verification_task, + pending_index_key, pending_index_prefix, torrent_key, torrent_prefix, + verification_lease_key, verification_lease_prefix, verification_locator_key, + verification_task_key, verification_task_prefix, + }, + repository::{ + StorageError, TorrentRepository, UpsertOutcome, VerificationEnqueueOutcome, + VerificationPriority, VerificationRequest, + }, +}; + +const DEFAULT_BLOCK_CACHE_BYTES: usize = 64 * 1024 * 1024; +type VerificationTaskEntry = (Box<[u8]>, InfoHash); + +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.migrate_schema()?; + Ok(repository) + } + + fn migrate_schema(&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) if value.as_ref() == 1_u32.to_be_bytes() => { + let mut batch = WriteBatch::default(); + batch.put(SCHEMA_VERSION_KEY, expected); + batch.put(MIGRATION_KEY, []); + self.db.write(batch)?; + self.migrate_v1_records() + } + 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, + }) + } + }?; + if self.db.get(MIGRATION_KEY)?.is_some() { + self.migrate_v1_records()?; + } + Ok(()) + } + + fn migrate_v1_records(&self) -> Result<(), StorageError> { + const BATCH_SIZE: usize = 1_000; + let prefix = torrent_prefix(); + let iterator = self + .db + .iterator(IteratorMode::From(&prefix, Direction::Forward)); + let mut batch = WriteBatch::default(); + let mut batch_len = 0; + for entry in iterator { + let (key, value) = entry?; + if !key.starts_with(&prefix) { + break; + } + let mut record = Self::decode_raw(&value)?; + record.normalize_schema(); + record.index_state = IndexState::Pending; + batch.put(&key, Self::encode(&record)?); + batch.put(pending_index_key(record.info_hash), []); + batch_len += 1; + if batch_len == BATCH_SIZE { + self.db.write(batch)?; + batch = WriteBatch::default(); + batch_len = 0; + } + } + batch.delete(MIGRATION_KEY); + self.db.write(batch)?; + Ok(()) + } + + fn encode(record: &TorrentRecord) -> Result, StorageError> { + rmp_serde::to_vec_named(record).map_err(Into::into) + } + + fn decode_raw(bytes: &[u8]) -> Result { + rmp_serde::from_slice(bytes).map_err(Into::into) + } + + fn decode(bytes: &[u8]) -> Result { + let mut record = Self::decode_raw(bytes)?; + record.normalize_schema(); + Ok(record) + } + + fn verification_queue_len_inner(&self) -> Result { + let Some(value) = self.db.get(VERIFICATION_QUEUE_COUNT_KEY)? else { + return Ok(0); + }; + let bytes: [u8; 8] = value + .as_slice() + .try_into() + .map_err(|_| StorageError::CorruptVerificationQueueCount)?; + Ok(u64::from_be_bytes(bytes).min(usize::MAX as u64) as usize) + } + + fn first_verification_task( + &self, + high_priority: bool, + ) -> Result, StorageError> { + let prefix = verification_task_prefix(high_priority); + let mut iterator = self + .db + .iterator(IteratorMode::From(&prefix, Direction::Forward)); + let Some(entry) = iterator.next() else { + return Ok(None); + }; + let (key, _) = entry?; + let Some((_, _, info_hash)) = decode_verification_task(&key) else { + return Ok(None); + }; + Ok(Some((key, info_hash))) + } +} + +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) + } + + fn enqueue_verification( + &self, + info_hashes: &[InfoHash], + priority: VerificationPriority, + requested_at: u64, + capacity: usize, + ) -> Result { + let _guard = self + .write_lock + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut count = self.verification_queue_len_inner()?; + let mut outcome = VerificationEnqueueOutcome::default(); + let mut seen = std::collections::HashSet::with_capacity(info_hashes.len()); + for info_hash in info_hashes { + if !seen.insert(*info_hash) { + outcome.deduplicated += 1; + continue; + } + let Some(record) = self.get(*info_hash)? else { + outcome.deduplicated += 1; + continue; + }; + if record.availability.next_check_at > requested_at { + outcome.deduplicated += 1; + continue; + } + let locator_key = verification_locator_key(*info_hash); + if let Some(existing_key) = self.db.get(locator_key)? { + let can_promote = priority == VerificationPriority::High + && decode_verification_task(&existing_key).is_some_and(|(high, _, _)| !high); + if can_promote { + let new_key = verification_task_key(true, requested_at, *info_hash); + let mut batch = WriteBatch::default(); + batch.delete(existing_key); + batch.put(new_key, []); + batch.put(locator_key, new_key); + self.db.write(batch)?; + } + outcome.deduplicated += 1; + continue; + } + if count >= capacity { + if priority == VerificationPriority::High { + if let Some((evicted_key, evicted_hash)) = + self.first_verification_task(false)? + { + let mut batch = WriteBatch::default(); + batch.delete(evicted_key); + batch.delete(verification_locator_key(evicted_hash)); + count = count.saturating_sub(1); + batch.put(VERIFICATION_QUEUE_COUNT_KEY, (count as u64).to_be_bytes()); + self.db.write(batch)?; + } else { + outcome.rejected_full += 1; + continue; + } + } else { + outcome.rejected_full += 1; + continue; + } + } + let task_key = verification_task_key( + priority == VerificationPriority::High, + requested_at, + *info_hash, + ); + let mut batch = WriteBatch::default(); + batch.put(task_key, []); + batch.put(locator_key, task_key); + count += 1; + batch.put(VERIFICATION_QUEUE_COUNT_KEY, (count as u64).to_be_bytes()); + self.db.write(batch)?; + outcome.accepted += 1; + } + Ok(outcome) + } + + fn claim_verification( + &self, + now: u64, + lease_secs: u64, + ) -> Result, StorageError> { + let _guard = self + .write_lock + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let lease_prefix = verification_lease_prefix(); + let iterator = self + .db + .iterator(IteratorMode::From(&lease_prefix, Direction::Forward)); + let mut recovery = WriteBatch::default(); + let mut recovered = false; + for entry in iterator { + let (key, value) = entry?; + let Some((lease_until, info_hash)) = decode_verification_lease(&key) else { + break; + }; + if lease_until > now { + break; + } + let high = value.first().copied() == Some(1); + let task_key = verification_task_key(high, now, info_hash); + recovery.delete(&key); + recovery.put(task_key, []); + recovery.put(verification_locator_key(info_hash), task_key); + recovered = true; + } + if recovered { + self.db.write(recovery)?; + } + + let selected = match self.first_verification_task(true)? { + Some((key, hash)) => Some((VerificationPriority::High, key, hash)), + None => self + .first_verification_task(false)? + .map(|(key, hash)| (VerificationPriority::Normal, key, hash)), + }; + let Some((priority, task_key, info_hash)) = selected else { + return Ok(None); + }; + let lease_key = verification_lease_key(now.saturating_add(lease_secs), info_hash); + let mut batch = WriteBatch::default(); + batch.delete(task_key); + batch.put( + lease_key, + [u8::from(priority == VerificationPriority::High)], + ); + batch.put(verification_locator_key(info_hash), lease_key); + self.db.write(batch)?; + Ok(Some(VerificationRequest { + info_hash, + priority, + })) + } + + fn finish_verification( + &self, + info_hash: InfoHash, + result: VerificationResult, + ) -> 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.apply_verification(result); + let locator_key = verification_locator_key(info_hash); + let queued_key = self.db.get(locator_key)?; + let count = self.verification_queue_len_inner()?.saturating_sub(1); + let mut batch = WriteBatch::default(); + batch.put(torrent_key(info_hash), Self::encode(&record)?); + batch.put(pending_index_key(info_hash), []); + if let Some(queued_key) = queued_key { + batch.delete(queued_key); + } + batch.delete(locator_key); + batch.put(VERIFICATION_QUEUE_COUNT_KEY, (count as u64).to_be_bytes()); + self.db.write(batch)?; + Ok(()) + } + + fn verification_queue_len(&self) -> Result { + self.verification_queue_len_inner() + } +} + +#[cfg(test)] +mod tests { + use serde::Serialize; + use tempfile::TempDir; + + use crate::domain::{ + AvailabilityStatus, IndexState, InfoHash, TorrentFile, VerificationResult, 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 + ); + } + + #[derive(Serialize)] + struct LegacyTorrentRecord { + schema_version: u16, + info_hash: InfoHash, + name: String, + total_size: u64, + files: Vec, + piece_length: u64, + source_peers: Vec, + content_key: [u8; 32], + first_seen: u64, + last_seen: u64, + seen_count: u64, + index_state: IndexState, + } + + #[test] + fn version_one_database_migrates_records_and_marks_them_pending() { + let directory = TempDir::new().unwrap(); + let record = test_record(9, 10); + let legacy = LegacyTorrentRecord { + schema_version: 1, + info_hash: record.info_hash, + name: record.name, + total_size: record.total_size, + files: record.files, + piece_length: record.piece_length, + source_peers: record.source_peers, + content_key: record.content_key, + first_seen: record.first_seen, + last_seen: record.last_seen, + seen_count: record.seen_count, + index_state: IndexState::Indexed { indexed_at: 20 }, + }; + { + let mut options = Options::default(); + options.create_if_missing(true); + let db = DB::open(&options, directory.path()).unwrap(); + db.put(SCHEMA_VERSION_KEY, 1_u32.to_be_bytes()).unwrap(); + db.put( + torrent_key(legacy.info_hash), + rmp_serde::to_vec_named(&legacy).unwrap(), + ) + .unwrap(); + } + + let repository = RocksTorrentRepository::open(directory.path()).unwrap(); + let migrated = repository.get(legacy.info_hash).unwrap().unwrap(); + assert_eq!(migrated.schema_version, 2); + assert_eq!(migrated.availability.status, AvailabilityStatus::Unknown); + assert_eq!(migrated.activity_updated_at, 10); + assert_eq!(migrated.index_state, IndexState::Pending); + assert_eq!( + repository.pending_index(10).unwrap(), + vec![legacy.info_hash] + ); + assert_eq!( + repository.db.get(SCHEMA_VERSION_KEY).unwrap().unwrap(), + DATABASE_SCHEMA_VERSION.to_be_bytes() + ); + assert!(repository.db.get(MIGRATION_KEY).unwrap().is_none()); + } + + #[test] + fn verification_queue_is_persistent_prioritized_and_updates_record_atomically() { + let directory = TempDir::new().unwrap(); + let repository = RocksTorrentRepository::open(directory.path()).unwrap(); + let first = test_record(1, 10); + let second = test_record(2, 10); + let high = test_record(3, 10); + repository.upsert(first.clone()).unwrap(); + repository.upsert(second.clone()).unwrap(); + repository.upsert(high.clone()).unwrap(); + repository + .enqueue_verification( + &[first.info_hash, second.info_hash], + VerificationPriority::Normal, + 100, + 2, + ) + .unwrap(); + let outcome = repository + .enqueue_verification(&[high.info_hash], VerificationPriority::High, 101, 2) + .unwrap(); + assert_eq!(outcome.accepted, 1); + assert_eq!(repository.verification_queue_len().unwrap(), 2); + + drop(repository); + let repository = RocksTorrentRepository::open(directory.path()).unwrap(); + let claimed = repository.claim_verification(102, 60).unwrap().unwrap(); + assert_eq!(claimed.info_hash, high.info_hash); + assert_eq!(claimed.priority, VerificationPriority::High); + repository + .finish_verification( + high.info_hash, + VerificationResult { + verified_at: 103, + discovered_peers: 4, + reachable_peers: 2, + }, + ) + .unwrap(); + assert_eq!(repository.verification_queue_len().unwrap(), 1); + let stored = repository.get(high.info_hash).unwrap().unwrap(); + assert_eq!(stored.availability.status, AvailabilityStatus::Active); + assert_eq!(stored.availability.reachable_peers, 2); + assert_eq!(stored.index_state, IndexState::Pending); + } + + #[test] + fn expired_verification_lease_is_recovered() { + let directory = TempDir::new().unwrap(); + let repository = RocksTorrentRepository::open(directory.path()).unwrap(); + let record = test_record(4, 10); + repository.upsert(record.clone()).unwrap(); + repository + .enqueue_verification(&[record.info_hash], VerificationPriority::Normal, 100, 1) + .unwrap(); + repository.claim_verification(100, 5).unwrap().unwrap(); + let recovered = repository.claim_verification(106, 5).unwrap().unwrap(); + assert_eq!(recovered.info_hash, record.info_hash); + assert_eq!(repository.verification_queue_len().unwrap(), 1); + } +} diff --git a/dht-search/src/telemetry.rs b/dht-search/src/telemetry.rs new file mode 100644 index 0000000..81573d6 --- /dev/null +++ b/dht-search/src/telemetry.rs @@ -0,0 +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}"); + } +} diff --git a/dht-search/src/verification.rs b/dht-search/src/verification.rs new file mode 100644 index 0000000..f6e8566 --- /dev/null +++ b/dht-search/src/verification.rs @@ -0,0 +1,274 @@ +// 负责按需调度种子可用性验证并以有界并发持久化验证结果 + +use std::sync::{ + Arc, + atomic::{AtomicU64, Ordering}, +}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use std::{collections::HashSet, net::SocketAddr}; + +use dht_crawler::DHTServer; +use dht_search::{ + domain::{InfoHash, VerificationResult}, + storage::{RocksTorrentRepository, TorrentRepository, VerificationPriority}, +}; +use tokio::task::{JoinHandle, JoinSet}; +use tokio_util::sync::CancellationToken; + +use crate::config::VerificationConfig; + +#[derive(Clone)] +pub(crate) struct VerificationIngress { + repository: Arc, + capacity: usize, + stats: VerificationStats, +} + +#[derive(Clone, Default)] +pub(crate) struct VerificationStats { + inner: Arc, +} + +#[derive(Default)] +struct VerificationStatsInner { + accepted: AtomicU64, + deduplicated: AtomicU64, + rejected_full: AtomicU64, + started: AtomicU64, + succeeded: AtomicU64, + failed: AtomicU64, + peers_discovered: AtomicU64, + handshakes_succeeded: AtomicU64, + queue_depth: AtomicU64, +} + +#[derive(Debug, Clone, Copy)] +pub(crate) struct VerificationSnapshot { + pub(crate) accepted: u64, + pub(crate) deduplicated: u64, + pub(crate) rejected_full: u64, + pub(crate) started: u64, + pub(crate) succeeded: u64, + pub(crate) failed: u64, + pub(crate) peers_discovered: u64, + pub(crate) handshakes_succeeded: u64, + pub(crate) queue_depth: u64, +} + +pub(crate) fn start( + repository: Arc, + server: DHTServer, + config: VerificationConfig, + cancel: CancellationToken, +) -> (VerificationIngress, JoinHandle>) { + let stats = VerificationStats::default(); + stats.inner.queue_depth.store( + repository.verification_queue_len().unwrap_or_default() as u64, + Ordering::Relaxed, + ); + let ingress = + VerificationIngress::new(repository.clone(), config.queue_capacity, stats.clone()); + let task = tokio::spawn(run(repository, server, config, stats, cancel)); + (ingress, task) +} + +impl VerificationIngress { + fn new( + repository: Arc, + capacity: usize, + stats: VerificationStats, + ) -> Self { + Self { + repository, + capacity, + stats, + } + } + + #[cfg(test)] + pub(crate) fn for_test(repository: Arc, capacity: usize) -> Self { + Self::new(repository, capacity, VerificationStats::default()) + } + + pub(crate) async fn enqueue(&self, hashes: Vec, priority: VerificationPriority) { + if hashes.is_empty() { + return; + } + let repository = self.repository.clone(); + let capacity = self.capacity; + let result = tokio::task::spawn_blocking(move || { + let outcome = + repository.enqueue_verification(&hashes, priority, unix_timestamp(), capacity)?; + let queue_len = repository.verification_queue_len()?; + Ok::<_, dht_search::storage::StorageError>((outcome, queue_len)) + }) + .await; + match result { + Ok(Ok((outcome, queue_len))) => { + self.stats + .inner + .accepted + .fetch_add(outcome.accepted as u64, Ordering::Relaxed); + self.stats + .inner + .deduplicated + .fetch_add(outcome.deduplicated as u64, Ordering::Relaxed); + self.stats + .inner + .rejected_full + .fetch_add(outcome.rejected_full as u64, Ordering::Relaxed); + self.stats + .inner + .queue_depth + .store(queue_len as u64, Ordering::Relaxed); + } + Ok(Err(error)) => tracing::error!(%error, "可用性验证任务持久化失败"), + Err(error) => tracing::error!(%error, "可用性验证入队任务异常"), + } + } + + pub(crate) fn stats(&self) -> VerificationStats { + self.stats.clone() + } +} + +impl VerificationStats { + pub(crate) fn snapshot(&self) -> VerificationSnapshot { + VerificationSnapshot { + accepted: self.inner.accepted.load(Ordering::Relaxed), + deduplicated: self.inner.deduplicated.load(Ordering::Relaxed), + rejected_full: self.inner.rejected_full.load(Ordering::Relaxed), + started: self.inner.started.load(Ordering::Relaxed), + succeeded: self.inner.succeeded.load(Ordering::Relaxed), + failed: self.inner.failed.load(Ordering::Relaxed), + peers_discovered: self.inner.peers_discovered.load(Ordering::Relaxed), + handshakes_succeeded: self.inner.handshakes_succeeded.load(Ordering::Relaxed), + queue_depth: self.inner.queue_depth.load(Ordering::Relaxed), + } + } +} + +async fn run( + repository: Arc, + server: DHTServer, + config: VerificationConfig, + stats: VerificationStats, + cancel: CancellationToken, +) -> Result<(), String> { + let mut active = JoinSet::new(); + let mut ticker = tokio::time::interval(Duration::from_millis(config.poll_interval_millis)); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + tokio::select! { + _ = cancel.cancelled() => break, + result = active.join_next(), if !active.is_empty() => { + if let Some(result) = result { + result.map_err(|error| error.to_string())??; + } + } + _ = ticker.tick() => { + while active.len() < config.max_active { + let claim_repository = repository.clone(); + let lease_secs = config.lease_secs; + let request = tokio::task::spawn_blocking(move || { + claim_repository.claim_verification(unix_timestamp(), lease_secs) + }) + .await + .map_err(|error| error.to_string())? + .map_err(|error| error.to_string())?; + let Some(request) = request else { break }; + stats.inner.started.fetch_add(1, Ordering::Relaxed); + active.spawn(verify_one( + repository.clone(), + server.clone(), + request.info_hash, + config.max_peer_attempts, + stats.clone(), + )); + } + } + } + } + while let Some(result) = active.join_next().await { + result.map_err(|error| error.to_string())??; + } + Ok(()) +} + +async fn verify_one( + repository: Arc, + server: DHTServer, + info_hash: InfoHash, + max_peer_attempts: usize, + stats: VerificationStats, +) -> Result<(), String> { + let peer_repository = repository.clone(); + let stored_peers = tokio::task::spawn_blocking(move || peer_repository.get(info_hash)) + .await + .map_err(|error| error.to_string())? + .map_err(|error| error.to_string())? + .map(|record| record.source_peers) + .unwrap_or_default(); + let lookup = server.lookup_peers(*info_hash.as_bytes()).await; + let dht_peers = lookup.map(|result| result.peers).unwrap_or_default(); + let mut unique = HashSet::with_capacity(dht_peers.len() + stored_peers.len()); + let mut peers = Vec::with_capacity(dht_peers.len() + stored_peers.len()); + for peer in dht_peers.into_iter().chain( + stored_peers + .iter() + .filter_map(|peer| peer.parse::().ok()), + ) { + if unique.insert(peer) { + peers.push(peer); + } + } + stats + .inner + .peers_discovered + .fetch_add(peers.len() as u64, Ordering::Relaxed); + let mut handshakes = JoinSet::new(); + for peer in peers.iter().copied().take(max_peer_attempts) { + let server = server.clone(); + let hash = *info_hash.as_bytes(); + handshakes.spawn(async move { server.verify_peer_handshake(hash, peer).await }); + } + let mut reachable = 0_u32; + while let Some(result) = handshakes.join_next().await { + if result.map_err(|error| error.to_string())? { + reachable = reachable.saturating_add(1); + } + } + stats + .inner + .handshakes_succeeded + .fetch_add(u64::from(reachable), Ordering::Relaxed); + if reachable > 0 { + stats.inner.succeeded.fetch_add(1, Ordering::Relaxed); + } else { + stats.inner.failed.fetch_add(1, Ordering::Relaxed); + } + let result = VerificationResult { + verified_at: unix_timestamp(), + discovered_peers: peers.len().min(u32::MAX as usize) as u32, + reachable_peers: reachable, + }; + let queue_len = tokio::task::spawn_blocking(move || { + repository.finish_verification(info_hash, result)?; + repository.verification_queue_len() + }) + .await + .map_err(|error| error.to_string())? + .map_err(|error| error.to_string())?; + stats + .inner + .queue_depth + .store(queue_len as u64, Ordering::Relaxed); + Ok(()) +} + +fn unix_timestamp() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() +}